chessmamba-v2

ChessMamba v2

A chess engine whose evaluation is a selective state space model that scans the board along its own geometry. Not alpha beta with a hand tuned evaluation, not NNUE, not a convolutional net with MCTS. Every layer sweeps each rank and each file as its own short sequence, forwards and backwards, with a learned decay deciding how far along a line influence carries. Rooks and queens act along ranks and files, so that is the axis the scan runs on.

4.3M parameters. No handcrafted evaluation: no piece values, no piece square tables, no opening book, no endgame tables, no engine generated labels. It learned from human games and from who won. Legal move generation and the board representation come from python-chess, so the rules are given; the judgement is not.

Rating

1936, plus or minus 105. One Bradley Terry maximum likelihood fit over 3,300 games, anchored at a single point: Stockfish UCI_Elo 1320.

Two consistency checks fall out of the fit that were not fitted for. The gap between v2 and Sunfish comes out at 168.4 points, and the separately measured head to head figure was 168. The gap between the fixed and unfixed engine comes out at 90 points, and the direct A/B between them measured 91. Those agreeing is the reason to believe the graph hangs together rather than being a pile of unrelated matches.

The ladder the number rests on, fixed engine, 100 games each:

The rest of the opponents, 200 games each. These were played before the repetition fix, so they understate the engine by roughly 90 points; they are left as measured rather than adjusted.

Two caveats that do not go away. The whole scale hangs on UCI_Elo 1320 being an honest 1320; Stockfish's limited strength mode is generally thought to overperform its label at fast time controls, which would make this estimate conservative, but it is an assumption rather than a measurement. And everything here is 500 ms a move on one RTX 3090. The search is the part that scales with hardware, so a different machine or a different clock moves the number.

Do not trust a rating on a label

The obvious way to anchor an engine is to play it against something with a published rating. That does not work as well as it looks, and checking it was worth the compute.

Maia is a family of networks named 1100 through 1900 after the human rating band each one imitates. It looks like an ideal anchor: seven labelled levels, all free. Played head to head against each other, 200 games a pairing, six pairings, the entire 800 point span of the labels is worth about 143 points of actual strength. Maia 1500 and Maia 1700 are a coin flip. The ordering is right and the scale is not, because a network trained to imitate a 1900 rated human still blunders like one, and with no search two of them hang pieces at similar rates.

Stockfish's UCI_Elo survives the same test. Measured the same way it tracks its own labels closely at the bottom and compresses toward the top, 880 nominal points coming out as 610 real ones. That is usable, which is why the rating above is anchored to it, using the measured scale and not the label.

Placed on that scale, Maia 1900 plays at about 1430.

The repetition bug

Worth writing down, because the engine shipped with it and it cost around 90 Elo.

The board encoding carries a three state repetition field: this position is new, this position has occurred once before, this position has occurred twice or more. The data pipeline counted prior occurrences, so a fresh position was state 0. The search counted occurrences including the current one, so a fresh position became state 1. Every ordinary node in the tree was being told it was a repeat.

The size of the mismatch is the part that matters. Across four million training positions, state 0 covers 99.279%, state 1 covers 0.716%. Instrumented from the opening position, the search sent state 1 on 597 evaluations out of 601 and state 0 on none. The engine was running almost entirely on an embedding it had seen on well under one percent of its training data, and one that means "drawish shuffling", because in training the positions carrying it genuinely were repetitions.

Measured on 600 real positions, holding everything else fixed and changing only that field: the value head is pulled toward a draw in 67% of positions, mean absolute shift 0.055 on a minus one to plus one scale, and the policy head picks a different best move in 15.7% of positions. The value distortion mostly cancels under search, since PUCT compares sibling Q values and a uniform squash preserves their order. The policy distortion does not, because priors decide where the search spends its nodes.

Fixed by passing the count minus one to the network while threefold detection keeps the raw count. Measured head to head against the unfixed engine over 200 games: 83 wins, 85 draws, 32 losses, 0.627, +91 ± 50.

All figures on this card are the fixed engine unless the row says otherwise.

What went wrong in v1

v1 read a game as a list of moves and never saw a board. Three faults, all measurable rather than matters of taste.

Its state space layers had a memory half life of 1.1 plies. dt_bias was initialised to zeros, which puts delta at softplus(0) = 0.693 and halves the slowest state every step. Measured on the released v1 checkpoint, every block, after its full training run: per step retention 0.529. Over a 96 ply window that keeps 1.4e-17 of what happened 56 moves earlier. A model whose only input is a move list, and whose memory reaches back about one ply, cannot hold a position in mind. Legal move masking and the search were carrying it.

Its value head regressed on plus or minus one game outcomes with MSE, from games at an 1800 floor. Leaf evaluation in the search is entirely the value head, so that capped everything.

Train and validation were split by position rather than by game, so positions from the same game sat on both sides and every validation number it reported was optimistic.

What v2 does instead

The board, not the move list. 64 square tokens carrying piece identity, square, and markers for the last two moves, presented from the side to move's point of view. Castling rights, en passant file, halfmove clock and repetition count enter as a conditioning vector added to every square.

Rank and file scans, not a raster. A 64 long scan over the squares in order makes h1 and a2 adjacent, which is wrong, and puts a1 and a8 fifty six steps apart when they are one rook move apart. Scanning each rank and each file as its own length 8 sequence fixes both. Measured receptive field: one layer reaches exactly the rook lines from a square, two layers reach all 64.

A bilinear policy head. Each square projects a from vector and a to vector, and a move scores as their dot product, so the 4096 move logits come from per square features rather than from one pooled vector.

A win draw loss head rather than MSE on a scalar. Time forfeit games, 19.5% of the corpus, keep their positions for policy training but have their value loss masked out, because the result there reflects the clock rather than the position.

Batched PUCT search. The tree collects a batch of leaves under virtual loss and evaluates them in one pass. v1 could not do this: its hidden state depended on the path taken to a position, so two move orders reaching the same board evaluated differently, which also ruled out a transposition table.

The finding worth keeping

How fast a square forgets its neighbours along a line is the hyperparameter that matters here. It is worth 1.88 points of move accuracy, roughly three and a half times the effect of doubling the number of scan directions, and more than tripling the parameter count achieved. It has a clean optimum at a half life of about four squares, half a rank, which is where a sliding piece actually exerts influence.

Mamba's own default puts that half life at 76 squares on an 8 square line, which is indistinguishable from uniform averaging. That default is tuned for long text. Taking it on faith would have cost 1.88 points silently.

Two things were rejected on measurement rather than taste: running all four scan directions in every layer, which was real but too small to justify 1.7 times the cost given the cheaper variant sees 70% more games in the same time, and a full rank dt_proj, which was three times the parameters for nothing.

Training

Supervised on the Lichess Elite database, both players 2400 and above, bullet excluded. 3,967,146 games, 355,073,653 positions. 100,000 steps at batch 1024, about 102M positions, roughly ten hours on one rented RTX 3090. No engine labels anywhere. Nothing is distilled from Stockfish and no evaluation function is borrowed.

Validation is split by game, so no position appears on both sides. Final: policy cross entropy 1.3236, top-1 move match 55.80%, win draw loss accuracy 61.6%.

Search

The search is memory bandwidth bound rather than compute bound, which decides where it runs. On a GPU it reaches 1512 nodes per second, against 191 for v1. On a CPU it manages 57 and plays badly, so the engine selects CUDA when it is available. Set CHESSMAMBA_DEVICE=cpu to override.

It wins while searching roughly eight times fewer positions than Sunfish, which is the point of evaluating a position with a network instead of a table.

Running it

pip install torch python-chess
python3 engine_uci_v2.py

Speaks UCI. Loads ckpt/model.pt, or point CHESSMAMBA_CKPT elsewhere. Also CHESSMAMBA_DEVICE, CHESSMAMBA_BATCH (leaves per network call, default 64), CHESSMAMBA_SIMS.

Play it in the browser at TobiasLogic/chess-arena-gpu.

What this is not

1936 is club strength. Full Stockfish would win every game, and so would most phone apps. It wants a GPU. The rating is anchored at one point, on one machine, at one fast time control, and the interval on it is wide. Read it as a band from roughly 1830 to 2040 rather than a number.

Layout

engine_uci_v2.py   UCI loop, entry point
model_v2.py        the network
mcts.py            batched PUCT search
board_io.py        board and move encoding
ckpt/model.pt      trained weights
training/          data pipeline, training, matches, benchmark
graphs/            the figures above
games_vs_v1.pgn    all 200 games against v1
games_vs_sunfish.pgn  all 200 games against Sunfish
PLAN_V2.md         the working notes, including what failed