The progress measures describe what changes during training. This chapter looks at the forces that change it: the gradient of the loss and the pull of weight decay. Because training is full-batch (chapter 7), the gradient the optimizer saw at any checkpoint can be recomputed from the saved weights, exactly up to floating-point rounding — which we did for every checkpoint.
Loading…
Two forces
Every AdamW step combines two pulls on the weights (chapter 7): the loss gradient, rescaled by Adam, and weight decay, -\gamma\lambda\theta, which always points straight back toward zero. The paper argues that weight decay drives both circuit formation and cleanup [1] (Sec. 5.2–5.3): the generalising circuit can reach the same training loss with much smaller weights than the memorised lookup table, so under constant pressure toward small weights it wins.
Nuance
Why we do not compare raw sizes. After memorisation the raw gradient is tiny — 1.12×10⁻⁸ to 2.09×10⁻⁸ of λ‖θ‖ in the paper's run — but Adam divides each coordinate by the root of its recent squared gradients, so even a tiny but consistent gradient produces updates of size ~γ. Raw gradient norms therefore say little about how strongly the loss pushes; directions (cosines) do, and they are what we plot.
Where the gradient points
The cosine between the loss gradient and the weights tells us how the two forces relate. A gradient-descent step moves along minus the gradient, so a cosine near −1 means the loss “wants” the weights to grow along exactly the direction weight decay shrinks them:
Where the gradient points
Cosine between the training-loss gradient and the weights, and between the training- and test-loss gradients.
Loading…
How to read it
<div class="fig-read"><p>A gradient-descent step moves along −gradient. cos(gradient, weights) ≈ −1 means the loss wants the weights to grow, exactly against weight decay's pull.</p>
What to expect: After grokking the gradient is almost exactly anti-parallel to the weights, i.e. it points against the pull of weight decay (whether Adam's step and weight decay balance is measured separately, chapter 19). On the paper's run the cosine between the train and test gradients is negative throughout circuit formation (-0.44 to -0.08) and over most of memorisation, and turns positive only during cleanup, after step 11,901.
· Computed by <code>gradients.py</code> (<a href="https://github.com/BurnyCoder/visualizing-grokking-modular-addition">code</a>)
</p>
On the paper's run, cos(gradient, weights) is -0.996 at step 40,001. Once the network is perfectly accurate, the cheapest way to lower cross-entropy is to scale all logits up (chapter 4), i.e. to grow the weights uniformly; weight decay shrinks them uniformly. The two cancel — the network sits in a balance between them, which is why its loss stops at ~10⁻⁷ instead of going to zero.
The second curve compares the gradient of the training loss with the gradient of the test loss. At the very start they are aligned (cosine 0.80): any improvement is a general improvement. The cosine is negative throughout circuit formation (-0.44 to -0.08; e.g. -0.27 at step 5,401, midway through it) and over most of memorisation: to first order, a small gradient step that lowers the training loss raises the test loss — memorisation works against generalisation. It turns positive only during cleanup and stays positive after step 11,901 (0.33 at step 40,001).
Which frequencies the gradient builds
Finally, which frequencies does the gradient push into the embedding at each moment? The share of the embedding gradient's energy per frequency, over training:
Which frequencies the embedding gradient pushes
Share of the embedding gradient's energy in each frequency (rows) at every checkpoint (columns).
Loading…
How to read it
<p class="fig-meta">
<a href="../chapters/19-gradients-weight-decay.html">Learn step by step →</a>
· Computed by <code>gradients.py</code> (<a href="https://github.com/BurnyCoder/visualizing-grokking-modular-addition">code</a>)
</p>
Size of the gradient, per parameter
Norm of the exact full-batch gradient of the training loss for each weight matrix, at every checkpoint.
Loading…
How to read it
<p class="fig-meta">
<a href="../chapters/19-gradients-weight-decay.html">Learn step by step →</a>
· Computed by <code>gradients.py</code> (<a href="https://github.com/BurnyCoder/visualizing-grokking-modular-addition">code</a>)
</p>
And the weights themselves, one line per parameter. On the paper's run most matrices grow while the network memorises and shrink through circuit formation and cleanup; the attention matrices W_K and W_Q shrink almost from the start, and the positional embedding W_pos keeps growing into circuit formation. By step 40,001 only four end smaller than they started (W_K 0.32×, W_Q 0.32×, W_in 0.55×, W_out 0.55× their initial norm), W_K and W_Q most; the others end larger (W_E 1.19×, W_pos 1.61×, W_V 1.08×, W_O 1.07×, W_U 1.08×). In total the weights still shrink: their sum of squares falls from 1,783 at initialisation to 955.
Size of the weights, per parameter
Frobenius norm (square root of the sum of squared entries) of each weight matrix and bias, at every checkpoint.
Loading…
How to read it
<div class="fig-read"><p>One line per parameter; the shading marks the three phases. Weight decay pulls every line toward 0, the loss gradient pushes back wherever a weight lowers the training loss.</p>
What to expect: On the paper's run most matrices grow while the network memorises, then shrink through circuit formation and cleanup; by the end the total size of the weights is below its value at initialisation. The attention matrices W_K and W_Q shrink almost from the start and the most; the positional embedding W_pos keeps growing into circuit formation.
· Computed by <code>gradients.py</code> (<a href="https://github.com/BurnyCoder/visualizing-grokking-modular-addition">code</a>)
</p>
What Adam actually does
The raw gradient is not the step. For our retrained runs we saved Adam's two running averages every 1,000 steps, so we can rebuild the direction Adam really moves in, \hat m/(\sqrt{\hat v}+\epsilon), and compare it with the raw gradient and with the weights:
Where Adam actually steps
Cosine between Adam's step direction and the raw gradient, and between Adam's step direction and the weights, every 1,000 steps.
Adam divides each coordinate of the gradient's running average by its running size, so its step can point elsewhere than the raw gradient. Computed from Adam's internal state, saved every 1,000 steps for our retrained runs only.
What to expect: After grokking, cos(Adam direction, weights) ≈ −1: the loss step pushes the weights outward exactly against weight decay's pull inward.
Both steps are multiplied by the same learning rate, so the ratio compares the two forces directly (unlike the raw gradient, which Adam rescales). Log scale; 1 means the two are equally large.
What to expect: The ratio settles at 1 after grokking: an equilibrium in which the loss and weight decay balance.
· Computed by <code>gradients.py (adam_direction)</code> (<a href="https://github.com/BurnyCoder/visualizing-grokking-modular-addition">code</a>)
</p>
The ratio is slightly above 1 at step 1,000 (the only saved point in memorisation), a little below 1 through circuit formation, and dips lowest right when the network groks: its minimum, 0.906 [0.897, 0.914] (mean and 95 % interval over seeds), falls within 1,000 steps of the grokking step in 10 of 11 seeds. At that moment weight decay outweighs the loss step, so the weights shrink — the cleanup that the paper attributes to weight decay, seen on the actual steps. (This pattern was noticed after the experiments, not predicted in advance, and the Adam state was saved only every 1,000 steps.)
At step 40,000, across the 11 weight-decay seeds, the loss step is 0.992 [0.972, 1.004] times the size of the weight-decay step (mean and 95 % interval), and its cosine with the weights is -0.991 [-1.000, -0.976] (bootstrap intervals). The two steps are nearly equal in size and nearly opposite in direction: the balance described above, measured on the steps that are actually taken. The raw gradient alone would not show this reliably — its cosine with the weights at the last checkpoint ranges from -0.996 to -0.157 across seeds, because Adam averages the gradient over many steps and rescales every coordinate. How much each weight matrix moves per step:
How much each parameter actually moves
Size of each parameter's AdamW update, logged every 10 steps during training (our retrained runs).
Loading…
How to read it
<div class="fig-read"><p>Log scale; one line per weight matrix. Hover to see values.</p>
<p class="fig-meta">
<a href="../chapters/19-gradients-weight-decay.html">Learn step by step →</a>
· Computed by <code>train.py (logged during training)</code> (<a href="https://github.com/BurnyCoder/visualizing-grokking-modular-addition">code</a>)
</p>
Where these forces take the weights: each run's path through (weight norm, loss) during training. With weight decay the norm first grows, then shrinks while the test loss collapses; without it the norm keeps growing and the test loss stays high.
Loss against weight norm
Each run's path through (weight norm, loss) during training, with a dot at the selected step: the selected run against the runs without weight decay.
Loading…
How to read it
<p class="fig-expect"><strong>What to expect:</strong> With weight decay the norm rises, then falls while the test loss collapses; without it the norm keeps growing and the test loss stays high.</p>
<p class="fig-meta">
<a href="../chapters/19-gradients-weight-decay.html">Learn step by step →</a>
· Source: <a class="cite" href="../references.html#nanda2023" title="Nanda et al. 2023">[1]</a> Sec. 5.3
· Computed by <code>checkpoint_metrics.py</code> (<a href="https://github.com/BurnyCoder/visualizing-grokking-modular-addition">code</a>)
</p>
Why does grokking happen? The paper's account
Combining the measurements, the paper's explanation (Sec. 5 and the speculative App. E) runs as follows [1]:
Memorisation is fast but its weight cost grows with the amount of data; the Fourier circuit has a fixed cost.
Weight decay therefore slowly favours the circuit, which grows while train loss stays flat (circuit formation).
Once the circuit carries the training data, the memorising weights are pure cost and weight decay removes them (cleanup) — only then does the test loss drop.
ARENA's bonus chapter adds intuition: when the model lacks the capacity to memorise all of its data, gradients from individual memorised examples point in unrelated directions and largely cancel, while gradients that serve the general pattern agree across examples and add up [3]. These accounts are explanations of the measurements, not proofs; the paper labels parts of its App. E as speculative.
Nuance
What we can and cannot measure here. For the paper's run we can recompute gradients but not Adam's internal running averages (they were not saved), so the size of the actual updates is unknown. For our own runs we saved Adam's state every 1,000 steps and logged each update's size, which the panels in “What Adam actually does” use; the gradient panels work for both.