Chapter 17 of 0–21
Progress measures and the three phases
Restricted loss, excluded loss, Gini and weight norm reveal progress hidden in the plateau.
We now know what the trained network computes. The paper's main contribution is to use that knowledge to look back at training and see progress that the loss curves hide.
Progress measures
Barak et al. (2022) proposed looking for hidden progress [20]; Nanda et al. define a progress measure as one of the [1]
“metrics that precede and are causally linked to the phase transition, and which vary more smoothly”— Nanda et al. 2023, Sec. 1, p. 1 (defining hidden progress measures, after Barak et al. 2022) · source
Knowing the final algorithm suggests two such measures. Both act on the logits of all 12,769 inputs, using the run's key frequencies 𝒦 (for the paper's run, those of the step-40,001 model):
- Restricted loss: keep only the cos/sin(w_k(a+b)) directions of the logit grids (10 directions for 5 key frequencies) plus each answer's average logit, and measure the loss on all inputs. It asks: how good is the network's Fourier-multiplication part on its own?
- Excluded loss: remove those directions and measure the loss on the training inputs. It asks: how much of the training performance would survive without the Fourier circuit? The idea, in the paper's words:
the memorizing solution should be spread out in the Fourier domain
(Nanda et al. 2023, Sec. 5.1, p. 8 · source)
Both definitions follow the paper's code exactly [6], and ARENA's reference implementations [4] reproduce to every digit we checked (see Reproduce).
Progress measures: restricted and excluded loss
Train and test loss with the paper's two progress measures computed at every checkpoint.
Loading…
How to read it · the math
restricted loss (green) uses only the key-frequency parts of the logits; excluded loss (purple) is the training loss after removing them.
What to expect: Restricted loss falls long before test loss does; excluded loss rises during circuit formation — the generalising circuit is built while the test loss still looks hopeless.
Reading the plot (paper's run): at step 5,001 the test loss is still 21 — worse than guessing — but the restricted loss is already down to 1.40: the Fourier part alone would already do far better than the whole network on unseen inputs. Meanwhile the excluded loss climbs from 6.94×10⁻⁵ at the end of memorisation to 8.10 at step 9,301: by then, without its Fourier circuit, the network could no longer even fit its training data — the memorised answers have been replaced. The restricted loss falls to a tenth of its value at the end of memorisation by step 8,101; the test loss gets there only at step 9,701 (the rule of hypothesis H1, chapter 20).
Three phases
From these curves the paper splits training into three phases [1] — it split training into three continuous phases: memorization, circuit formation, and cleanup
(Nanda et al. 2023, abstract, p. 1 · source)
The paper drew the boundaries by eye (1.4k, 9.4k, 14k). To compare runs fairly we defined objective rules in advance (their constants calibrated once on the paper's run, then frozen before our runs were analysed), using only the train/test loss and the weight norm (so they cannot be circular with respect to the progress measures): see docs/operationalizations.md. On the paper's run they give 1,567, 9,301 and 14,001.
- Memorisation (0 – 1,567): train loss and excluded loss fall together, test and restricted loss stay high: the network memorises without using the key frequencies.
- Circuit formation (1,567 – 9,301): excluded loss rises, restricted loss falls, the weights shrink, while the train loss barely moves and the test loss falls only gradually, from about 26 to 4.85 (still above the guessing level).
Notably, the circuit is formed well before grokking occurs.
(Nanda et al. 2023, Sec. 5.2, p. 9 · source) - Cleanup (9,301 – 14,001): the test loss suddenly drops, the weights shrink fast and the Fourier spectra become sparse:
weight decay encourages the network to shed the memorized solution in favor of focusing on the Fourier multiplication circuit
(Nanda et al. 2023, Sec. 5.2, p. 9 · source)
The test/restricted-loss ratio shows the “lag” between having the circuit and using it cleanly, which ARENA highlights [3]:
Test loss ÷ restricted loss
How much worse the full model is on the test set than its key-frequency part alone.
Loading…
How to read it
<p class="fig-expect"><strong>What to expect:</strong> On the paper's run <span class="num" data-num="paper.ratio_cf_range">9–356</span> during circuit formation, then far larger around grokking: above 10⁴ at steps <span class="num" data-num="paper.ratio_1e4_steps">9,901–10,701</span>, peaking at <span class="num" data-num="paper.ratio_peak">1.91×10⁴</span> at step <span class="num" data-num="paper.ratio_peak_step">10,201</span>. The key-frequency part already works while the full model still carries memorisation noise, which cleanup removes.</p>
<p class="fig-meta">
<a href="../chapters/17-progress-measures.html">Learn step by step →</a>
· Source: <a class="cite" href="../references.html#arena152" title="ARENA 1.5.2">[3]</a> section 3, “Small lag to clean up noise”
· Computed by <code>checkpoint_metrics.py</code> (<a href="https://github.com/BurnyCoder/visualizing-grokking-modular-addition">code</a>)
</p>
Sparsity and weight norm
Two more measures track the same story. The Gini coefficient coefficient of the Fourier spectra of W_E and W_L (chapter 11) measures how few frequencies they use:
Sparsity of the weights in the Fourier basis (Gini)
Gini coefficient of the 113 Fourier-component norms of W_E and of W_L at every checkpoint.
Loading…
How to read it · the math
<p class="fig-expect"><strong>What to expect:</strong> Low and flat during memorisation, rising during circuit formation and sharply during cleanup.</p>
<p class="fig-meta">
<a href="../chapters/17-progress-measures.html">Learn step by step →</a>
· Source: <a class="cite" href="../references.html#nanda2023" title="Nanda et al. 2023">[1]</a> Fig. 7; <a class="cite" href="../references.html#hurley2009" title="Hurley & Rickard 2009">[15]</a>
· Computed by <code>fourier.gini</code> (<a href="https://github.com/BurnyCoder/visualizing-grokking-modular-addition">code</a>)
</p>
and the sum of squared weights — the quantity weight decay pushes down — rises early in memorisation to a peak of 3,658 (step 901), is already down to 3,198 when memorisation ends, keeps falling through circuit formation, and falls fastest on average during cleanup, to 995 at step 14,001:
Sum of squared weights
Σθ² over all parameters (bold) and per parameter matrix (toggle in the legend).
Loading…
How to read it
<p class="fig-expect"><strong>What to expect:</strong> Rises during memorisation, falls steadily during circuit formation and steeply during cleanup — the fingerprints of weight decay.</p>
<p class="fig-meta">
<a href="../chapters/17-progress-measures.html">Learn step by step →</a>
· Source: <a class="cite" href="../references.html#nanda2023" title="Nanda et al. 2023">[1]</a> Fig. 7; <a class="cite" href="../references.html#arena152" title="ARENA 1.5.2">[3]</a> section 3
· Computed by <code>metrics.sum_sq_weights</code> (<a href="https://github.com/BurnyCoder/visualizing-grokking-modular-addition">code</a>)
</p>
Per frequency, the excluded loss shows each key frequency taking over part of the job:
Excluded loss, one key frequency at a time
Training loss after removing a single key frequency's cos/sin(w_k(a+b)) components from the logits.
Loading…
How to read it
<p class="fig-expect"><strong>What to expect:</strong> Each curve rises above the training loss during circuit formation: more and more of the training performance flows through that frequency.</p>
<p class="fig-meta">
<a href="../chapters/17-progress-measures.html">Learn step by step →</a>
· Source: <a class="cite" href="../references.html#nanda2023" title="Nanda et al. 2023">[1]</a> Fig. 15; <a class="cite" href="../references.html#arena152" title="ARENA 1.5.2">[3]</a> section 3, Excluded loss
· Computed by <code>metrics.excluded_loss_per_freq</code> (<a href="https://github.com/BurnyCoder/visualizing-grokking-modular-addition">code</a>)
</p>
Computing the restricted loss separately on the training and the test questions shows why it measures hidden progress: the two curves nearly coincide at every step, so the part of the logits that uses only the key frequencies generalises from the start, while the ordinary test loss is still dominated by memorised answers.
The restricted loss on training and test pairs
The restricted loss (logits kept only along the key frequencies) on training pairs and on test pairs separately, the ordinary loss on all pairs, and the excluded loss of the worst single frequency.
Loading…
How to read it
<div class="fig-read"><p>The restricted loss keeps only the key-frequency part of the logits (docs/operationalizations.md); computing it on training and on test pairs separately shows whether that part generalises. The dotted line is the largest training loss after removing any one frequency (the paper's K-free excluded loss, footnote 3).</p>