Step 3 says the logits are a linear read-out of \cos(w_k(a+b)) and \sin(w_k(a+b)) with weights \cos(w_k c) and \sin(w_k c). Because the skip connection around the MLP barely matters (chapter 6), the logits are essentially W_L\,\mathrm{MLP}(a,b), so the claim is about the 113 × 512 matrix W_L.
W_L contains only the key frequencies
If \Logits_c \approx \sum_k \cos(w_k c)\,[\ldots] + \sin(w_k c)\,[\ldots], then every column of W_L, read as a function of c, must be a mix of the waves \cos(w_k c), \sin(w_k c) for k ∈ 𝒦 — ten waves in total. The Fourier transform of W_L along c:
The neuron–logit map in the Fourier basis
Length of each Fourier component (over the output c) of W_L = W_U W_out at the selected step.
Loading…
How to read it · the math
\big\lVert (F W_L)_{i,:}\big\rVert_2
<p class="fig-expect"><strong>What to expect:</strong> Only the key frequencies survive, ~10 non-zero bars: W_L is approximately rank 10.</p>
<p class="fig-meta">
<a href="../chapters/15-evidence-readoff.html">Learn step by step →</a>
· Source: <a class="cite" href="../references.html#nanda2023" title="Nanda et al. 2023">[1]</a> Fig. 3 (right)
· Computed by <code>metrics.wl_fourier_norms</code> (<a href="https://github.com/BurnyCoder/visualizing-grokking-modular-addition">code</a>)
</p>
Over the whole run: how strongly each frequency's wave appears in the logits (or how much the training loss rises without it), with the selected frequency drawn below.
Every frequency's strength in the logits, over training
For all 56 frequencies at every checkpoint: how strongly cos(w_k(a+b−c)) appears in the logits, or how much the training loss rises without frequency k; below, the selected frequency over time.
Loading…
How to read it
<div class="fig-read"><p>Rows are frequencies k = 1 (bottom) to 56 (top), columns are checkpoints. The outlined column is the current step and the outlined row the selected frequency; click a cell to select both. The line below draws the selected frequency up to the current step over a faint copy of the whole run.</p>
What to expect: Only the key frequencies' rows light up, and they grow during the plateau, before the test loss drops.
Learn step by step →
· Source: [1] Fig. 14
· Computed by metrics.cos_coefficients, metrics.excluded_loss_per_freq (code)
Only the ten key-frequency components are visible: W_L has rank about 10. Approximating W_L by its projection onto those ten waves leaves a residual whose squared Frobenius norm is 0.552 % of W_L's at step 40,001 (0.520 % at the released run's last checkpoint, step 49,901); its norm itself is 7.4 % of W_L's. The paper reports a residual norm “under 0.55 %”: that is close to our squared ratio (marginally above 0.55 % at step 40,001), not to the norm ratio; we flag the wording on the errata page.
Writing the ten rows as W_L \approx \sum_{k}\cos(w_k)\,u_k^\top + \sin(w_k)\,v_k^\top (paper eq. 2) defines, for each key frequency, two directions u_k, v_k in the 512-dimensional neuron space: whatever lies along u_k is written into the logits as \cos(w_k c), whatever lies along v_k as \sin(w_k c).
The neurons hold cos(w_k(a+b)) and sin(w_k(a+b))
The read-off u_k^\top\,\mathrm{MLP}(a,b) is one number per input (a, b) — a 113 × 113 grid. The paper's Table 1 shows that this grid is almost exactly a multiple of \cos(w_k(a+b)), and v_k^\top\mathrm{MLP}(a,b) a multiple of \sin(w_k(a+b)), with 93–98 % of the variance explained [1]. On the paper's network we measure an average of 96.4 % over the ten directions (range 93.2–98.2 %). This is the MLP outputting the result of step 2.
ARENA shows the same thing from the neurons' side [3]: each neuron writes only the two waves of its own frequency into the logits — a block structure once the neurons are sorted by cluster:
W_logit in the Fourier basis, neurons sorted by cluster
Each row is one neuron's column of the neuron–logit map, written in the Fourier basis over answers c.
Loading…
How to read it
<div class="fig-read"><p>Rows are sorted by the neuron's frequency cluster (unclustered neurons at the bottom). Click a row to select that neuron.</p>
What to expect: Block structure: a neuron of cluster k writes only cos k and sin k.
Learn step by step →
· Source: [3] “W_logit in Fourier basis”; [1] Fig. 5 (right)
· Computed by browser (GrokModel) (code)
Linear terms in the neurons, none in the logits
The neurons hold more than the products: chapter 13's 3 × 3 blocks also contain the linear terms cos(w_k a), sin(w_k b), … Averaging the squared 2D Fourier coefficients over all 512 neurons shows both kinds — on the paper's network 76.2 % of the neurons' variance sits in linear terms (row and column 0 of the map):
All neurons in the 2D Fourier basis
Mean over the 512 neurons of the squared 2D Fourier coefficients of their activations (log scale).
Loading…
How to read it
<div class="fig-read"><p>Cell (i, j) = how much the neurons, on average, contain the wave (basis i in a) × (basis j in b). Row 0 (Const in a) holds the terms in b alone, cos/sin(w_k b), and column 0 those in a alone; the cells cos k × cos k, sin k × sin k, … hold the products. Compare with the same map for the logits (group E).</p>
What to expect: Bright cells only at the key frequencies — both linear terms on the edges and product terms on the diagonal blocks.
Learn step by step →
· Source: [1] Sec. 4.2, Fig. 4; [3] “Neuron activations in Fourier basis”
· Computed by geometry.spectra (code)
The same picture for the logits keeps only the products. In the logits, linear terms carry 1.0 % of the variance: the read-off W_L adds the neurons up so that their linear parts cancel, leaving cos(w_k(a+b)) and sin(w_k(a+b)) — exactly the terms that step 3 of the algorithm needs. The products of the key frequencies hold 97.0 % of the logits' variance:
The logits in the 2D Fourier basis
Size of each 2D Fourier component (in a and in b) of the logits, combined over all 113 answers (log scale).
Loading…
How to read it
<div class="fig-read"><p>Same layout as the neuron map in group C. Each answer's mean is removed first, so the Const × Const cell is empty.</p>
What to expect: The product cells of the key frequencies (cos k × cos k, sin k × sin k, cos k × sin k) are by far the brightest — on the paper's network they hold 97.0 % of the logits' variance. The faint lines along key rows and columns are small (the colour scale is logarithmic), and the neurons' large linear terms have almost cancelled.
Played through training, watch which components light up, and when.
The logits are a sum of cos(w_k(a+b−c))
Putting the pieces together, the paper fits all 113 × 113 × 113 logits with five numbers, \Logits(a,b)_c \approx \sum_k \alpha_k\cos(w_k(a+b-c)), by least squares (we add a constant). On the paper's network this simple formula explains 96.1 % of the variance of the 1,442,897 logits (paper: 95 %), with coefficients α = 5.87, 5.96, 5.94, 9.6, 8.96 for k = 14, 35, 41, 42, 52. Replacing the network's logits by the fitted formula gives loss 4.40×10⁻⁸ — lower than the network's own 2.41×10⁻⁷ (paper: 4.7 × 10⁻⁸ vs 2.4 × 10⁻⁷).
Worked example
What a fit of five numbers means. For the input (100, 50) the formula predicts that the logit of c is a sum of five cosines peaking at c = 37 (chapter 12's interference picture), and nothing else about the input matters — not whether (100, 50) was a training question, not any memorised detail. 96.1 % of the variance of all 113³ ≈ 1.4 million logits is captured that way.
The skip connection: why we could ignore it
The paper reasons with \Logits \approx W_L\,\mathrm{MLP}, ignoring the residual path that carries x^{(1)} (the attention output plus the “=” embedding) straight to the unembedding. Chapter 6 measured that zero-ablating that path raises the loss only from 2.41×10⁻⁷ to 8.88×10⁻⁷ (the paper reports accuracy unchanged), so the approximation is justified for the trained network.