A transformer trained with Muon

What this notebook trains

A transformer reads a sequence of tokens (symbols from a fixed list, the vocabulary) and, at every position, predicts a token. What lets it use the rest of the sequence is attention: each position computes a weighted average over the positions before it, and the weights are learned.

This one is as small as a transformer gets: one layer, one attention head, no feed-forward block, 336 numbers in all. Everything runs in this browser tab, forward and backward, with no machine-learning framework. The layers come from lib/nn and their backward passes, written by hand, from lib/nn_adjoints, the bookkeeping of the parameters from lib/params, and the optimizer, Muon, from lib/optim. Muon is the optimizer used in the modded-nanogpt speedrun, a public contest to train a small GPT as fast as possible.

The task. Sequences of six tokens from a vocabulary of six. At every position the model must name the token before it. The tokens are random, so the token at a position says nothing about the one before it, and a model that cannot look at other positions can do no better than guessing. (Position zero is the one exception; the section on what guessing scores comes back to it.)

The first cell sets the sizes. D is the width of the vector that stands for each token inside the model, H the width of the attention head's vectors, and MU, LR and STEPS are the optimizer's settings, explained where they are used.

The data

Sixteen sequences of six random tokens. The random numbers come from a linear congruential generator: multiply by a large constant, add another constant, keep the remainder. It starts from a fixed seed, so every run trains on the same data and prints the same numbers. The target at each position is the token one step back; position zero has nothing before it, so it names itself.

The targets are the tokens shifted one place to the right, with the first token repeated at the front.

What guessing scores

At each position the model outputs six numbers, one per token, called logits. A softmax turns them into probabilities: exponentiate each and divide by the sum. The loss is the cross-entropy, $-\log p$ where $p$ is the probability given to the right token, averaged over every position. It is $0$ for a model that puts all its probability on the right answer, and it grows without limit as that probability goes to zero.

A model that knows nothing about the answer does best by spreading its probability evenly, $p=1/6$ for every token. Before running: what loss is that?

$-\log (1/6)=\log 6=1.792$ (chance). That is the loss to beat. A model that sees only the current token could do slightly better than this by getting position zero right, since its answer is its own token, but it would still be guessing at the other five positions out of six. Only attention to the previous position gets the loss near zero.

The parameters

Six weight matrices, all learned:

• the token embedding E, $6\times 8$: a row of eight numbers for each token;
• the position embedding P, $6\times 8$: a row for each position;
• the attention head's three projections Wq, Wk and Wv, $8\times 8$ each;
• the output projection Wo, $8\times 6$, which turns the head's output into six logits.

That is $48+48+3\times 64+48=336$ numbers. lib/params keeps all six in one flat vector, so the gradient is one vector too and the optimizer's memory is two more. Each matrix starts as the corner of a fixed pseudo-random $8\times 8$ matrix, scaled down so that the first logits are small and the first guesses near even.

The forward pass starts with the embedding

The model's input at position $t$ is a row of eight numbers: the token embedding's row for the token there, plus the position embedding's row for $t$. The position row is the only place the model learns where a token sits. Without it, the six positions would look the same to the attention head.

The forward pass continues through one attention head

Each position's row $x_t$ is multiplied by three matrices to make a query $q_t=x_t{W}_q$, a key $k_t=x_t{W}_k$ and a value $v_t=x_t{W}_v$. Position $i$ scores every position $j$ by how well its query matches that key, $q_i\cdot k_j/\sqrt{H}$. The scores of positions after $i$ are set to $-\infty$ (the causal mask: a position may not look ahead). A softmax along each row turns the scores into weights $A_{ij}$ that are zero or positive and add up to 1, and the head's output at $i$ is the weighted average ${\sum }_j{A}_{ij}{v}_j$.

The attention matrix $A$ is where the model's looking shows. To solve the task, row $i$ has to put its weight on column $i-1$. Before running: with the starting random weights, where does each row look?

Everywhere it may, evenly. Row 0 can see only itself, so its weight is 1. Every later row spreads its weight almost equally over the positions it can see: row 1 gives $0.50$ to each of two, row 5 about $0.167$ to each of six (attention_before). There is no preference for the position before. Everything above the diagonal is exactly zero: that is the mask.

The forward and backward pass for one sequence

The head's output times Wo gives the logits, and cross-entropy scores them against the targets. That is the forward pass.

The backward pass computes the gradient, the derivative of the loss with respect to each of the 336 parameters, by the chain rule. It runs the forward steps in reverse. Each _bwd function from lib/nn takes the gradient with respect to its step's output and turns it into gradients with respect to that step's inputs and weights: cross-entropy first, then the output projection, then attention. An embedding's gradient goes only to the rows that were looked up.

(Viv can also write this pass itself: value_and_grad differentiates the forward pass, and lib/nn's tests check its gradients. This notebook keeps the pass written by hand. The two gradients differ in the last digit, and sixty Muon steps at a constant step size carry that difference into the sixth digit of constant_at_end below.)

Loss and accuracy over the training set

The training loss is the mean over all sixteen sequences, and the gradient is the sum of theirs. The accuracy is the fraction of the 96 positions where the token the model rates most likely is the right one. Before running: the untrained model's loss should be close to $1.792$, and its accuracy close to $1/6=0.167$, the rate at which a guess is right by luck.

The loss starts at $1.785$ (start_loss), a hair under chance. The accuracy starts at $0.146$ (start_accuracy), 14 of the 96 positions, about what luck gives.

The optimizer, Muon

Plain gradient descent steps each weight against its gradient. Momentum keeps a running sum instead, $M\leftarrow \mu M+G$ with $\mu$ = MU = 0.95, and steps along $M$, so a direction the gradient keeps pointing in builds up speed and one it flips back and forth on cancels.

Muon adds one thing. It treats each weight matrix's momentum as a matrix and, before stepping, replaces it with a matrix that has the same singular vectors (the directions the matrix acts along) but every singular value (how strongly it acts along each) close to 1. Then no direction of the weight matrix gets a much larger step than any other, however lopsided the gradient was. It does this with five steps of a Newton-Schulz iteration, a matrix polynomial,

which acts on each singular value separately and pushes it toward 1 without computing any of them. The classical Newton-Schulz polynomial, with coefficients $15/8$, $-10/8$ and $3/8$, converges to exactly 1 but slowly; Muon's coefficients are tuned to get close fast.

The cell below builds a $5\times 8$ matrix whose singular values run from $1$ down to $0.001$, a spread of a thousand, and runs five steps of each polynomial on it. Before running: after five steps, how far apart are the largest and smallest?

Muon's five steps leave the largest singular value $2.60$ times the smallest (muon_spread); the classical polynomial's five leave $45.1$ (classical_spread). Muon's values are not all 1; they run from $0.45$ to $1.17$ (after_muon), and they do not need to be 1: what the step needs is that no direction dominates.

Training

A training step computes the loss and gradient over all sixteen sequences, then takes one Muon step. Muon's step length is set by the step size LR, not by the size of the gradient, so near the minimum a step of fixed length can walk over it. For that reason the step size falls in a straight line from LR = 0.12 to a tenth of that over the run. The function below takes a decay switch, 1 for that schedule and 0 for a constant step size, so the two can be compared later.

Before running: after sixty steps, will the loss be below chance, near zero, or somewhere between? And how many of the 96 positions will be right?

Near zero. The loss is $1.785$ at the start (loss_at_start), $0.904$ after ten steps (loss_at_step_10), $0.0072$ after thirty (loss_at_step_30) and $0.00058$ at the end (loss_at_end), a factor of three thousand from start to end. The accuracy is 1 (got): all 96 positions right.

The figure plots the loss on a logarithmic scale, against the line that guessing gets.

It does not get there smoothly

After step thirty the loss does not slide down steadily. It sits at a few thousandths and then jumps. The cell below lists every step after thirty where the loss is above a hundredth, and how many times larger it is than one step earlier.

Three steps: the loss jumps to $0.046$ at step 33, $0.016$ at step 39 and $0.133$ at step 51, the last about seventy times its value one step earlier (times_the_step_before). None of the three lasts: the step after each is not on the list, so the loss is back under a hundredth one step later. These are the overshoots described above: near a minimum, a step of fixed length can land somewhere much worse, and the momentum brings it back.

What the shrinking step size buys

Before running: the same sixty steps from the same starting weights, but with the step size held at LR the whole way. Better or worse at the end?

Worse. With a constant step size the loss is $0.0026$ at step thirty (constant_at_30), lower than the schedule's $0.0072$ there, but it ends at $0.047$ (constant_at_end), eighty times the schedule's $0.00058$. Full-length steps get down fast and then keep stepping over the minimum.

Where it ended up looking

Nothing in the setup mentions the previous position. The model gets a position embedding and a loss. Before running: has the head learned to put its attention one position back, and how much of it?

All of it, to three decimal places: averaged over the five rows that have a previous position, the weight on that position is above $0.999$ (the test; back_mass prints it rounded). The picture is the attention matrix: a bright line just below the diagonal, position zero's weight on itself in the top corner, and zeros above the diagonal, which are the causal mask, not something learned.

Every backward pass used here is checked in lib/nn_adjoints against the gradient value_and_grad generates, which lib/nn checks against central differences of the forward pass (nudge each input up and down a little and see how the output moves), including the four in attention. A subtly wrong backward pass still runs and can still make the loss fall for a while, so training alone would not reveal the error.