Grokking, Visualized
Chapters
  1. The whole picture, first
  2. Counting on a clock
  3. The task as a dataset
  4. Vectors and matrices
  5. From scores to probabilities
  6. The model, step by step
  7. Simplifying the model
  8. How the model learns
  9. Grokking
  10. Angles and waves
  11. The Fourier basis
  12. Measuring fit and sparsity
  13. The Fourier multiplication algorithm
  14. Evidence 1: everything is periodic
  15. Evidence 2: neurons and heads
  16. Evidence 3: reading off the answer
  17. Evidence 4: interventions
  18. Progress measures and the three phases
  19. Watch it form
  20. Gradients and weight decay
  21. Is it robust? Seeds and controls
  22. The whole picture, again

Chapter 20 of 0–21

Is it robust? Seeds and controls

22 retrained networks, preregistered tests, and what changes from seed to seed.

Sources for this chapter
  • docs/preregistration.md
  • [1] App. C.2, D.1

Everything so far came from one trained network. The paper reports four more seeds and many other settings in its appendix and finds the same algorithm with different key frequencies [1] (App. C.2). We retrained the paper's architecture 22 times to test the main claims with pre-specified statistics.

Testing claims, not stories

A hypothesis test asks how surprising the data would be if the claimed effect were absent. The answer is a p-value: the probability of results at least this extreme under “no effect”; small values are evidence for the effect. A confidence interval gives a range of plausible values for a quantity; a 95 % interval is built so that, over many repetitions of the experiment, 95 % of such intervals would contain the true value.

Our main tool is the sign test: if each of n independently trained networks either shows an effect or not, and “no effect” would make each outcome a coin flip, then all n showing it has probability 0.5^n. Testing several hypotheses at once raises the chance of a false alarm, so we apply a multiple-comparison correction (Holm's method) to the five tests.

To avoid choosing tests after seeing the results, we wrote the hypotheses, metrics and pass/fail thresholds down before training — a preregistration, frozen in the repository as git tag prereg-v1 (docs/preregistration.md). The number of runs was chosen for statistical power: with 11 seeds, even one disagreeing seed (10 of 11) still passes the strictest corrected threshold, and a 95 %-reliable effect is detected with 90 % probability.

Worked example

Why 11 seeds and not 5? With 5 seeds the smallest possible sign-test p-value is 0.5^5 = 0.031 — not small enough for the strictest Holm threshold (0.05/5 = 0.01), even if all five agree. With 11 seeds, 11/11 gives 0.5^{11} = 0.00049 and 10/11 gives 0.0059 — both below 0.01.

The runs

Group Runs Weight decay Task Purpose
wd1 11 (seeds 0–10) 1 addition the paper's setting; seed 0 uses the paper's exact data split
wd0 5 0 addition control: no weight decay
rand 3 1 random answer table control: nothing to generalise
split0 3 1 addition the paper's data split with new initial weights

All runs use the paper's architecture, optimiser and initialisation rule for 40,000 full-batch steps (chapter 7).

Every run's test loss

Test loss of the paper's run and all our runs (seeds and controls) on one log axis; the selected run is highlighted.

Loading…

How to read it
<p class="fig-expect"><strong>What to expect:</strong> All weight-decay runs grok, at different times; no-weight-decay and random-label runs never do.</p>
<p class="fig-meta">
  <a href="../chapters/20-robustness.html">Learn step by step →</a>
  
  · Computed by <code>train.py</code> (<a href="https://github.com/BurnyCoder/visualizing-grokking-modular-addition">code</a>)
</p>

The phases of every run on one time axis, and every weight-decay seed's clock at the same moment (it plays while on screen): different seeds grok at different times and with different frequencies, but in the same order and into the same circle.

Every run's phases on one time axis

Memorisation, circuit formation and cleanup of every run whose phases were detected, with the step at which test accuracy passes 99 %; on the paper's row, ticks mark the paper's own boundaries.

Loading…

How to read it
<div class="fig-read"><p>Phases come from the objective rules of docs/operationalizations.md applied to each run's curves. The vertical cursor is the selected step; click a row to select that run.</p>

What to expect: The weight-decay seeds grok at different times, but always in the same order: memorisation, a long circuit formation, then cleanup.

Learn step by step → · Source: [1] Sec. 5.1 · Computed by phases.py (code)

The clock in every seed

Every weight-decay seed (and the paper's run) at the selected step: the 113 number embeddings on the plane of the run's first key frequency.

Loading…

How to read it
<p class="fig-expect"><strong>What to expect:</strong> Different seeds choose different key frequencies, but every one of them ends with the numbers on a circle in clock order.</p>
<p class="fig-meta">
  <a href="../chapters/20-robustness.html">Learn step by step →</a>
   · Source: <a class="cite" href="../references.html#nanda2023" title="Nanda et al. 2023">[1]</a> App. C.2
  · Computed by <code>geometry.key_plane_coords</code> (<a href="https://github.com/BurnyCoder/visualizing-grokking-modular-addition">code</a>)
</p>

Grokking is not always the end of the story. In 2 of our 14 runs that grokked, the test loss later jumps more than 100-fold above its lowest value so far, again and again, and falls back each time: split0_init102 (6 spikes; test loss up to 0.64, test accuracy down to 87 %); wd1_s4 (4 spikes; test loss up to 0.25, test accuracy down to 93 %). Similar loss spikes late in training with adaptive optimisers like Adam were described as the "slingshot mechanism" [23]; we did not investigate their cause, and none of the preregistered tests depends on them (H1 and H2 are measured before grokking, H3–H5 at step 40,000, where both runs are at 100 % test accuracy).

Results

Preregistered hypotheses

Results of the five tests written down before training (docs/preregistration.md).

Loading…

How to read it
<p class="fig-meta">
  <a href="../chapters/20-robustness.html">Learn step by step →</a>
  
  · Computed by <code>hypotheses.py</code> (<a href="https://github.com/BurnyCoder/visualizing-grokking-modular-addition">code</a>)
</p>

Each seed's value of every preregistered metric against its threshold, and the sparsity test (H3) run by run:

Each seed against each preregistered threshold

For every weight-decay seed, the value of each preregistered metric of H1, H2 and H4 against its threshold.

Loading…

How to read it
<div class="fig-read"><p>Filled dots meet the preregistered condition, hollow ones do not. H1: the test loss drops after the restricted loss (steps). H2: the excluded loss grows more than 10× during circuit formation. H4: removing any key frequency multiplies the loss by at least 100, removing any other frequency by at most 2, and keeping only the key frequencies by at most 1.5. H3 is shown in the Gini panel; H5 compares groups of runs (table above).</p>

What to expect: Every seed meets H1 and H2. H4 fails in several seeds, through condition (i) in some (one key frequency matters less than 100×) and condition (iii) in others (the key frequencies alone are not quite enough), which is why H4 is not supported.

Learn step by step → · Source: docs/preregistration.md · Computed by hypotheses.py (code)

Sparsity in the Fourier basis: start, end and the bar to clear

For every run, the Gini coefficient of the W_E (or W_L) frequency norms at initialisation and at the end, against the preregistered bar of H3.

Loading…

How to read it
<div class="fig-read"><p>The bar is max(initial Gini + 0.3, 99th percentile of the same Gini computed in 100 random orthonormal bases): sparsity has to be large and specific to the Fourier basis (docs/preregistration.md, H3).</p>

What to expect: Runs that grokked end far above the bar (0.61–0.88); runs without weight decay or with random labels end at 0.10 or less.

Learn step by step → · Source: docs/preregistration.md; [1] Fig. 7 · Computed by hypotheses.py (code)

H1 supported (11/11 seeds; Holm-adjusted p = 0.002); H2 supported (11/11 seeds; Holm-adjusted p = 0.002); H3 supported (11/11 seeds; Holm-adjusted p = 0.002); H4 not supported (6/11 seeds; Holm-adjusted p = 0.5); H5 supported (11/11 weight-decay runs vs 0/5 without grokked; Holm-adjusted p = 0.0011).

Four of the five claims pass. H4 does not: it required, in at least 10 of 11 seeds, all three of: every key frequency necessary (ablating it multiplies the loss by ≥ 100), every other frequency dispensable (≤ 2×), and all other frequencies jointly dispensable (keeping only the key frequencies costs ≤ 1.5×). Splitting the preregistered criterion into its parts (a description after the fact, not a new test): the second part held in 11/11 seeds (worst case 1.84×), the first in 9/11, the third in 8/11 (worst case 2.78×).

Nuance

Why H4 failed. Both failing parts concern which frequencies count as "key", not whether the network uses a few frequencies. Our preregistered rule calls a frequency "key" if its W_L Fourier norm is at least 25 % of the largest. Twice it included a frequency that matters much less than that seed's other key frequencies: wd1_s1: k = 2 costs 5.9× (27 % of the largest W_L norm; its other key frequencies ≥ 145,740×); wd1_s8: k = 46 costs 70.4× (58 % of the largest W_L norm; its other key frequencies ≥ 144×). So a large weight is not the same as causal importance. Of the 6 seeds with at most four key frequencies, 3 have the opposite problem: every other frequency is individually dispensable, but removing all of them together still multiplies the (already tiny) final loss by more than 1.5: wd1_s4 (3 key frequencies, 2.78×), wd1_s9 (3 key frequencies, 1.99×), wd1_s10 (4 key frequencies, 1.64×). The preregistered verdict stands: the paper's model of "a few key frequencies, all else irrelevant" is an approximately correct description of every seed, but not as exact a one as our thresholds demanded.

The phase boundaries vary from seed to seed: across the 11 weight-decay runs, memorisation ends at step 1,576 [1,567, 1,586] (mean and 95 % interval), circuit formation at 7,445 [5,874, 9,017], and test accuracy reaches 99 % at 8,455 [6,816, 10,093].

Different seeds, different frequencies, same algorithm

Each run picks its own key frequencies:

Which frequencies each run chose

Key frequencies of every run (click a run to select it everywhere).

Loading…

How to read it
<p class="fig-meta">
  <a href="../chapters/20-robustness.html">Learn step by step →</a>
  
  · Computed by <code>metrics.key_frequencies</code> (<a href="https://github.com/BurnyCoder/visualizing-grokking-modular-addition">code</a>)
</p>

The weight-decay runs use 4.364 [3.743, 4.985] key frequencies on average (mean and 95 % interval over the 11 seeds); on the same footing, single-frequency neurons make up 80.2 % [75.3 %, 85.2 %] of the MLP and the cosine formula of chapter 15 explains 96.5 % [95.3 %, 97.7 %] of the logit variance. The counterfactual patching test of chapter 16 recovers 0.997 [0.995, 0.998] of the logit difference for the key-frequency subspace versus 0.058 [0.049, 0.067] for a random subspace of the same size.

Controls

Who chooses the frequencies?

Seed 0 turned out to start from exactly the paper's initial weights (largest difference 2.98×10⁻⁸, i.e. float32 rounding) and to use the paper's data split — our code evidently draws the initial weights from the same random generator in the same order as the paper's. It is therefore a re-run of the paper's network, with our code on our GPU. The two training losses agree to five significant digits for the first 52 steps and then drift apart — tiny rounding differences are amplified by training (see the experiment report) — and by step 40,000 the two weight vectors have cosine similarity only 0.74. Yet the re-run ended with the same key frequencies: 14, 35, 41, 42, 52. The three runs with the same data but new initial weights chose different sets (above). This is an exploratory observation on one pair of runs, not a tested hypothesis, but it suggests that the initial weights, not rounding noise, decide which frequencies a network uses.

Nuance

Scope of these checks. All 22 runs share the paper's architecture, data fraction and optimiser; robustness to other settings (data fraction, depth, modulus, dropout) is reported in the paper's appendix, not re-tested here. A pre-specified threshold that a run misses by a little counts as a miss — that is the price of not tuning criteria after the fact.