Architecture Note · Draft Efficient & On-Device AI Agents

ChessReasoner120M

A transformer trained from scratch to play chess and say why — in natural language, at about 120 megabytes, with no engine and no search at inference.

The design principle in one line: verbalize the search, internalize the evaluation.

drag to orbit
119,566,080backbone parameters
~120MBint8 on device
18.9MBKV cache, full context
~500tokens per decision
01
Why this exists

A frontier model, a 40-line heuristic, and a tie

Two results bound the problem. Karvonen's Chess-GPT showed a 25M-parameter transformer trained on nothing but move lists reaching roughly 1500 Elo, with its internal board state linearly probeable at ~99%. DeepMind's Grandmaster-Level Chess Without Search put 270M parameters on Stockfish-annotated games and reached 2895 blitz Elo with no search at all.

So playing chess at small scale is solved. What is not solved is a small model that plays well and narrates why. Both results above are policy networks: they emit a move and nothing else.

The measurements below are from this project's own runs on a 1000-position held-out set. They are the reason the architecture looks the way it does.

Fig. 1 — Accuracy on held-out move selection (n = 1000)binary task · chance = 50%
Chance50.0%
Bag-of-words classifier
explanation text only, never sees the board
63.7%
Gemma 4 E2B, 4-bit
~1,150 output tokens per position
62.0%
40-line hand-written heuristic
two plies of captures, no engine, no ML
79.1%
DeepSeek V4 Flash, thinking
~19,345 output tokens per position
91.0%
Stockfish depth 1
quiescence only — no real search
92.0%
Stockfish depth 8100%

Three things fall out of this. A 2B model scores below a word-frequency classifier that never sees a chessboard, so its number carries no evidence of chess ability. The candidates differ by a median of 851 centipawns — this is blunder detection, not strategy. And depth-1 Stockfish, which is quiescence resolution plus a static evaluation, already beats the frontier model that spends 19,345 tokens per position.

02
The design principle

Verbalize the search.
Internalize the evaluation.

Explicit text can express a search tree, but not a deep one. Each ply of verbalized line costs 30–50 tokens, so five candidates at four ply is already ~800 tokens. Ten ply is not expressible at any sane length.

A strong static evaluation is the mirror problem: thousands of tuned terms, none of them verbalizable. You cannot write down "this knight is worth 3.4 because of the pawn structure" in a way that generalizes.

Fig. 1 says which half needs which treatment. Depth-1 Stockfish scores 92% — so the part that must be said out loud only has to reach a few forcing plies. Everything else lives in the weights.

The model says the capture sequence. It does not say the evaluation — it is one.

This has a direct architectural consequence. There is a second constraint that shapes it: the input must stay raw. Expanding the FEN into a piece list, or pre-applying the candidate moves, hands the model the exact sub-skill under test. So instead of scaffolding the input, this design scaffolds the training objective — auxiliary heads that force an accurate internal board during training and are deleted before inference.

03
The model, in three dimensions

Board in, prose out

Sixty-four squares enter as one fixed-length plane, flow through eighteen decoder blocks, and leave through four heads — only one of which survives to inference. Drag to orbit; click a stage to isolate it.

Fig. 2 — Interactive modeldrag to orbit
Whole model. The board enters as 72 tokens, is consumed by an 18-layer stack of 119.6M parameters, and exits through four heads. The three cold heads are auxiliary — they exist only while training.
04
Anatomy

Every part, and what it is for

Click any component. The rail explains what it does and why it is shaped that way.

Fig. 3 — Architectureclick a component
position 72 tokens 1 + 64 + 1 + 4 + 1 + 1 fixed board span 18 × decoder block d_model 768 GQA 12 Q / 4 KV SwiGLU 2048 119.6M parameters LM head next token · tied Board head 64 × 13 occupancy Value head 128-bin HL-Gauss Policy head 1968 moves training only — removed at inference kept at inference L = L_LM + λ_b·L_board + λ_v·L_value + λ_p·L_policy λ annealed to 0 over the second half of training
05
The tokenizer

Reading the board

The baseline's failure was not strategic. Given a real position it wrote "White pieces: K(g2), Q(h5), R(e1,f1), B(c1,e2)…" — the king was on h2, the bishop on g2, and c1, e2 and g3 were empty. Then it read the move f5f6 as the single square f5 and spent four hundred tokens confused that "the Queen is already on h5."

Both failures are tokenizer failures, and both are unrepresentable here. Sixty-four square names, twelve piece symbols and the structural markers are reserved as atomic tokens before BPE ever sees the corpus. A move is exactly (from, to[, promo]).

Fig. 4 — Board ⇄ token spanhover either side
Ruy Lopez, after 3…a6 4.Ba4 Nf6
Hover a square or a token.

Every square is emitted, including the empty ones, and always in the same a1 → h8 order. The span is 72 tokens whatever the position. That costs a little more than a raw FEN's ~80 subword tokens would — the win is not compression, it is positional regularity: square i always sits at a computable offset, so file (+1), rank (+8) and diagonal (+9/+7) neighbours are constant strides an attention head can learn once and reuse everywhere. Run-length FEN makes every square's position depend on decoding the digits before it.

Square tokens are also shared across roles. <f3> is the same embedding whether it names a board slot, a piece location in prose, or the endpoint of a move. The model's notion of the square is one object, not three.

06
Inside a block

Eighteen of these, stacked

Pre-norm, no biases, tied embeddings. Eighteen layers at width 768 is deeper than the usual 120M configuration — GPT-2 small and Pythia-160M are both 12 × 768. The depth is deliberate: serializing a search procedure is sequential composition, and depth buys composition steps that width does not.

Fig. 5 — One decoder blockx18
01 · Normalization

RMSNorm, fp32

Accumulated in float32 regardless of the activation dtype. The T4 is Turing and has no bfloat16, so fp16 training without fp32 reduction diverges around step two thousand.

02 · Attention

GQA · 12 Q / 4 KV

Four key-value heads instead of twelve. At full context that is a 18.9 MB KV cache rather than 56.6 MB — which is the difference between comfortable and awkward on a phone.

03 · Position

Standard 1-D RoPE

A two-dimensional board-aware variant was designed, measured, and cut. The fixed raster already puts neighbours at constant offsets; the 2-D scheme only repaired file wrap-around, for half the rotary dimensions. See §10.

04 · Feed-forward

SwiGLU 768→2048→768

Three matrices, no bias. 4.72M parameters per layer, which is 75% of the block — the usual ratio, and the reason depth is cheaper than width here.

per layer attention 1,572,864
SwiGLU 4,718,592
RMSNorm 1,536
─────────
6,292,992 × 18 = 113,273,856
embeddings (tied) 6,291,456
final norm 768
─────────
backbone 119,566,080
07
The heads

Three of the four are thrown away

The auxiliary heads are 2,251,632 parameters that never ship. They exist to shape what the residual stream carries, and are annealed to zero weight across the second half of training so the model never learns to lean on them.

Fig. 6 — Head placement and purpose
HeadShapeFires atTargetWhy
LM — kept768 → 8192every positionnext tokenThe whole inference path. Tied to the input embedding.
Board768 → 64×13</LINE>occupancy at the leafForces an accurate internal board.
Value768 → 128</FEN>HL-Gauss win probabilityThe internalized evaluation from §02.
Policy768 → 1968</FEN>engine best moveCandidate prior — stops the model analysing moves nobody would play.

The board head's placement is the load-bearing decision. At </FEN> the board is verbatim in context and the head is a copy — it teaches nothing. At the leaf of an analysed variation the model must have simulated the moves, and that position exists nowhere in the input. That is what makes it a world-model objective rather than a parsing objective.

A caveat that has to be measured

The policy head fires at </FEN>before any reasoning tokens exist. Training it there teaches the model the answer before it thinks, which is the textbook setup for post-hoc rationalization: the trace becomes decoration over a decision already made.

This is defensible as a System-1 / System-2 split, and it is why the head is deleted at inference. But it cannot be assumed. The required experiment is to compare the policy head's own top-1 accuracy against the full model's post-reasoning accuracy. If they match, the reasoning is decorative, and the paper has to say so.

08
The corpus

3.2 billion tokens, none of them human

Everything is generated, and everything is CC0-sourced. The unlock is that the Lichess evaluation database is already Stockfish-annotated — roughly 500 million positions with centipawn score, depth and principal variation. Annotation stops being the bottleneck.

Every factual claim is generated from python-chess ground truth, so the corpus is correct by construction. That is what distillation from a frontier model cannot offer — the DeepSeek traces collected for this project contain fabricated engine evaluations on positions never analysed.

Fig. 7 — Corpus tiers≈ 3.2B tokens
TierContentExamplesTokensWhat it teaches
1Board literacy6.0M570MWhat is on e4. Who attacks d5. What is hanging. Built.
2Tactical primitives5.0M775MExchanges verbalized ply by ply; motifs grounded in puzzle themes.
3Full reasoning traces3.0M960MCandidates, forcing lines, comparison, choice.
4Game continuity1.5M420MPlans across moves, not isolated tactics.
5General English480M15% FineWeb-Edu, so it is not mute outside chess.

Three rules that decide whether it works

Randomize everything structural. Candidate count, discussion order, which features get mentioned, phrasing. If the best move is always discussed last, the corpus has a positional shortcut — precisely the pathology Fig. 1 measured in the public dataset.

Only machine-checkable predicates. python-chess can verify occupancy, legality, attacks, captures and material. It cannot verify "gains space" or "the initiative is gone". Evaluative language is excluded so the correctness guarantee is exact rather than approximate.

Include corrective traces. Some fraction must explore a candidate, find the refutation, and reject it. A corpus of only-correct-first-guess traces teaches confident assertion, not search.

Fig. 8 — Amortizing the board planemeasured on a 20k-board shard
PackingBoard shareAnswer shareTokens/example
One question per board81.6%8.0%88
6–10 questions per board38.3%26.1%188

Re-emitting the 72-token board for every question spends most of the compute budget re-reading the same position. Packing several questions onto one board gives 3.3× more supervised signal for the same FLOPs, and forces the model to hold the position across a long span instead of answering off the most recent tokens.

09
Training

Forty hours, give or take twenty

Fig. 9 — Curriculum
StageTokensContentActive heads
1570MBoard literacyboard (high λ), value
2775MTactical primitivesboard, value, policy
3960MFull tracesall three, annealing
4420MGame continuityannealed to zero
5Rejection-sampling SFTnone

Stage 5 is where the free verifier pays off: sample k traces, keep only those where the move is right and every stated board fact verifies. That double filter is the difference between reinforcing reasoning and reinforcing lucky guesses with hallucinated boards.

Optimizer
AdamW · β 0.9/0.95
Peak LR
3e-4 → 3e-5 cosine
Precision
fp16 + loss scaling
Effective batch
262,144 tokens
Steps
≈ 12,200
Attention kernel
SDPA mem-efficient
6ND = 6 × 1.196e8 × 3.2e9 ≈ 2.3e18 FLOPs

T4 peak fp16 65 TFLOP/s
realistic MFU 15–20% (measured, not assumed)
─────────────────────────────
≈ 40–60 hours on one T4

Two Turing-specific traps, stated in advance because both cost a wasted run to discover. FlashAttention-2 requires Ampere (sm_80+); the T4 is sm_75, so it will not build — the SDPA memory-efficient backend is the one that works. And there is no bfloat16, so RMSNorm and the loss reduction must accumulate in fp32.

The Kaggle weekly quota is 30 GPU-hours with 9–12 hour sessions, so a 40-hour run is about five sessions across two calendar weeks — or half the wall-clock on a dual-T4 kernel.

10
Adversarial review

What broke, and how it was found

The design was attacked before it was implemented, and the implementation was attacked after. Three of four checked design claims came back against the design; four real bugs surfaced in code that already had a passing test suite. All of them are recorded here because a design note that only lists its successes is not evidence of anything.

R1 · design · demoted

The board-aware positional encoding was solving a solved problem

A two-dimensional rotary scheme was designed to make board geometry legible to attention. But the fixed a1→h8 raster already does that — file neighbours are +1, rank +8, diagonals +9/+7, all constant 1-D offsets that standard RoPE handles. The 2-D scheme repairs exactly one thing the raster does not: file wrap-around, where h1→a2 is also +1 and is not a file step.

affects 10.9% of file adjacencies · costs half the rotary dimensions · cut from v1
R2 · design · fixed

Most of the gradient was going to empty squares

57% of the 64 board-plane tokens are the empty symbol, so an unweighted loss spends the majority of its gradient on near-zero-information tokens. Masking the plane entirely is also wrong — predicting it teaches a prior over plausible piece configurations. The answer is a per-token weight, not a mask.

board 0.1 · prompt 0.0 · answer 1.0
B · code · fixed

Attention leaked the future through the KV cache

is_causal=True was passed whenever query and key lengths matched. With a populated cache and more than one query token — a prefill continuation, or speculative decoding — that silently lets a token attend to its own future. Single-token decoding never triggered it, so the test suite was green.

leak 2.9e-017.2e-07
C · code · fixed

The position sampler ignored its own random seed

The puzzle-database reader accepted an rng and never used it, so a shard was a contiguous prefix of a file sorted by puzzle id — and id correlates with rating and theme. This is the same mistake that once produced a task set in which every single position had an empty a8, because it was taken in FEN-lexicographic order.

different seeds returned identical positionsindependent samples
D · code · fixed · the serious one

The model was being trained to answer questions about boards it could not see

Slicing a concatenated token stream into fixed windows cuts examples in half. Measured on a real corpus, 15.9% of supervised answer tokens ended up in a window with no board span before them. The loss still demanded those answers, so the model was being trained to invent them — which is precisely the hallucination this entire project exists to remove. Packing now fills windows with whole examples and pads the tail.

orphaned answers 4,440 of 26,7730 · 88.7% utilization

Two further findings did not become bugs but changed the numbers. The claim that the corpus is "100% factually correct by construction" only holds for machine-checkable predicates, so the generator was restricted to them. And the compute estimate assumed 25–30% MFU, which is optimistic for a model this small on a memory-bandwidth-bound card — hence the honest 40–60 hour range, pending an actual measurement on an actual T4.

11
Evaluation

The metric to lead with

Accuracy on a saturated binary benchmark is not the interesting claim. The interesting claim is trace factuality — the fraction of checkable assertions in a generated trace that verify against python-chess.

It is fully automatable, nobody reports it, and it measures faithful reasoning rather than leaderboard position. A 120M model producing 95%-factual traces against a frontier model at 70% and seventeen times the tokens is a stronger result than beating it on accuracy — and it is a claim about reasoning, which is what this is actually about.

Fig. 10 — Evaluation batterytrained on none of these
MetricWhy it earns a place
Trace factualityCheckable claims that verify. Directly measures whether the narration is real.
Puzzle rating spreadHeld-out Lichess puzzles across the full range — a number on a known scale.
Playing EloVersus Stockfish levels 1–8. No position-selection benchmark can be gamed into this.
Board-probe accuracyLinear probe on frozen residuals for 64-square occupancy. Evidences the world model.
Tokens per decision~500 here against 19,345 for the frontier reference.
Shortcut controlsAlways-A, longest-explanation, text-only-TF-IDF, explanation-stripped. Mandatory.
The honest uncertainty: whether 120M parameters on 3.2B tokens of synthetic chess text will reason, or merely pattern-match, is not known in advance.

The ablations and the factuality metric are designed to detect the difference rather than paper over it. Given that a 25M model reaches 1500 Elo from move lists alone, a defensible expectation is somewhere around 1800–2200 puzzle rating — and that is a guess, flagged as one.

The next step is not a training run. It is a four-hour probe: a 20M model on 200M tokens of board-literacy data. If it cannot learn to read a board and answer what is on e4, nothing downstream will work, and that is worth knowing before generating three billion tokens.