This chapter is a guided walk through training. Use the bar below to pick a run and move through its checkpoints (or press ▶); every figure on the page follows. The dashboard shows all of these side by side.
Loading…
The spectrum sharpens
First, the embedding's Fourier spectrum over the whole run at once — each column is one checkpoint, each row one frequency:
The embedding's spectrum over the whole of training
Rows: frequency k (1 bottom, 56 top). Columns: checkpoints. Colour: √(cos² + sin²) norm in W_E.
Loading…
How to read it
<div class="fig-read"><p>This compresses the animation of the bar chart above into one picture; click a cell to jump there.</p>
What to expect: Early on all frequencies carry some weight; during circuit formation and cleanup everything except a few rows fades out.
Learn step by step →
· Source: [3] section 3, “Embedding in Fourier basis”
· Computed by checkpoint_metrics.py (code)
Early on every frequency has some weight. During memorisation the picture stays noisy; during circuit formation a few rows brighten while the rest fade; during cleanup everything except the key frequencies disappears. ARENA observes the same and notes that the singular values of W_E go from a random-matrix profile to a few dominant ones [3]. The neuron–logit map follows the same pattern, slightly earlier:
The neuron–logit map's spectrum over training
Same as above for W_L.
Loading…
How to read it
<p class="fig-meta">
<a href="../chapters/18-watch-it-form.html">Learn step by step →</a>
· Computed by <code>checkpoint_metrics.py</code> (<a href="https://github.com/BurnyCoder/visualizing-grokking-modular-addition">code</a>)
</p>
Circles emerge from noise
Project each number's embedding onto the plane that carries cos(w_k a) and sin(w_k a), for each key frequency. Move the bar from step ~1,000 to ~15,000:
Numbers on circles: the embedding in each key plane
Every number's embedding projected onto the 2D plane that carries cos(w_k a) and sin(w_k a), one panel per key frequency.
Loading…
How to read it
<div class="fig-read"><p>Each number n is coloured by the clock position the algorithm predicts for it at frequency k, namely k·n mod 113. A ring whose colours run smoothly around it means the embedding places n at angle w_k·n — the clock picture. Hands point at a, b and a + b; hovering a number enlarges it on every clock. While playing, the dots move between stored frames along the circle (polar blending); paused, the picture is a stored frame.</p>
What to expect: A cloud early in training that becomes a clean circle during circuit formation.
At initialisation the 113 points are a random cloud. As the circuit forms they spread into a ring whose colours (each number coloured by its predicted clock position k·a mod 113) run smoothly around it — the number a has become the angle w_k·a, the clock of chapter 9. The same happens for the answer: the MLP's average activation for each answer c = (a + b) mod 113, read along W_L's cos/sin directions, also traces circles:
The answer on a circle: MLP activations by answer
Average neuron activation for each answer c = a+b mod p, read along W_L's cos/sin(w_k c) directions.
Loading…
How to read it
<p class="fig-expect"><strong>What to expect:</strong> Circles: the MLP represents (a+b) mod p as an angle at each key frequency.</p>
<p class="fig-meta">
<a href="../chapters/18-watch-it-form.html">Learn step by step →</a>
· Computed by <code>geometry.frame</code> (<a href="https://github.com/BurnyCoder/visualizing-grokking-modular-addition">code</a>)
</p>
Pairs of numbers live on a torus: put a at angle w_k·a around the ring and b at angle w_k·b around the tube. Coloured by the probability of the right answer, the whole surface lights up as the network groks; coloured by the answer, the inputs that share an answer lie on one helix.
Every input on a torus
All 12,769 inputs (a, b) on a torus, at angles w_k·a around the ring and w_k·b around the tube; coloured by the probability of the right answer, by training/test pair, or by the answer a + b.
Loading…
How to read it
<div class="fig-read"><p>A pair of angles is a point on a torus, the way one angle is a point on a circle. The network's answer depends on the sum of the two angles, w_k(a + b), so the inputs that share an answer lie on one helix winding around the torus; colour by the answer to see the helices. Played through training, the colour by probability shows memorisation (a speckle of training pairs) give way to the whole surface lighting up. While playing, probabilities are blended between stored frames (every 200 steps for the paper's run). Drag to turn the torus.</p>
What to expect: Early in training only the training pairs are bright; after grokking every point is.
The logits learn the trig terms first, the neurons later
How much of the variance of the neurons and of the logits lies in the cos/sin(w_k(a+b)) directions? ARENA finds that the logits acquire this structure smoothly and early, the neurons later and more abruptly [3]:
Variance explained by cos/sin(w_k(a+b))
Fraction of the (centred) variance of neuron pre-activations, neuron activations and logits lying in the key cos/sin(w_k(a+b)) directions.
Loading…
How to read it
<p class="fig-expect"><strong>What to expect:</strong> Logits pick up the trig structure early and smoothly; neurons later and more abruptly.</p>
<p class="fig-meta">
<a href="../chapters/18-watch-it-form.html">Learn step by step →</a>
· Source: <a class="cite" href="../references.html#arena152" title="ARENA 1.5.2">[3]</a> section 3, “Development of computing trig components”
· Computed by <code>metrics.trig_fraction</code> (<a href="https://github.com/BurnyCoder/visualizing-grokking-modular-addition">code</a>)
</p>
Strength of each key frequency in the logits
Coefficient of the unit-normalised cube cos(w_k(a+b−c)) in the logits for each key frequency, over training.
Loading…
How to read it
<p class="fig-expect"><strong>What to expect:</strong> Each coefficient grows smoothly from zero long before the test loss drops.</p>
<p class="fig-meta">
<a href="../chapters/18-watch-it-form.html">Learn step by step →</a>
· Source: <a class="cite" href="../references.html#nanda2023" title="Nanda et al. 2023">[1]</a> Fig. 14
· Computed by <code>metrics.cos_coefficients</code> (<a href="https://github.com/BurnyCoder/visualizing-grokking-modular-addition">code</a>)
</p>
How many neurons belong to each frequency, over training
Number of neurons with more than 85 % of their variance at each key frequency, stacked.
Loading…
How to read it
<p class="fig-meta">
<a href="../chapters/18-watch-it-form.html">Learn step by step →</a>
· Computed by <code>checkpoint_metrics.py</code> (<a href="https://github.com/BurnyCoder/visualizing-grokking-modular-addition">code</a>)
</p>
How well the Fourier story explains the network, over training
Fraction of logit variance explained by Σα_k cos(w_k(a+b−c)), fraction of neurons dominated by one frequency, test accuracy of the logits restricted to the key frequencies, and the share of W_L outside the key frequencies.
Loading…
How to read it
<p class="fig-expect"><strong>What to expect:</strong> Restricted test accuracy reaches 100 % before grokking; the share of W_L outside the key frequencies falls from <span class="num" data-num="paper.wl_outside_init">91.3 %</span> at initialisation to <span class="num" data-num="paper.wl_outside_t3">0.3 %</span> at the end of cleanup.</p>
<p class="fig-meta">
<a href="../chapters/18-watch-it-form.html">Learn step by step →</a>
· Source: <a class="cite" href="../references.html#nanda2023" title="Nanda et al. 2023">[1]</a> Sec. 4.2–4.3
· Computed by <code>metrics.logit_cos_fit / neuron_frequencies</code> (<a href="https://github.com/BurnyCoder/visualizing-grokking-modular-addition">code</a>)
</p>
Confidence on every input
Each cell below is one question (a, b); its brightness is the probability the network gives the correct answer. During memorisation the training questions are bright and the test questions dark — the speckle is the random train/test split of chapter 2. During cleanup the test questions light up, all together:
Confidence on every input
Probability the model assigns to the correct answer, for all (a, b) at the selected step.
Loading…
How to read it
<p class="fig-expect"><strong>What to expect:</strong> During memorisation the training pairs are bright and the test pairs dark (a speckle pattern = the random split); after grokking everything is bright.</p>
<p class="fig-meta">
<a href="../chapters/18-watch-it-form.html">Learn step by step →</a>
· Computed by <code>geometry.frame</code> (<a href="https://github.com/BurnyCoder/visualizing-grokking-modular-addition">code</a>)
</p>
Uncertainty on every input
Normalised entropy of the prediction (0 = certain, 1 = uniform guess) for all (a, b).
Loading…
How to read it
<p class="fig-meta">
<a href="../chapters/18-watch-it-form.html">Learn step by step →</a>
· Computed by <code>geometry.frame</code> (<a href="https://github.com/BurnyCoder/visualizing-grokking-modular-addition">code</a>)
</p>
The geometry of the output
Averaging the network's predicted distribution over all inputs with the same correct answer gives 113 points; PCA draws them in two dimensions. As the network becomes certain, each average distribution approaches a one-hot vector, and the points move to the corners of the probability simplex:
Output geometry: predicted distributions
Mean softmax output of all inputs sharing the same answer, in two PCA directions.
Loading…
How to read it
<p class="fig-expect"><strong>What to expect:</strong> Collapses to the corners of a simplex as the model becomes certain; during memorisation the test inputs are far from their corners.</p>
<p class="fig-meta">
<a href="../chapters/18-watch-it-form.html">Learn step by step →</a>
· Computed by <code>geometry.frame</code> (<a href="https://github.com/BurnyCoder/visualizing-grokking-modular-addition">code</a>)
</p>
Output geometry: logits
Mean logit vector of all inputs sharing the same answer, in two PCA directions.
Loading…
How to read it
<p class="fig-meta">
<a href="../chapters/18-watch-it-form.html">Learn step by step →</a>
· Computed by <code>geometry.frame</code> (<a href="https://github.com/BurnyCoder/visualizing-grokking-modular-addition">code</a>)
</p>
The first three principal components in 3D (drag to turn):
The representations in 3D
The first three principal components of the number embeddings, of the mean logits per answer, or of the mean predicted distribution per answer, at the selected step.
Loading…
How to read it
<div class="fig-read"><p>Principal components are the directions of largest variance. A principal axis has no inherent sign, so while two stored frames are blended the second is flipped to agree with the first; paused, the stored frame is shown as computed.</p>
What to expect: Rings and loops appear as the circuit forms: the numbers are arranged by the clock.
Training does not end with cleanup at step 14,001. From around step 43,000, in the paper's run, two attention heads gradually start attending to “=” again (on average 11.3 %, 0.1 %, 9.1 %, 0.1 % for heads 0–3 at step 45,001, up from 0.44 %, 0.11 %, 0.44 %, 0.12 % at 40,001), and the sum of squared weights shows a small kink (955 → 942 in ARENA's description) [3]. Neither the paper nor ARENA explains it; it is a reminder that even this “fully understood” network has unexplained behaviour.
Average attention per head
Mean over all inputs of how much '=' attends to a, b and itself.
Loading…
How to read it
<p class="fig-expect"><strong>What to expect:</strong> About half to a and half to b; almost nothing to '='.</p>
<p class="fig-meta">
<a href="../chapters/06-simplifying-the-model.html">Learn step by step →</a>
· Source: <a class="cite" href="../references.html#nanda2023" title="Nanda et al. 2023">[1]</a> App. A.1; <a class="cite" href="../references.html#arena152" title="ARENA 1.5.2">[3]</a> section 3, “Development of commutativity”
· Computed by <code>checkpoint_metrics.py</code> (<a href="https://github.com/BurnyCoder/visualizing-grokking-modular-addition">code</a>)
</p>
The same averages over the whole run, one chart per head: the return of attention to “=” late in training stands out.
Where each head looks, over training
Mean attention from '=' to a, b and '=' for each head at every checkpoint.
Loading…
How to read it
<p class="fig-expect"><strong>What to expect:</strong> Attention to '=' drops from about a third (near-uniform attention over the three tokens at initialisation) to below 1 % by step <span class="num" data-num="paper.attn_eq_drop_step">101</span>; in the paper's run two heads attend to '=' again late in training.</p>
<p class="fig-meta">
<a href="../chapters/18-watch-it-form.html">Learn step by step →</a>
· Source: <a class="cite" href="../references.html#nanda2023" title="Nanda et al. 2023">[1]</a> App. C.1.3
· Computed by <code>checkpoint_metrics.py</code> (<a href="https://github.com/BurnyCoder/visualizing-grokking-modular-addition">code</a>)
</p>