jaxverse
Contents

jaxverse

A little universe of learning machines, designed and built by Neo Mohsenvand.

Chapter 8 · Everything at once · ≈17 min

Rook

One small language model learns chess three ways: pretrained on games, fine-tuned on better ones, then sharpened by verifiable reward. The modern pipeline, in miniature.

Every large language model you have talked to was made in three acts. Pretraining: predict the next token on an enormous, indiscriminate corpus, until the model speaks the language of its data. Supervised fine-tuning: keep training, but on a small corpus somebody curated, until it speaks the way you want. And a last act of reinforcement learning: let the model generate, let a judge score what comes back, and move the weights toward whatever scored well. The whole modern pipeline fits in one sentence — pretrain on everything, fine-tune on the good stuff, reinforce on verified outcomes.

This chapter runs all three acts on one small transformer, live, in this page. Not a diagram of the pipeline — the pipeline, with every gradient computed on your GPU. The model is called Rook, and it learns chess.

Chess, first, because moves are tokens. A move written in UCI notation — e2e4, g8f6 — is a short string, and only 1,930 distinct ones ever occur in Rook's world. Add one marker token, ⟨game⟩, for “a new game begins”, and you have a vocabulary of 1,931. A game is a sentence written in it. The game of Chapter 5 — predict the next token, score it with cross-entropy, descend — carries over without changing a comma.

Plate I Moves are tokens
one token, one move point at any token
e2e4from squareto squareno separator, no piece letter, nothing else a8b8c8d8e8f8g8h8a7b7c7d7e7f7g7h7a6b6c6d6e6f6g6h6a5b5c5d5e5f5g5h5a4b4c4d4e4f4g4h4a3b3c3d3e3f3g3h3a2b2c2d2e2f2g2h2a1b1c1d1e1f1g1h1

Every square carries a name — the letter of its column, then the number of its row. That is the whole notation, and the only thing a move token is made of.

the vocabulary · 1,931 tokens
0 ⟨game⟩ the marker

Sorted alphabetically, so an id carries no meaning of its own — the model has to learn what 1013 does from the company it keeps.

a game is a sentence written in them
⟨game⟩ 0
9 ? next token

That column is the whole input: no board, no piece list, no rules. Rook reads the numbers and has to name the next one.

one token runs longer

142 of the strings run to five characters, like h7h8q: a pawn reaching the last rank, and the piece it turns into. Point at that one and the anatomy grows a third bracket.

Rook's entire language: 1,930 move strings and one marker. A token is the square a piece leaves and the square it lands on, glued together — point at any of them and the board says where those squares are, which is the moment e2e4 stops being a code and becomes a move. Nothing here says which moves are playable; that is a different question, and it comes next.

Chess, second, because every claim is checkable. A move is legal in a position or it is not, and a small library of rules (this page carries chess.js, a complete referee) settles the question instantly, every time. Nothing in the vocabulary above knows this: the same three tokens Rook can emit at any moment are, in one particular position, two moves and one piece of nonsense.

after ⟨game⟩ e2e4 e7e5 — three tokens Rook could name next
g1f3 ✓ legal here
the knight leaves the back rank for an empty square
f1c4 ✓ legal here
the diagonal opened when the e-pawn stepped out of the way
e4e5 ✕ illegal here
e5 is occupied — a pawn captures diagonally, never ahead
Every one of these is a perfectly ordinary token: e4e5 is a fine move in thousands of other positions, and nothing in the vocabulary marks it as wrong in this one. Legality is a property of the board, not of the string — and the only thing on this page that knows the difference is chess.js, which pretraining never asks.

That gap is what the V means in RLVR — reinforcement learning from verifiable rewards1 : the judge is not anyone's taste but a verifier that cannot be flattered, bribed, or fooled by confident nonsense. When frontier models are trained to reason about mathematics and code, this is the same move: reward what can be checked.2

One more thing, and it shapes everything that follows: Rook never sees a board. No 8×8 grid, no piece list, no rules — only text, one move-token after another, 128 plies of context at most. Everything it appears to know about squares and pieces has to be squeezed out of the statistics of move strings. Keep that in mind as you watch it play: the little diagrams on this page are for you; the model gets the sentence.

Meet the student, then: four transformer blocks, 128 numbers wide, four attention heads — 1.3 million parameters. The same weights thread through this entire chapter; what pretraining builds is exactly what fine-tuning bends and reinforcement sharpens. The plates need WebGPU (a current Chrome or Edge) and wake on their own as you reach them — the whole chapter shares one model, running on your GPU.

Plate II Pretraining — predict the next move
waking Rook — corpus + trained checkpoint, ≈2 MB…

one 1.3M-parameter transformer on your GPU, shared by every plate in this chapter — nothing leaves this page

Nothing in the loss says “chess”: the model is only playing predict-the-next-token, exactly as in the chapter on next-token prediction. Legality arrives as a side effect — the gauge asks, at 32 positions from real games, whether Rook's top-choice move is legal there; rewind to step 0 and watch it climb. Beside it is one game the model wrote, replayed against the referee: point at any move to see it played, and at the ✕ to see the board where the judge stopped believing it.

What did pretraining buy? At step 0 the time machine shows you noise: roughly one sampled move in a hundred is legal, and that one by accident. By step 2,600, about 87 in a hundred are legal — played in positions the model has never seen, under rules nobody stated. The only pressure ever applied was next-token prediction over some six thousand games. Legality condensed out of statistics, the way grammar condenses out of text: not because the loss asked for it, but because predicting the corpus well is impossible without it.

Legality was not the only thing that condensed, and the rest can be checked rather than argued about.3 Fit the simplest readout you can build — one linear layer, no bend — to the numbers running through a trained model, and ask it where the pieces are. They are in there, more clearly with every block. In the cleaner case of Othello the board comes out sharply enough to edit: change it inside the activations, let the model carry on, and its moves follow the board you planted rather than the one its own moves implied.4 That first result needed a probe with a bend in it, and what fixed it was the question rather than the model — ask each square black or white and a straight readout struggles; ask it mine or theirs and one works fine.5 Which is why the plate below is drawn from the side to move's point of view. On chess models trained the way Rook is, the readout recovers more than the position — it recovers an estimate of how strong the player being imitated was.6

Plate III Where the pieces are

reading the probe results…

A measurement rather than a training run, and made offline: the probe is fitted on a hundred thousand games against a Rook of the same design but eight times the size. Fit the simplest readout there is — one linear layer, no bend — to the numbers running through the model at each block, and ask it for the occupancy of all sixty-four squares. Above: one position, then the board that readout returns from each block in turn, with a dot on every square it gets wrong and the count of them underneath — which is for this one position and will wobble from block to block, because one position is never evidence. Below, on the same left-to-right axis of depth: how often it is right, against the two numbers that decide whether that means anything. Always answering 'empty' scores the dashed line. The same probe, fitted the same way to the same architecture with its weights left random, scores the grey one — a network that has learned nothing still carries some trace of the moves it was fed, and an honest claim has to clear that rather than merely the dashed line. What training adds is the gap between the grey line and the ultramarine one, and it widens with depth.

None of which is proof of understanding, and the plate above is careful about the difference. A probe that scores well might be reading the model, or it might be reading the moves that were fed into it — which is what the control line measures, and why it sits above the dashed one rather than on it. The claim the experiment actually supports is narrower and still remarkable: a machine that has only ever seen sentences ends up carrying something shaped like the thing the sentences are about.

Now look at what pretraining did not buy. Read the sampled games: pieces drift, nothing is hunted, material hangs unpunished. Of course it does — the corpus is random legal play, and imitation is faithful to a fault. Rook plays like its diet, because its diet is the only definition of “good move” it has ever been shown. The fix is not more data. It is better data. Before we administer it, sit across the board from your student and feel what 87% legal actually plays like.

Plate IV Play it — you are White
waking Rook — corpus + trained checkpoint, ≈2 MB…

one 1.3M-parameter transformer on your GPU, shared by every plate in this chapter — nothing leaves this page

You are playing a language model: every reply is a sampled sentence-continuation, masked to legality. The gauge is the honest score — how much probability Rook put on legal moves before the mask saved it. And this board always plays the CURRENT weights: fine-tune or reinforce below, come back, and feel the difference.

Better data, same machine

To see why the second act exists, look at what the first one actually produces. A pretrained model is a completion machine: ask a raw base model a question and it may answer — or continue with three more questions, or a survey form, because that is what question-shaped text does in its corpus. The assistant on the other end of a chat window is made from one by supervised fine-tuning7 : continued training on a curated set of (instruction → good response) pairs, written and vetted by people. It adds little knowledge. What it teaches is the format of being helpful — that in this house, a question is followed by an answer. One pair from such a corpus looks like this:

instruction
Explain gravity to a child.
curated response
The Earth is always gently pulling everything toward itself — you, your ball, the rain. That pull is called gravity. It is why things fall down instead of up, and why you come back when you jump.

Mechanically, nothing changes: the same weights, the same cross-entropy, the same update θθηL\htmlClass{eq-model}{\theta} \leftarrow \htmlClass{eq-model}{\theta} - \htmlClass{eq-knob}{\eta} \htmlClass{eq-world}{\nabla \mathcal{L}} — only the corpus is now chosen on purpose. For Rook the mapping is direct. The “instruction” is the game so far; the “good response” is what a competent player did next. Where a frontier lab curates tens of thousands of demonstrations, we curate 2,381 games played by a greedy little bot that grabs material whenever it can, and half the time gives check when it cannot. Its style is loud: 38% of its moves are captures, against 12% in the random-play corpus underneath, and one game in twenty ends in checkmate rather than one in seventy.

One paragraph of honesty before the button. Fine-tuning shifts style, and it charges a price: as the weights lean toward greedy play, the old random-play corpus becomes more surprising to the model, and its validation loss there quietly rises. That drift is not a malfunction — it is what specialization looks like from the old distribution's point of view. It is also why fine-tuning is run at a smaller step size than pretraining (we ease η\htmlClass{eq-knob}{\eta} from 1.2·10⁻³ down to 3·10⁻⁴): nudge the weights and the style shifts; blast them and the old competence goes too. The plate below measures both sides of the trade, because a curve that only shows the win is an advertisement, not an experiment.

Plate V Fine-tuning — the same weights, a better diet
waking Rook — corpus + trained checkpoint, ≈2 MB…

one 1.3M-parameter transformer on your GPU, shared by every plate in this chapter — nothing leaves this page

Nothing about the machine changed — same parameters, same loss, same update rule — only the corpus. Style follows the diet: within a hundred-odd steps the capture rate of Rook's own games climbs from the random-play 12% toward the greedy corpus's 38%, while the held-out loss on the corpus it left drifts up. That drift is the price, printed beside the prize. Stop while you are ahead: over-tuning keeps buying style and starts paying with competence.

A judge instead of examples

Before moving on, scroll back and play your fine-tuned student — Plate IV always drives the current weights, and you will feel the difference: it comes for your pieces now.

The last act needs no examples at all. It needs a judge. Here is one full iteration of RLVR, as this page runs it: take a real opening from the corpus — four to ten plies, different every time — and have Rook finish it G=8G = 8 different ways, up to fourteen plies each, sampled hot enough to disagree with itself. Then chess.js replays every continuation from that exact position and issues a verdict no one can argue with. Rollout ii earns

ri  =  ini  +  12[i=ni]\htmlClass{eq-world}{r_i} \;=\; \frac{\htmlClass{eq-world}{\ell_i}}{\htmlClass{eq-world}{n_i}} \;+\; \tfrac{1}{2}\,\big[\,\htmlClass{eq-world}{\ell_i} = \htmlClass{eq-world}{n_i}\,\big]

where i\htmlClass{eq-world}{\ell_i} counts the consecutive legal plies before the first illegal one, ni\htmlClass{eq-world}{n_i} counts the plies attempted, and the bracket pays a half-point bonus only when the whole rollout survives the judge. Play twelve clean plies out of fourteen and you score well; break on ply two and you score badly; stay perfect and you are paid extra.

Raw rewards are not used directly. Within the group of eight, each reward is compared to the others:

A^i  =  rimean(r1,,rG)std(r1,,rG)\hat{A}_i \;=\; \frac{\htmlClass{eq-world}{r_i} - \htmlClass{eq-op}{\operatorname{mean}}(\htmlClass{eq-world}{r_1},\ldots,\htmlClass{eq-world}{r_G})}{\htmlClass{eq-op}{\operatorname{std}}(\htmlClass{eq-world}{r_1},\ldots,\htmlClass{eq-world}{r_G})}

This A^i\hat{A}_i is the advantage: how much better or worse rollout ii did than its own siblings. Chapter 6 stabilized learning by measuring outcomes against a baseline instead of trusting raw reward; this is the same idea with the scaffolding removed — the group is the baseline. That trick — the heart of the GRPO family of methods8 — needs no second network to estimate values: eight siblings, sampled from the same position, are estimate enough.

The update is REINFORCE, exactly as the policy-gradient chapter promised it would be:

θJ    i=1GA^i  θlogπθ(rollouti)\nabla_{\htmlClass{eq-model}{\theta}} J \;\approx\; \sum_{i=1}^{G} \hat{A}_i \;\nabla_{\htmlClass{eq-model}{\theta}} \log \htmlClass{eq-model}{\pi_\theta}(\text{rollout}_i)

— raise the log-probability of every token in the above-average rollouts, lower it for the below-average ones, and touch only the generated tokens: the prefix was given, so it gets no credit. And one honest edge case, which you will see happen: if all eight rollouts earn the same reward, the standard deviation is zero, the advantages are undefined, and the step is skipped. When every answer is equally good, there is no gradient — reinforcement learning learns from differences, and a group of equals teaches nothing.

Now the two things this update does not have, because chapter 7 just spent a chapter installing both, and a plate that quietly omits them would be teaching you the wrong lesson. There is no importance ratio and no clip. That is not a simplification — it is a consequence: we take exactly one gradient step per group of rollouts, so the policy that generated the data is the policy being updated, every ratio is 1, and a fence at 1 ± ε would never be touched. Clipping machinery is the price of reusing a batch, and this loop does not reuse one. Read any paper in this family with that in hand: how much clipping it carries is a direct readout of how many updates it takes per generation.

And there is no KL leash. A frontier run would carry one — the same β\htmlClass{eq-knob}{\beta} you slid a moment ago — pinning the policy to the model it started from. Here the leash's job is done by a very small step size instead, and the choice was forced rather than principled: at pretraining rates this policy collapses into repeated tokens within a dozen updates. That is measured, not feared. What you are watching is the cheapest thing in the family that still works, and knowing exactly which safety rails have been taken off is the difference between running a demo and reading one.

Plate VI RLVR — reinforcement from a judge
waking Rook — corpus + trained checkpoint, ≈2 MB…

one 1.3M-parameter transformer on your GPU, shared by every plate in this chapter — nothing leaves this page

No example ever said “this move is good”. A verifier said “this rollout held up longer”, the group's own average set the bar, and the gradient did the rest — which is why the eight are ranked here with that average drawn as a line: everything above it gets reinforced, everything below it discouraged. Point at any rollout to see the board where the judge stopped it. Sample, verify, standardize, reinforce: a small, honest cousin of how frontier models learn to reason.

Line them up

You have now made three different chess players out of one set of weights, and the honest way to compare them is not a curve — it is a decision. The page has been quietly photographing your work: every time the fine-tuning or RLVR loop pauses, it snapshots the weights, so all three stages survive even though the resident model keeps moving on. The arena below puts them at the same board. Play a move, and every fielded Rook is asked the same question — what would you do here? — each answer drawn as an arrow in its stage's color, with the numbers that explain it: how much of its belief was even legal, and what it wanted most before the mask. Watch the pretrained arrow wander while the fine-tuned one aims at your pieces; watch the RLVR column hold its legal mass in positions that make the earlier stages guess.

Plate VII The arena — every stage, same position
waking Rook — corpus + trained checkpoint, ≈2 MB…

one 1.3M-parameter transformer on your GPU, shared by every plate in this chapter — nothing leaves this page

Same board, three students. After each of your moves, every fielded Rook is asked for its next-move beliefs — the weights are swapped in place on the one engine, questioned, and restored — and each decision lands on the board as an arrow: ultramarine pretrained, violet fine-tuned, green RLVR. The columns keep score honestly: legal mass climbs stage by stage, and the fine-tuned Rook starts hunting your pieces. The stage you pick under “plays Black” is the one that actually answers.

Where the judge works today

The loop you just ran is the loop, at reduced scale. When today's models are trained to write code, the verifier is a bank of unit tests: the program compiles and passes, or it does not, and no amount of confident prose changes the verdict. That is a move being legal or not, wearing different clothes. In mathematics, the judge checks the final answer, or a proof checker walks the argument line by line. Same iteration, exactly: sample a group of attempts from one problem, let the verifier score each, standardize within the group, reinforce what held up. The graders got bigger; the gradient did not change.

This is also why the verifiable domains are the ones moving fastest, and you have already felt the other half of that sentence. A judge made of rules is cheap, instant, tireless, and incorruptible; you can ask it a million times a day and it never lowers its standards, so no curve of yours ever peaks and turns over. A judge made of taste is expensive, slow, noisy, and has holes an optimizer will find. So coding and mathematics improve at a pace that essay-writing and advice do not — the pipeline is identical, and only some subjects come with an answer key.

The caveat is that "no feelings to flatter" is not "no exploits": a model rewarded for passing tests will sometimes learn to game the tests — patching the assertions, special-casing the inputs, exiting early with the expected code — and building verifiers that cannot be lawyered is a live problem rather than a solved one. There is a nastier turn to it, too. Read the model's own reasoning and you can often catch it planning the exploit in plain language; train against that signal and it does not stop exploiting, it stops saying so.9

There is a second consequence of grading on a curve, and it decides how a run is fed. A group only teaches when its members disagree: all eight legal and the advantages vanish, all eight hopeless and they vanish too. So a group-relative method learns only at the frontier of what the model can currently almost do — which makes the difficulty distribution of your data not a detail of the pipeline but arguably its most important hyperparameter. Feed it problems it has already mastered and you will burn a week of compute on text that teaches nothing, while the reward curve sits reassuringly pinned near the top. Neither number on your dashboard will look broken. The industrial fix is as blunt as it sounds: keep sampling a prompt until its group contains both successes and failures, and drop the prompts that never manage both.10

And notice how the three stages fit: reinforcement could not have started from noise — a group of rollouts that are all hopeless earns identical rewards, and as Plate VI showed, a group of equals has no gradient. Pretraining built enough competence for the judge to have something to grade; fine-tuning aimed it; only then could verified reward pull it the rest of the way. Each act feeds the next. That ordering is not a convention — it is the reason the pipeline has this shape.

Step back and look at what you ran. One set of weights was pretrained on everything available, fine-tuned on a curated slice, then reinforced against a verifier. Swap chess for language. Swap legality for human preference and verified answers. Swap 1.3 million parameters for a trillion, and a browser tab for a datacenter. The recipe you just ran is the recipe. The ideas do not change — only the appetite. And the appetite is not a hand-wave: this same text-only treatment of chess, given enough of both, plays grandmaster-level blitz without searching a single move ahead.11

And that is every machine in this book that answers a question. A loss surface and a step downhill; a neuron, a bump of influence; space bent until classes come apart; handwriting read out of 784 numbers and no picture; a map that draws itself; a game of guess-the-next-word; a policy beating to windward with nothing to go on but a clock; a judge fitted to your own eye and then broken by an optimizer — and here, all of it at once, teaching a pocket of numbers to play a game whose rules nobody told it. None of it stayed mysterious once you could watch the gradients move. That was the point.

Two chapters are left, and they are about the other thing a machine can do with what it has learned. Everything so far was handed something and asked about it. What remains is handed nothing at all.

Sources

  1. 1 Tülu 3: Pushing Frontiers in Open Language Model Post-Training Lambert et al. · 2024 · arXiv preprint · arXiv:2411.15124 Where the phrase reinforcement learning from verifiable rewards is coined, in the course of writing down a full post-training recipe openly enough to reproduce.
  2. 2 DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning Guo et al. · 2025 · arXiv preprint · arXiv:2501.12948 Verifiable reward applied to mathematics and code at full scale, with almost no demonstrations — and the model grows longer deliberations on its own, because longer ones score better.
  3. 3 Chess as a Testbed for Language Model State Tracking Toshniwal et al. · 2022 · AAAI 2022 · arXiv:2102.13249 Exactly this chapter's setup, studied on purpose: a transformer fed nothing but move strings, asked afterwards where the pieces are.
  4. 4 Emergent World Representations: Exploring a Sequence Model Trained on a Synthetic Task Li et al. · 2023 · ICLR 2023 · arXiv:2210.13382 Othello-GPT. Train on move lists alone, then read the board out of the activations — and, decisively, edit the board in the activations and watch the model's moves change to match.
  5. 5 Emergent Linear Representations in World Models of Self-Supervised Sequence Models Nanda, Lee & Wattenberg · 2023 · arXiv preprint · arXiv:2309.00941 The follow-up that changed what the Othello result means. Li et al. needed a probe with a bend in it; ask each square 'mine or theirs' instead of 'black or white' and a straight line finds the board. The model's picture of the world is relative to whoever is about to move — which is how this chapter's probe reads it too.
  6. 6 Emergent World Models and Latent Variable Estimation in Chess-Playing Language Models Karvonen · 2024 · COLM 2024 · arXiv:2403.15498 The same experiment in chess, on models trained the way Rook is. The board is recoverable from the residual stream — and so, uncomfortably, is an estimate of how strong the player being imitated is.
  7. 7 Training Language Models to Follow Instructions with Human Feedback Ouyang et al. · 2022 · NeurIPS 2022 · arXiv:2203.02155 InstructGPT — the pipeline in this chapter, run at product scale, and the reason the assistant you have talked to answers instructions rather than continuing them.
  8. 8 DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models Shao et al. · 2024 · arXiv preprint · arXiv:2402.03300 Where GRPO is introduced: drop the value network and let a group of samples from the same prompt be each other's baseline.
  9. 9 Monitoring Reasoning Models for Misbehavior and the Risks of Promoting Obfuscation Baker et al. · 2025 · arXiv preprint · arXiv:2503.11926 Models trained against test suites learn to defeat the test suites — and punishing them for saying so out loud teaches them to hide it rather than to stop.
  10. 10 DAPO: An Open-Source LLM Reinforcement Learning System at Scale Yu et al. · 2025 · arXiv preprint · arXiv:2503.14476 The industrial answer to the dead group this chapter runs into: keep sampling until a prompt produces both successes and failures, and throw away the prompts that cannot.
  11. 11 Amortized Planning with Large-Scale Transformers: A Case Study on Chess Ruoss et al. · 2024 · NeurIPS 2024 · arXiv:2402.04494 How far the text-only approach goes when you stop being polite about scale: grandmaster-level blitz from a transformer with no search at all.