Train a transformer in your browser

A real transformer, one layer and one attention head, learns to reverse and to sort five digits while you watch. Within a few hundred steps its attention goes from noise to a clean pattern you can read, and that pattern is exactly the algorithm that solves the task.

Key idea: A training loop (batch, forward, loss, backward, update) on a tiny attention model, and the attention pattern that emerges
Real world: A new shelf-stacker learns to put tins in date order. At first they grab at random; after a few corrections they learn the rule “find the oldest one left”, and you can see the rule in where their eyes go.
In an LLM: GPT-2 was trained the same way, on text instead of digits. Its heads learned patterns such as “look at the previous token” and “look at what followed this token last time”. Nobody wrote them; they emerged because they lower the loss.

The task

Each problem is five different digits, a separator “›”, and then the answer. Two tasks:

Reverse: 3 8 1 6 4 › 4 6 1 8 3
Sort: 7 2 9 4 0 › 0 2 4 7 9

The model reads the sequence left to right and, like any language model, predicts the next token at every position. We only score it on the five answer digits. At test time it writes the answer one digit at a time, feeding each guess back in, exactly as an LLM writes a reply (the generation loop).

Neither task can be done by looking only at the current token. To write the first answer digit, the position of “›” has to fetch a digit from somewhere earlier in the sequence. Moving information between positions is the one thing attention does, so these tasks need it.

The model

This is the smallest thing that is still a transformer:

Embed. Each token gets a learned vector of 16 numbers, plus a learned vector for its position (0 to 9). Added together, they make the row of $X$ for that token.
Attend. One attention head, as in queries, keys and values: $Q = XW_Q$ , $K = XW_K$ , $V = XW_V$ , causal mask, softmax, $Z = AV$ .
Add. The head’s output, mapped back by $W_O$ , is added to $X$ : the residual stream.
Unembed. A last matrix turns each position’s 16 numbers into 11 logits, one per token (0–9 and “›”), and softmax gives probabilities.

That is 1,547 numbers in total. GPT-2 small has 124 million and Llama 3 70B has 70 billion, but the pieces are the same, just wider and stacked many times, with an MLP and normalization in every block.

Watch it learn

Press Train. Every frame the page runs a few training steps, then redraws.

learning rate0.003speed30 steps/s

step0

loss–

solved (of 100, step 0)0%

startseed 4

loss (cross-entropy, smoothed)random guessing (2.40)

test problem

38164›8›876

want 4 6 1 8 3 · 0/5 right

outline the right answer

attention weightsanswer positions × tokens

38164›8›87›0.10.10.20.10.30.20.00.00.00.080.20.20.10.30.10.10.20.00.00.0›0.10.10.20.20.10.10.10.10.00.080.10.20.10.10.10.10.10.10.10.070.10.10.10.10.20.10.10.10.10.1

row = token whose position writes the next answer digit · column = token it reads

The model has random weights. Its attention is spread out (the outlined cells are where it should look), its answers are guesses, and its loss would be about 2.40, the same as guessing evenly among 11 tokens. Press Train.

A real transformer (one layer, one attention head, 1,547 numbers to learn) training in your browser. Press Train and watch the loss fall, the attention settle on the right digits, and the answers turn green.

Each step of training is the loop from the previous pages, written out by hand in about 200 lines of TypeScript:

Batch. Make 16 fresh random problems.
Forward. Run the model on each and compute the cross-entropy loss on the answer digits: minus the log of the probability it gave the right digit.
Backward. Backpropagation works out how the loss changes with each of the 1,547 numbers: the gradient.
Update. Nudge every number a little against its gradient. The size of the nudge is the . We use the Adam optimizer, which adapts the step size for each number separately.

Things to try:

Compare the two tasks. Reverse is solved in about 150 steps. Sort takes several hundred, and its loss often drops in stages.
Turn the learning rate up to 0.1 and reset. Steps that are too big overshoot: the loss bounces or stalls. Turn it down to 0.0003 and learning crawls.
Press New start a few times. A different random start gives a different run, and now and then sorting gets stuck for a long time. Real training runs have the same problem, which is why their settings are tuned with care.
Press New problem once it is trained. The model has never seen that exact problem (there are 30,240 possible ones), yet it gets it right: it has learned the rule, not memorised answers.

example · the shelf-stacker

A new shelf-stacker is told “put these tins in date order” and gets nothing else: no rule, just a manager who says “wrong” or “right”. At first they grab tins at random. Each correction nudges how they work. After a while you notice where their eyes go before each pick: they scan for the oldest tin still on the trolley. Nobody told them that rule; it is simply what stopped the corrections. The tiny transformer is in the same position, with the loss as the manager.

Why that pattern solves it

Look at the attention matrix once the model is trained. Each row is a position that has to write an answer digit; each column is a token it can look at.

Reverse: an anti-diagonal. The position of “›” looks at position 4, the next position at 3, then 2, 1 and 0. Whatever digit sits there is copied, through the values, into the answer position, and the unembedding reads it off. The rule only depends on positions, never on which digits are present. Switch the view to Position scores to see it inside the head: the query from each answer position, compared with the key of every input position using only the position embeddings, is largest exactly on the anti-diagonal.

Sort: “the smallest digit bigger than the one I just wrote”. The first answer position (“›”) has to find the smallest digit. After that, the position holding the digit it just wrote (say 2) must find the smallest digit larger than 2. That is a rule about content, so the outlined cells jump around from problem to problem. Switch to Digit scores: for most digits, the query–key score is low for smaller digits, highest at or near the next digit up, and lower again for larger ones. (Position embeddings add a little on top, which is why a few rows are messier.) Softmax then picks out the best match that is actually present. One matrix product, $W_Q W_K^\top$ , has learned a comparison operator.

The same model with the same code learned two different algorithms, because the data was different. That is the point of training: the architecture only gives the model ways to move information around, and gradient descent finds the way that lowers the loss.

example · an induction head in GPT-2

Real language models learn the same kind of thing at scale. When GPT-2 small reads “The quick brown fox jumps over the lazy dog” twice, heads 5.1 and 5.5 learn to look, from each token in the second copy, at the token that followed the same word in the first copy, and so predict the rest of the sentence. Like our sorting head, the pattern is decided by content (“where did this token appear before?”) and nobody programmed it. See real attention patterns.

What this toy leaves out

Depth. One layer can only do one round of looking. Many tasks need two: GPT-2’s induction heads rely on an earlier head that first looks at the previous token. Real models stack 12 to over 100 blocks.
The MLP and normalization. Each real block also has an MLP, where much of a model’s knowledge is stored, and normalization layers that keep training stable.
Scale. Our model sees 32,000 problems in its 2,000 steps. Llama 3 was trained on about 15 trillion tokens, on thousands of GPUs for months. The steps are the same; for the largest Llama 3 model the arithmetic is about 10¹⁶ times bigger. See what a forward pass costs.

going deeper

The exact model

With $d = 16$ , $T = 10$ positions and a vocabulary of 11:

X = E[\text{tokens}] + P, \qquad A = \operatorname{softmax}\!\Big(\frac{XW_Q (XW_K)^\top}{\sqrt{d}} + M\Big), \qquad H = X + A\,XW_V\,W_O, \qquad \text{logits} = HU + b

Parameters: $E$ 11×16, $P$ 10×16, $W_Q, W_K, W_V, W_O$ 16×16 each, $U$ 16×11, $b$ 11: 176 + 160 + 1,024 + 176 + 11 = 1,547. Loss: mean cross-entropy over the five answer positions (teacher forcing: during training the model is given the correct earlier answer digits). Optimizer: Adam, $\beta_1 = 0.9$ , $\beta_2 = 0.98$ , batch 16, default learning rate 0.003, seeded initialisation (embeddings $\mathcal{N}(0, 0.5^2)$ , weight matrices $\mathcal{N}(0, 1/d)$ , scaled down for the unembedding so the first loss is close to $\ln 11 \approx 2.40$ ).

Checking the backward pass

The gradients are written by hand (no autodiff library). They were checked against finite differences, $\big(L(\theta + \epsilon) - L(\theta - \epsilon)\big) / 2\epsilon$ with $\epsilon = 10^{-5}$ , for all 1,547 parameters: the largest error was about $3 \times 10^{-9}$ of the largest gradient in the same matrix. On ten seeded starts at the default settings, reverse solved every one of 100 fresh problems by step 200, every time. Sort solved at least 95 of 100 by step 900 on nine of the ten starts; the tenth was still at 77 of 100 after 2,000 steps.

The softmax step of the backward pass

For a row of attention weights $a = \operatorname{softmax}(s)$ and incoming gradient $g = \partial L / \partial a$ :

\frac{\partial L}{\partial s_j} = a_j \Big(g_j - \sum_k a_k g_k\Big)

then divide by $\sqrt{d}$ and pass it to $Q$ and $K$ : $\partial L/\partial Q = (\partial L/\partial S)\,K$ and $\partial L/\partial K = (\partial L/\partial S)^\top Q$ .

In papers

Transformers learning algorithmic tasks, and reading the circuits they form: Elhage et al., 2021, A Mathematical Framework for Transformer Circuits and Olsson et al., 2022, In-context Learning and Induction Heads.
Sorting and reversing as test tasks for attention: Weiss et al., 2021, Thinking Like Transformers.
Adam: Kingma and Ba, 2014.
A minimal transformer trained from scratch, for reading: Karpathy’s nanoGPT.