Grokking, From the Characters
Up
A self-contained course for the Map of Math project: why small
neural networks trained on finite algebra rediscover representation
theory, math first. Prerequisites: group basics (§3 of
elliptic-curves.md); the Legendre symbol is met again from
quadratic-reciprocity.md. Canonical connections graph: the
map’s lens files.
How this doc fits the map
This is the study side of the training-dynamics track
(R26_training_dynamics/). Every chapter ends in a lab, and
every lab is a real record in the repo. A lab is either a scored run,
which you read against the chapter, or an open probe, which the chapter
prepares you to seal. The course does not score anything. Where a
chapter and a record disagree, the record wins and the chapter is
fixed.
| ch |
subject |
lab |
record state (2026-09-28) |
| 1 |
Characters: the Fourier basis of a finite group |
run 037, multiplicative characters mod 113 |
scored, confirmed |
| 2 |
Networks from scratch: perceptron, MLP, backpropagation by hand |
a tiny network on Z/7 addition — exercise, no record yet |
no record yet |
| 3 |
Embeddings and attention: the one-layer transformer, piece by
piece |
037’s control arm (a + b mod 113) |
scored with 037 |
| 4 |
Training dynamics: gradient descent, weight decay, plateaus as
singular values, the edge of stability |
P-027 (staircase as spectral readout), P-028 (edge of stability);
animations/learning-curves.html |
open |
| 5 |
Grokking: memorization against circuit; watching the circuit form
before the jump |
run 037’s loss curves and grokking steps |
scored |
| 6 |
Product groups: the Chinese Remainder Theorem inside an
embedding |
P-031 (units of Z/561); P-030 (pseudoprimes as error fingerprint;
561 is the first Carmichael number) |
open |
| 7 |
Non-abelian groups: irreducible representations and S5 |
grok.table_task takes S5’s 120 × 120 table; no probe
yet |
none |
| 8 |
Curves as spectra: double descent, scaling laws, which learning
curves read out an eigenvalue |
P-029 (double descent), P-035 (curve shapes);
animations/learning-curves.html |
open |
How to use a chapter: read one section, then do its
exercises by hand before moving on. The calculations are kept small on
purpose, because the pattern is the point.
Chapter 1
— Characters: How a Finite Group Hears Itself
1.1 The clock and its waves
A character of
is a way of placing the clock onto the unit circle that respects
addition. There are exactly
of them, and that one fact is the whole engine of this course.
Start with the clock
,
added mod
.
Write
for the first
-th
root of unity.
Definition. A character of
is a map
from
to the nonzero complex numbers with
Theorem 1.1. The characters of
are exactly the
maps
Proof. A character is fixed by its value on 1, since
.
Also
,
so
is an
-th
root of unity, one of
.
Each choice gives a character (Exercise 4). ∎
Picture
as a hand that moves
notches around the circle for every tick of the clock.
never moves.
goes round once as
runs from
to
,
and
goes round
times. The number
is the frequency, the same word as in music.
The real form. A network stores real numbers, not
complex ones. Splitting
into its real and imaginary parts gives the pair
So “the network learned frequency
”
means that two coordinates somewhere in its embedding of the token
hold this pair. Section 1.6 puts that sentence to work.
Distinct characters are orthogonal, so any function on the clock
splits into frequencies in exactly one way. That split is the discrete
Fourier transform.
Theorem 1.2 (orthogonality).
Proof. The sum is
with
.
If
then
and the sum is
.
Otherwise
but
,
and the geometric series gives
.
∎
So the
characters form an orthonormal basis for functions on
.
The coordinates of
in that basis are its Fourier coefficients:
The energy is shared out exactly (Parseval):
Worked by hand: n = 5
Take
,
a bump on the first two hours. Then
,
and since
,
|
|
|
squared |
| 0 |
1 |
0.4000 |
0.1600 |
| 1 |
0.8090 |
0.3236 |
0.1047 |
| 2 |
0.3090 |
0.1236 |
0.0153 |
| 3 |
−0.3090 |
0.1236 |
0.0153 |
| 4 |
−0.8090 |
0.3236 |
0.1047 |
The squares sum to 0.4000, and the right side of Parseval is
.
It balances.
Notice two things: - The mass sits at low frequency, because a bump
two hours wide is smooth on a five-hour clock. -
and
have the same size, because
is real. Frequency
and frequency
are the cos/sin pair of Section 1.1 seen twice.
Every real embedding in this course has that symmetry. That is why
run 037 folds frequencies into conjugate-pair bins,
,
and counts 56 of them.
1.3 Why characters
diagonalize addition
In the character basis, adding two clock readings becomes multiplying
two numbers on the circle. A network that learns characters has turned a
lookup table into one multiplication and one comparison.
Convolution. For functions on
,
define
.
With the normalization of Section 1.2,
Proof. Expand the sum, substitute
,
and use
to split the double sum into a product. ∎
Here is why this matters. The answer to
“”
is the point mass at
,
which is the convolution of the point masses at
and at
.
So the Fourier transform turns the addition table into pointwise
multiplication, one frequency at a time. That is what “diagonalize”
means.
The angle-addition trick. Suppose a network holds
the cos and sin of
for its input
,
and the same for
.
It can then form
and likewise the sine. This is
written in real numbers. It needs products of inputs, and those come
from the nonlinear layer; Chapter 2 shows how.
The readout: interference. To pick the answer
,
score every candidate by how well it lines up with
,
summed over a few frequencies
:
At
every term is 1, so the score is
.
Anywhere else the cosines point in different directions and partly
cancel. One frequency gives a correct but blunt peak, with near
neighbours scoring almost as high. Several frequencies interfere until
one spike stands alone. Exercise 5 finds a case where the cancellation
is perfect. The explorer in animations/learning-curves.html
draws exactly this sum.
1.4 The other clock:
multiplication
The nonzero residues mod a prime
form a second clock, with
hours, and multiplication runs it. Its characters are the Dirichlet
characters. They are periodic in a different coordinate, the discrete
log.
Theorem 1.3. For
prime, the group
of nonzero residues under multiplication is cyclic of order
.
Proof sketch. A polynomial of degree
has at most
roots mod
,
so for every
dividing
,
at most
elements satisfy
.
A counting argument over element orders then forces an element of order
exactly
.
∎
An element
of order
is a primitive root. Every nonzero
is a power of
,
.
The exponent
is the discrete log of
,
defined mod
.
The key identity is
The discrete log turns the multiplicative clock into the additive
clock
,
and everything in Sections 1.1 to 1.3 then carries over.
Worked by hand: p = 7, g = 3
The powers of 3 mod 7 are
,
,
,
,
,
,
and
closes the loop.
|
1 |
2 |
3 |
4 |
5 |
6 |
|
0 |
2 |
1 |
4 |
5 |
3 |
Check:
mod 7, and
mod 6, which is
.
It works.
Read the table in ordinary order and the logs jump around: 0, 2, 1,
4, 5, 3. That scrambling is the whole difficulty. A pattern that is
periodic in the log looks like noise when you list
Dirichlet characters. The characters of
are
The order of
is
.
Changing the primitive root multiplies every frequency
by a unit mod
.
That relabels the characters but leaves each order unchanged.
The Legendre symbol lives here. Take
.
Then
,
which is
exactly when
is a square mod
.
For
the squares are 1, 2 and 4, and their logs 0, 2 and 4 are the even ones.
This is the only character of order 2, and it does not depend on the
choice of
.
Quadratic reciprocity is a theorem about this one character (see
quadratic-reciprocity.md).
1.5 When the clocks meet: Gauss
sums
Look at a nontrivial multiplicative character through additive eyes,
on
,
and it is spread perfectly evenly over all nonzero frequencies. The two
clocks are as incompatible as two bases can be, and one short proof
shows it.
The residues mod
carry both clocks at once: addition on all of
,
and multiplication on the nonzero part. The Gauss sum
of a multiplicative character
pairs it against the first additive character:
Theorem 1.4. If
is not the trivial character, then
.
Proof. Expand
as a double sum over
and
,
and substitute
(as
runs over the nonzero residues, so does
):
The inner sum is
when
and
otherwise, because the full sum over all
,
including
,
vanishes. So
.
The values of a nontrivial character sum to zero over the group (Theorem
1.2 with
),
so
.
That leaves
.
∎
What it means. Replace
by
for any
and you get
,
which has the same size
.
So extend
to all of
by setting
.
Its Fourier transform, taken at length
,
then has exactly equal magnitude at every nonzero frequency and zero at
frequency 0. Additive analysis sees none of the multiplicative
structure; all it sees is perfect flatness.
Worked by hand: p = 5,
the Legendre symbol
The squares mod 5 are 1 and 4, so
on
.
With
,
Now
and
,
which are
and
,
where
is the golden ratio. The sum is
.
The theorem and the pentagon’s golden ratio appear in a single line.
The fine print, which matters in the lab. Theorem
1.4 describes the
length-
transform over
,
with the zero slot included. A network trained on
has tokens only for
.
Run 037 therefore read those tokens with a
length-
transform in ordinary order. That transform is not the character
transform of any group law on the tokens, so the theorem says nothing
exact about it. Section 1.6 shows what the difference costs.
1.6 Lab: what a network found
The record is
lab/runs/037-grokking-multiplicative-characters.md, scored
confirmed on 2026-09-28. Run 037 trained a one-layer
transformer on
mod 113 and found Dirichlet characters in the discrete-log coordinate of
Section 1.4.
The setup: - Two tasks, three seeds
each. The control is
mod 113, which should reproduce the published additive result (Nanda et
al. 2023, as cited in the record). The main task is
on the 112 nonzero residues. - The network and
training. One layer, 4 heads,
.
It trained on 30% of the table, with weight decay 1.0, AdamW and full
batch. - Every seed grokked. On the multiplicative
task, test accuracy reached 0.99 at steps 10,900, 12,300 and 23,200,
long after training accuracy did (step 200). - The
measurement. The analysis took each learned embedding and
measured the share of its non-DC power in the top 10 of 56
conjugate-pair bins, in two orderings: ordinary
()
and discrete log (primitive root
).
| Task and ordering |
Transform |
Top-10 share |
Section that explains it |
| Addition, ordinary order |
length 113 |
0.956 to 0.966 (3 seeds) |
1.3: characters of
|
| Multiplication, discrete-log order |
length 112 |
0.949 to 0.970 (3 seeds) |
1.4: Dirichlet characters |
| Multiplication, ordinary order |
length 112 |
0.231 to 0.270 (3 seeds) |
1.4: the scrambling (see below) |
| Gaussian noise (floor) |
either |
0.202 |
none |
| Reference: one exact Dirichlet character, ordinary
order |
length 112 |
0.344 (mean over
) |
1.4 |
| Reference: span of three characters, ordinary order |
length 112 |
0.275 (range 0.239–0.301) |
1.4 |
| Reference: one exact character, zero slot added |
length 113 |
0.179 exactly
|
1.5: Gauss sums |
Read the table with the chapter in hand: - Addition.
The addition network is the angle-addition circuit of Section 1.3: a
handful of frequencies carry nearly all of its structure. -
Multiplication in discrete-log order. The
multiplication network is the same circuit after the change of
coordinates in Section 1.4. The tokens arrive in ordinary order, so it
could not have seen the discrete log directly; it built that coordinate
itself. - Multiplication in ordinary order. It reads
almost like noise, but not quite (1.14–1.34× the floor). The reference
rows say why. The record’s post-hoc reference, which is not part of the
sealed test, reads a perfect multiplicative embedding under the same
length-112 transform. One character reads 0.344 and a three-character
span 0.275, both above the noise floor. A character indexed by value and
transformed additively is a mismatched-coordinate object, not a
structureless one. The trained seeds sit inside or just below the
three-character range. - Where Gauss sums apply. Only
the last row is Section 1.5’s theorem. Add the zero slot and transform
at length 113, and a single character comes out exactly flat, at
.
That is lower than noise, because a flat spectrum has none of
the order-statistic surplus that noise has.
Loose threads, recorded in 037 as unsealed observations and
not part of the verdict: - The Legendre symbol
(frequency 56, order 2) carried 11% of the power in one seed and almost
none in the other two. The seeds show no stable preference. -
Character order. There was no preference for low-order
characters: full-order characters took 53 to 66% of the power in every
seed. - Odd characters. All three seeds leaned toward
odd characters, those with
,
by growing margins (odd share 0.536, 0.661, 0.763; noise gives 0.505).
Since
,
a character is odd exactly when its frequency
is odd. Nobody has explained this yet.
That last thread is an open question with your name on it (Exercise
11).
Exercises
One star is routine, two stars need an idea, and three stars are
proofs to rebuild without looking. Do the hand exercises with a pencil
and a table of cosines; that is where the patterns stick.
- ★ Write the 4 × 4 character table of
(its entries are
).
Check by hand that rows 1 and 3 are orthogonal.
- ★ Redo the
computation of Section 1.2 for
.
Where does the mass sit now, and why? Check Parseval.
- ★ Find the smallest primitive root mod 11 and build its discrete-log
table. Circle the squares mod 11 and confirm that their logs are the
even ones.
- ★★ Fill in the “each choice works” step of Theorem 1.1: show that
is well defined on
when
.
- ★★ The perfect-interference exercise. On
,
compute
for
with the frequencies
,
then
,
then
.
Record the margin between
and the best wrong answer each time. The last case has a surprise;
explain it with Theorem 1.2.
- ★★ Compute the Gauss sum of the Legendre symbol mod 3 exactly, then
mod 7. In the second case, show
without using decimals.
- ★★ Prove that the order of a Dirichlet character does not depend on
the primitive root chosen, and that the Legendre symbol’s frequency
never moves.
- ★★★ Rebuild the proof of Theorem 1.4 from memory.
- ★★★ Using
,
show that the quadratic Gauss sum is real when
mod 4 and purely imaginary when
mod 4. (Which sign it has is a different, famous story: it took Gauss
years.)
- Computational. Plot the Legendre symbol mod 113 twice: in ordinary
order, then in discrete-log order with
.
Describe each picture in one sentence. Take the top-10 share of each
under a length-112 transform. Then take it again under a length-113
transform with the zero slot added, and compare all four numbers with
the table in Section 1.6. (
grok.analyse in
R26_training_dynamics/ does the reading.)
- Computational, open. Run 037’s seeds leaned toward odd characters.
Propose one reason why a network trained on
mod 113 might prefer
,
and one experiment that would test it. If the experiment is worth
running, it becomes a probe in the training-dynamics block, with a
sealed prediction written before any compute.
Hints. - (3) Try 2, and check that its order is 10,
not 2 or 5. - (5) With the frequencies
,
add the conjugates 4, 5 and 6 and use orthogonality. - (6) Mod 3 the
answer is
.
- (9) Conjugate
and substitute
.
- (10) The length-113 reading of the Legendre symbol should come out at
exactly
.
Why?
Summary and bridge to Chapter
2
A finite cyclic group has exactly as many characters as elements.
They are orthogonal, and in their basis the group operation becomes
multiplication on the circle.
-
has
characters,
;
a network stores each one as a cos/sin pair.
- Orthogonality makes the Fourier transform exact and invertible, with
energy conserved.
- Addition becomes multiplication frequency by frequency, and angle
addition is that fact written in real numbers.
- Mod a prime, multiplication is addition in the discrete-log
coordinate, and its characters are the Dirichlet characters. The
Legendre symbol is the one of order 2.
- Gauss sums prove that, on
,
the two clocks are maximally incompatible:
.
A network’s instrument, without the zero slot, sees a blurred version of
that flatness.
Next. This chapter assumed a network can multiply
two cosines. Chapter 2 builds that ability from nothing: a perceptron,
then a two-layer network, then backpropagation worked by hand on a
3-neuron example. Its lab trains a tiny network on
addition, small enough to read its weights directly, and asks whether
the character table of
appears in them.
Chapter
2 — Networks From Scratch: How a Nonlinearity Multiplies
A linear unit can add its inputs but cannot multiply them, and
Chapter 1 showed that a group law, read in characters, is
multiplication. This chapter builds the missing piece: a hidden layer
with a nonlinearity, trained by backpropagation, which you will run once
by hand.
2.1 The perceptron
A perceptron takes inputs
,
weights
and a bias
,
and outputs
The set
is a hyperplane, and the perceptron answers one question: which side of
it is
on? A task is linearly separable when some hyperplane
puts all the 1s on one side and all the 0s on the other.
The learning rule. Show the perceptron a labelled
example
.
If it answers correctly, do nothing. If it is wrong, move the hyperplane
toward the right answer:
On linearly separable data this rule stops after finitely many
mistakes. (Exercise 1 runs it by hand on AND. The perceptron is due to
Rosenblatt, 1958; citation unverified.)
Theorem 2.1. No perceptron computes XOR on
.
Proof. Suppose one did. Then
gives
.
The inputs
and
give
and
;
adding them,
.
But
needs
.
Contradiction. ∎
2.2 XOR
is addition mod 2, and addition mod 2 is a product
XOR is not an arbitrary hard example. It is the group law of
:
.
Now read it in characters, as in Chapter 1.
has two characters, and the nontrivial one is
,
sending
and
.
Section 1.3 says the group law becomes multiplication:
So in the
encoding, XOR is the product
.
Theorem 2.1 says a linear unit cannot compute it, and here is a sharper
reason.
Theorem 2.2. On the four inputs
,
the functions
,
,
and
are orthogonal. So the best linear fit to XOR, in least squares, is the
zero function.
Proof. These four functions are the characters of
,
and characters of a finite abelian group are orthogonal (Theorem 1.2
applied coordinate by coordinate). Least squares projects
onto the span of
,
and that projection is zero. ∎
A linear unit is not just unable to get XOR right. It cannot see XOR
at all: every weight it could learn has zero correlation with the
target. Multiplication has to come from somewhere else.
2.3 Two layers
Add a hidden layer: several units, each computing
and passing it through a nonlinear activation
.
The output is a linear combination of the hidden values. This chapter
uses the ReLU,
XOR with two hidden units, on
:
|
|
|
|
| (0, 0) |
0 |
0 |
0 |
| (1, 0) |
1 |
0 |
1 |
| (0, 1) |
1 |
0 |
1 |
| (1, 1) |
2 |
1 |
0 |
The second unit switches on only when both inputs are on, and
subtracts exactly enough. A sum of ReLUs is a piecewise-linear function;
each unit contributes one bend. That is all a one-hidden-layer ReLU
network is (Exercise 8).
2.4 How a nonlinearity
multiplies
Products come from squares. The polarization
identity is
So two hidden units with activation
,
one fed
and one fed
,
compute the product exactly. A ReLU is not a square, but it can make one
well enough where it is needed.
In the
encoding it is exact. Since
,
,
and
is 2 or 0. So
Two ReLU units compute the character product of Section 2.2 exactly.
Check all four inputs by hand; it takes a minute.
Back to Chapter 1. For
,
the angle-addition trick of Section 1.3 needs products like
.
The product-to-sum identity
and polarization say the same thing: a layer that can square can
multiply. That is the job the nonlinear layer does in a network that
adds on a clock. Nanda et al. (arXiv:2301.05217) reverse-engineered a
trained transformer and found it computing modular addition with
discrete Fourier transforms and trigonometric identities. Chapter 3
finds where in the network each step happens.
2.5 Loss and gradient descent
A network with parameters
produces an output
.
A loss
measures how wrong that output is. This chapter uses the squared loss
for one example with target
:
(Classification over
answers uses softmax and cross-entropy instead; Chapter 3 introduces
them.)
Gradient descent moves every parameter a small step
downhill:
where
is the learning rate. Near a point, the loss changes by about
per step. That estimate is only first-order. When the step is large, the
curvature of the loss takes over, and Chapter 4 is about what happens
then.
The only hard part is computing
for every weight at once. Backpropagation is the chain rule, organized
so that each quantity is computed once (Rumelhart, Hinton and Williams,
Nature 323, 533–536, 1986, DOI 10.1038/323533a0).
2.6 Backpropagation
by hand: a 3-neuron network
Two hidden ReLU units and one linear output: three neurons. Do this
section with a pencil. Every number is chosen to stay exact.
The network.
Starting weights:
,
,
;
,
,
;
,
,
.
The example: the XOR input
with target
.
Forward pass
| quantity |
computation |
value |
|
|
0.5 |
|
|
0.5 |
|
|
0.5 |
|
|
0.5 |
|
|
0 |
|
|
0.5 |
Backward pass
Write
.
Start at the loss and walk back, reusing each
once it is known.
| step |
rule |
value |
|
|
−1 |
|
|
−0.5, −0.5 |
|
|
−1 |
|
|
−1, +1 |
|
,
both
|
−1, +1 |
|
|
−1, 0, −1 |
|
|
+1, 0, +1 |
Notice that
:
an input that is off carries no blame.
One step, then check
Take
.
The new weights are
,
,
,
,
,
and
;
and
are unchanged. Run the forward pass again:
The loss falls from
to
.
Compare the first-order prediction. The squared gradient sums to
,
so the prediction is a drop of
.
That is more than the whole loss of 0.5. The actual drop is 0.399. A
step this large leaves the region where the linear estimate holds, a
first taste of Chapter 4.
The general rule
For a layer with pre-activations
,
weights
and activation
,
The table above is this rule with one hidden layer and every sum
written out.
2.7 Lab: a tiny network
on
addition
No record exists for this lab. It is an exercise
until someone seals it as a run, and this chapter makes no prediction
about its outcome.
The network. Encode
and
as one-hot vectors in
and concatenate them (14 inputs). Use one hidden ReLU layer and 7 output
logits, trained with cross-entropy (Chapter 3) and weight decay, on part
of the 49-entry table. Small is the point: the first-layer weights split
into two blocks of 7 columns, one per input token, and you can print
them on one page.
Reading the weights. Take the block for input
.
Each hidden unit has a row of 7 numbers, one per token. Take that row’s
length-7 DFT over the token index
.
Because the row is real, fold the frequencies into the DC term and three
conjugate-pair bins,
,
and
,
exactly as run 037 did at length 113.
Questions this lab asks, open: - Does the table get
learned at all from a fraction of 49 examples, and does test accuracy
jump late, as in run 037? - After training, how is each hidden unit’s
power spread over the three bins? - Do the hidden units for input
and for input
use the same bins? - How do the answers change with the hidden width and
the weight decay?
A hand warm-up comes first (Exercise 9): identify the frequency of a
given row of 7 numbers using nothing but three inner products.
Exercises
- ★ Run the perceptron rule by hand on AND, starting from
,
,
cycling through the four inputs in order. How many mistakes does it make
before it stops?
- ★ Redo the proof of Theorem 2.1 for the
encoding: no
give
on all four inputs.
- ★ Verify
on all four inputs. Draw it as a network with two hidden units.
- ★★ Using polarization, build a network with square activations that
computes
from inputs
and
.
How many hidden units does it need?
- ★★ Redo the backpropagation of Section 2.6 on the input
with target
,
from the same starting weights. Which weights get no gradient, and
why?
- ★★ The dead unit. Change
to
and redo the forward and backward passes for
.
Show that no gradient reaches
,
or
.
What does that mean for training?
- ★★ Derive the general rule of Section 2.6 from the chain rule, for
one hidden layer written in matrices.
- ★★★ Show that every continuous piecewise-linear function
with finitely many breakpoints is
.
Deduce that one hidden layer of ReLUs represents it exactly.
- ★★ Frequency by hand. A hidden unit’s row over the
tokens
is
.
Take its inner product with
for
,
and name the frequency. Why is no sine term needed?
- Computational. Check the gradients of Section 2.6 by finite
differences: nudge each weight by
and compare
with the table.
- Computational, open. Build the lab network of Section 2.7 and read
its weights as described. Record what you see; the questions there are
open.
Hints. - (1) Only the input
has target 1, so the bias must end up negative. - (5) With
both inputs are on, so no input carries zero blame this time. Look at
the signs. - (6)
for
.
- (9) Orthogonality (Theorem 1.2): only one inner product survives, and
it equals
.
Summary and bridge to Chapter
3
A linear unit adds; a group law in characters multiplies; a nonlinear
hidden layer supplies the products.
- XOR is addition on
,
and in the
encoding it is the character product
.
- The characters of
are orthogonal, so a linear unit has zero correlation with XOR.
- Two ReLU units compute
exactly; squares and polarization give products in general.
- Backpropagation is the chain rule, with each
computed once and reused.
- A large step outruns the first-order estimate of how much the loss
will fall.
Next. The lab encodes tokens as one-hot vectors and
lets the first layer learn one column per token. That column is an
embedding. Chapter 3 builds the one-layer transformer
of run 037 piece by piece: embeddings, attention, the MLP, and the
softmax readout. Its lab reads 037’s control arm,
,
through that anatomy.
References for Chapter 2
- Rumelhart, D. E., Hinton, G. E. and Williams, R. J. Learning
representations by back-propagating errors. Nature 323, 533–536
(1986). DOI 10.1038/323533a0.
- Nanda, N., Chan, L., Lieberum, T., Smith, J. and Steinhardt, J.
Progress measures for grokking via mechanistic interpretability.
arXiv:2301.05217.
- Rosenblatt, F. (1958), the perceptron. Unverified: no DOI
checked against the source.
A one-layer transformer that adds on a clock is Chapter 1’s circuit
built from learned parts. The embedding writes each token onto circles.
Attention and the MLP multiply those circles together, and the
unembedding sums the interference of Section 1.3. This chapter builds
each part and then says plainly which parts the repo has measured: so
far, only the embedding.
Sources for this chapter. Everything said here about
what the attention and MLP layers compute inside a trained network comes
from Nanda et al. (arXiv:2301.05217), about their own trained models.
Run 037 measured none of those layers. Section 3.6 lists what the repo
has and has not measured.
3.1 Tokens,
embeddings and the residual stream
The input is a sequence of three tokens,
,
and
,
each a one-hot vector. Position 3 is where the answer is read.
The embedding. A matrix
with one column per token value sends the one-hot vector of
to that column, a vector in
.
This is exactly Chapter 2’s first-layer block: a one-hot input picks out
one column of weights, and that column is the embedding. A learned
positional embedding
is added, so the network can tell the first slot from the second:
The residual stream. Each later layer reads the
current vector at a position and adds its output back. The
running sum is called the residual stream. Nothing is overwritten;
layers communicate by writing into shared coordinates.
What Chapter 1 predicts the embedding should hold.
If the network uses the angle-addition circuit, the embedding of
must contain the pairs
for a few frequencies
.
Nanda et al. report exactly this for their mainline model: a Fourier
transform along the token index of
shows significant norm at 6 frequencies, of which 5,
,
are used later in the network. They call these the key
frequencies.
3.2 Attention
Attention lets the vector at one position read from the vectors at
other positions. A head has three learned maps: a query
,
a key
and a value
,
plus an output map
.
The position reading (here,
)
forms a query; every position offers a key and a value.
The softmax (Section 3.5) turns scores into weights
that sum to 1. Several heads run side by side and their outputs are
added into the residual stream. The architecture is from Vaswani et
al. (arXiv:1706.03762).
Why attention can multiply. Look at what a head
computes: a weight
,
which depends on the tokens, times a value
,
which also depends on the tokens. That is a product of two
token-dependent quantities, the very thing Chapter 2 said a linear layer
cannot make. Attention is approximately bilinear, which is a second
route to products, alongside the MLP’s nonlinearity.
What Nanda et al. found in their heads (their
Appendices A.1 and C.1.2–C.1.3). Attention from
to itself was negligible, so each head’s softmax runs over just two
positions,
and
.
A softmax over two items is a sigmoid of their score difference,
.
The score
was close to a single wave,
,
and the sigmoid operated in a nearly linear range. So each head’s weight
is roughly
plus a single-frequency wave in
minus the same wave in
.
Multiplied by a value that is itself a wave, this gives degree-2
polynomials in sines and cosines of one frequency. In their model, two
heads did this; the other two mainly amplified the key frequencies
already in the embedding.
3.3 The MLP
After attention, the vector at position 3 passes through an MLP: a
wide linear map, a ReLU, and a linear map back.
This is Chapter 2’s two-layer network, applied at one position. Its
job in the circuit is to finish the products that Section 1.3’s
angle-addition formula needs, so that
and
are available as linear directions.
What Nanda et al. found in their MLP (their Sections
4.2–4.3). Of 512 neurons, 433 (84.6%) had over 85% of their variance
explained by a degree-2 polynomial of sines and cosines of a
single key frequency. And projecting the MLP activations onto
the directions the output reads gave, for each key frequency, a
near-exact multiple of
or
.
The network had computed Chapter 1’s angle-addition formula in its
neurons.
3.4 The
unembedding: the interference sum, learned
A final matrix
maps the residual stream to one logit per candidate
answer
.
In their model the skip connection around the MLP barely mattered, so
the logits are approximately
with
.
Now the payoff. To score candidate
,
the network needs
This is linear in the two quantities the MLP produced, with
coefficients
and
that depend only on
.
So a fixed matrix can do it: the rows of
just need to be those cosines and sines. Nanda et al. found
was approximately rank 10, one cosine and one sine direction for each of
the 5 key frequencies. They also found the logits were well approximated
by
with this form explaining 95% of the logits’ variance.
This is the interference sum of Section 1.3, with
learned weights
in place of 1. Every term peaks at
and the terms disagree elsewhere. Their Appendix B shows why one
frequency is not enough:
,
so frequency 14 alone nearly confuses
with
.
But
,
so adding frequency 35 cancels that false peak (Exercise 4 checks both
by hand).
3.5 Softmax and cross-entropy
Chapter 2 deferred these. The network outputs logits
,
one per answer. Softmax turns them into
probabilities:
Adding the same constant to every logit changes nothing, so only
differences between logits matter. Cross-entropy for
the correct answer
is
Its gradient is simple:
.
It pushes every wrong logit down in proportion to its probability and
the right one up by what is missing.
Tie to the interference sum. Suppose the logits are
,
a scale
times the Section 1.3 sum, and let
be the margin between the correct answer and the best wrong one. Every
wrong logit sits at least
below the right one, so
Two ways to lower the loss follow. Grow the scale
,
or grow the margin
by adding frequencies that cancel false peaks. Weight decay charges for
scale, so under weight decay a wider margin looks like the cheaper
route. That is one reason, in this model, to expect several frequencies
rather than one. It is a heuristic, not a consequence of the
bound: adding a frequency also adds weight norm — a cosine and
a sine direction in
,
and the neurons to feed them — so weight decay charges for margin too,
and the inequality above says nothing about which cost wins. Chapter 5
returns to the competition this sets up.
The same softmax appears inside attention (Section 3.2), where it
weights positions instead of answers.
3.6 The whole
network, and what the repo has measured
Putting the pieces together, following Nanda et al.’s Appendix A:
| Part |
What it computes in the Fourier circuit |
Source of that claim |
Measured in the repo? |
| Embedding
|
cos/sin of
at a few frequencies |
Nanda et al. §4.1 |
Yes: run 037 (Section 3.7) |
| Positional embedding
|
tells slot
from slot
|
Nanda et al. App. A |
No |
| Attention heads |
degree-2 polynomials of one frequency; amplification |
Nanda et al. App. C.1.2–C.1.3 |
No |
| MLP
(,
) |
,
as linear directions |
Nanda et al. §4.2–4.3 |
No |
| Unembedding
,
|
the interference sum over key frequencies |
Nanda et al. §4.2 |
No |
| Loss and accuracy |
the training and test curves |
— |
Yes: run 037 logs them |
Read the last column carefully. Run 037 saved only
— R26_training_dynamics/emb/*.npy holds the final embedding
matrix and nothing else; there is no checkpoint, no
state_dict, no saved optimizer. The positional embedding,
the four heads, both MLP matrices and the unembedding were never written
to disk, so they are not merely unanalyzed, they are
gone: the trained models no longer exist. Every claim
in this chapter about attention, the MLP and the unembedding is Nanda et
al.’s finding about their models, not a measurement of 037.
3.7 Lab: 037’s
control arm through this anatomy
The record is
lab/runs/037-grokking-multiplicative-characters.md, scored
confirmed. Its control arm trained on
and exists to reproduce the additive result that Nanda et al. published.
This lab reads that arm with Chapter 3 in hand.
Setup, side by side.
|
Nanda et al. mainline (their §3) |
Run 037 control arm (the record) |
| Task |
|
on
(all 12,769 ordered pairs) |
| Layers, heads, head dimension |
1 layer, 4 heads of dimension 32 |
1 layer, 4 heads of dimension 32 (MultiheadAttention’s
embed_dim // num_heads = 128/4) |
| Attention implementation |
written out from
(their App. A) |
torch.nn.MultiheadAttention(128, 4, bias=False) —
a different parametrization, not a re-implementation of
theirs: one packed in-projection rather than four named
matrices, so a head’s
and
have to be unpacked from it before anything in Sections 3.2–3.4 can be
looked for |
| Positional embedding; causal mask |
learned; their App. A |
learned W_pos over the 3 slots; strictly
upper-triangular mask, so = may attend to
,
and itself — their finding that
attention is negligible (0.1–0.4%) is empirical, and nothing here
forbids it |
| Residual width
,
MLP width |
128, 512 |
128, 512 |
| LayerNorm; biases; tied embed/unembed |
none; untied |
none; no biases anywhere; untied
(
and
are separate parameters) |
| Training fraction |
30% |
30% — 3,831 of 12,769 pairs, 8,938 held out |
| Optimizer, learning rate, weight decay |
AdamW, 0.001, 1.0, full batch |
AdamW
(),
0.001, 1.0, full batch |
| Seeds reported |
5 (mainline plus 4) |
3 (seeds 0, 1, 2); all three grokked |
What 037 measured on this arm. For each seed, the
final
was Fourier-transformed along the token index with a length-113
DFT over
.
Non-DC power was folded into 56 conjugate-pair bins
,
and the share in the top 10 bins was recorded.
| Control seed |
Top-10 share (length-113 DFT, 56 bins) |
Top bins |
Test accuracy 0.99 at step |
| 0 |
0.9656 |
44, 21, 34, 26, 52 (the top 5 of the 10) |
21,900 |
| 1 |
0.9558 |
26, 16, 47, 17, 34 (the top 5 of the 10) |
21,600 |
| 2 |
0.9636 |
10, 3, 25, 49, 50 (the top 5 of the 10) |
8,100 |
| Gaussian floor |
0.2022 (Gaussian, 113 × 128) |
none |
none |
Reading it. - The embedding row of Section
3.6’s table is confirmed for 037’s own seeds. A handful of bins
carry nearly all the non-DC power of
,
which is what the Fourier circuit requires of the embedding. It is the
only row 037 can confirm. - Which frequencies is a separate
question from how many. Compare the top bins across 037’s three
seeds and against Nanda et al.’s key frequencies
.
Their paper found different key frequencies in different seeds of the
same architecture: Appendix C.2.1 (Table 3) lists four other runs of it,
using
,
,
and
— no two alike, and none equal to the mainline’s five. So a mismatch
between 037 and the published set is not by itself a disagreement;
agreement would be the surprise. - Everything downstream of
is unmeasured here, and cannot be recovered from what was
saved. Only
was written out (R26_training_dynamics/emb/*.npy), so
whether 037’s heads, neurons and
match Sections 3.2–3.4 is not merely unknown from the record — the
weights that would answer it were never stored.
What an answer would require, stated as a cost
rather than as a plan, and with no prediction attached:
- A retrain that saves every weight, not just
— the model’s whole
state_dict at the end, and ideally at
intervals, since Chapter 5’s question is when the circuit forms.
- A check that the retrained
reproduces the saved one, seed for seed, before any
head or
is read. Without it a mismatch downstream cannot be told from a
different model. This check is not free, and 037’s own numbers say why:
the record’s only same-seed device comparison is the pair of 300-step
sizing runs for seed 0
(
metrics/add-transformer-s0-size100-{mps,cpu}.json), and
the two agree to
at step 0 and have already diverged to
in training loss by step 100. Identical seed, identical hyperparameters,
different device, different trajectory. So the reproduction check has to
be run on the same device, and even then it is a check,
not an assumption — and whether it is sufficient is open,
because same-device determinism on MPS has not been tested here. The
measurement, its limits and the cheap check that would close the first
gap are filed as N-031
(notes/mps-cpu-training-divergence.md).
- Unpacking
MultiheadAttention into
per-head
,
since 037’s attention is parametrized as one packed projection and
Sections 3.2–3.4 are written in the four-matrix form.
Open questions. None of these has a record, and this
chapter makes no prediction about any of them. - Do the frequencies in
037’s
match the frequencies in its
,
seed by seed? - In 037’s attention heads, is there the same split
between single-frequency heads and amplifying heads? - For the
multiplicative arm, what does each layer look like in discrete-log
coordinates?
Exercises
- ★ Write out the softmax probabilities for logits
and for
.
Why are they the same?
- ★ Show that a softmax over two positions with scores
equals
and
,
where
.
- ★ Derive
for cross-entropy after softmax.
- ★★ Checking Appendix B by hand. Compute
and
.
Hint: reduce
and
mod 113 first, and look at how close each is to
or to
.
- ★★ Prove the bound
of Section 3.5. When is it an equality?
- ★★ Expand
.
Which terms are degree 2? Which of them are pieces of
,
via the product-to-sum identity?
- ★★ Show that if the rows of a matrix
are
and
for
,
then
applied to
gives the vector
over
.
This is the unembedding of Section 3.4.
- ★★★ Using Theorem 1.2, show that the logits
,
summed over
together with the conjugates
,
have margin exactly
when
covers every nonzero frequency. What is that margin as a fraction of the
peak?
- Computational. For the frequency set
and
,
compute the interference sum of Section 3.4 with all
over
.
Report the margin, then drop one frequency at a time and report which
one costs the most margin.
- Computational, open. Load the saved final embeddings of 037’s
control arm and repeat its measurement yourself: a length-113 DFT along
the token index, conjugate-pair bins, top-10 share. Check your numbers
against the record. (They are not in the checkout —
.gitignore excludes
R26_training_dynamics/emb/, since the metrics JSON is the
record and the matrices are regenerable. Regenerate them with
python run037.py --task add --seeds 0 1 2 from
R26_training_dynamics/, using that directory’s own venv;
each seed writes emb/add-transformer-s<seed>.npy.
About six minutes a seed.)
Hints. - (1) Only differences between logits matter.
- (4)
,
and
,
which is close to
.
- (8) Chapter 1, Exercise 5 is the case
,
.
Summary and bridge to Chapter
4
A one-layer transformer can build Chapter 1’s circuit. Its embedding
writes circles, its attention and MLP multiply them, and its unembedding
sums the interference.
- The embedding is Chapter 2’s first-layer block: one column per
token.
- Attention multiplies a token-dependent weight by a token-dependent
value, so it is a second source of products.
- In Nanda et al.’s model, the MLP carried
and
as linear directions, and the unembedding read them off.
- Softmax plus cross-entropy rewards margin, and under weight decay
margin is cheaper than scale.
- The repo has measured only the embedding and the training curves;
the internal layers of 037 are unmeasured.
Next. So far every network has been examined after
training. Chapter 4 studies the training itself: gradient descent as a
dynamical system, what weight decay does to it, why loss curves have
plateaus, and what happens when the learning rate is too large. Its labs
are the open probes P-027 and P-028, and
animations/learning-curves.html draws both.
References for Chapter 3
- Nanda, N., Chan, L., Lieberum, T., Smith, J. and Steinhardt, J.
Progress measures for grokking via mechanistic interpretability.
arXiv:2301.05217. (Sections 3, 4.1–4.3; Appendices A, A.1, B,
C.1.2–C.1.3, C.2.1.) Checked against v3, the current
version (submitted 19 October 2023; v1 12 January 2023, v2 13 January
2023): every figure quoted in this chapter, and every section and
appendix title cited, was read from it on 2026-09-28. The PDF is in
reading/ as 2301.05217v3.pdf.
- Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L.,
Gomez, A. N., Kaiser, Ł. and Polosukhin, I. Attention is all you need.
arXiv:1706.03762.
Chapter
4 — Training Dynamics: Gradient Descent as a Dynamical System
Training is a map applied over and over,
,
and the shape of a loss curve is the behaviour of that map. This chapter
reads two shapes exactly. A staircase of plateaus is a
spectrum being learned one mode at a time. A period-doubling
cascade is what the map does when the step is too large for the
curvature. Both are exact only in small models, and that boundary is
where the chapter’s two labs, P-027 and P-028, begin.
4.1 Gradient descent is a map
Chapter 2 wrote one step of gradient descent. Iterate it:
This is a discrete dynamical system, and the vocabulary of dynamics
applies directly.
- A fixed point of
is a point with
:
a minimum, a maximum or a saddle of the loss.
- A fixed point is stable when nearby starts converge
to it. For a map, that is decided by the Jacobian
,
where
is the Hessian of
:
the fixed point is stable when every eigenvalue of
has absolute value below 1.
- A plateau in the loss curve is time spent near a
fixed point that is not a minimum, usually a saddle, before the iterate
finds the way out.
Small
makes
track the continuous gradient flow
,
with
steps of size
corresponding to time
.
Section 4.3 works in that limit. Section 4.4 is about what happens when
is not small.
4.2 The
quadratic: stable exactly when
Take a quadratic loss with its minimum at the origin:
Then
,
and gradient descent is linear:
.
Theorem 4.1. Gradient descent on this quadratic
converges to 0 from every starting point if and only if
.
Proof. Expand
in the eigenvectors
of
:
.
Each component evolves on its own,
,
so after
steps it is
.
This goes to 0 for every
exactly when
for every
,
that is
.
The binding case is the largest eigenvalue. ∎
Read the factor
component by component.
| Step size |
Factor
|
What that component does |
|
between 0 and 1 |
shrinks monotonically |
|
exactly 0 |
vanishes in one step |
|
between −1 and 0 |
shrinks while flipping sign every step |
|
below −1 |
grows while flipping sign |
The instability arrives through the sign-flipping row. At
the sharpest direction stops shrinking and starts bouncing between two
sides of the valley. That bounce is the seed of Section 4.4’s period
doubling. The quantity
is called the sharpness.
Two consequences: - The slowest direction sets the speed. With
just under
,
the flattest direction shrinks by
per step. The ratio
,
the condition number, is the price of a curved valley. - On a quadratic
the threshold is exact. On a real loss,
changes as
moves, so
is a moving target. Cohen, Kaur, Li, Kolter and Talwalkar
(arXiv:2103.00065) report that full-batch gradient descent on neural
networks typically runs with the sharpness hovering just above
,
while the loss still decreases over long timescales. They call this
regime the edge of stability. Theorem 4.1 says a
quadratic cannot do that. Why a network can is part of what Section 4.4
and P-028 are about.
4.3 Plateaus and the
singular-value staircase
The cleanest source of plateaus is a network with no nonlinearity at
all. Saxe, McClelland and Ganguli (arXiv:1312.6120) found exact
solutions for gradient descent in deep linear networks. This section
derives the heart of it for one mode, by hand.
The setup. A two-layer linear network maps
.
Train it with squared loss on inputs that are whitened
(their covariance is the identity). Write the input-output correlation
matrix of the data as its singular value decomposition,
,
with singular values
.
If the weights start aligned with the singular vectors, they stay
aligned, and each mode
is carried by one scalar strength per layer,
in
and
in
.
The network’s strength on mode
is the product
.
The modes then decouple, and the loss is, up to a constant,
One mode, by hand. Drop the index. Gradient flow on
gives
Note that
.
So
is conserved. With balanced initialisation,
at the start, it stays balanced, and
.
Then
This is the logistic equation. Separate variables with
and integrate from
:
The staircase. For small
,
the mode sits near
for a long time, then rises to
in a burst. It reaches half its final value at
Three features follow directly: - Big modes go
first.
falls as
grows, so modes are learned in order of singular value. - The
wait is logarithmic in the initialisation. Halving
adds
to the wait. The plateau is time spent escaping the saddle at
,
which is unstable but has a small gradient when
is small. - The burst width is about
.
The rise from 10% to 90% takes
.
The loss is a sum of one drop per mode, each of size
,
each at its own
.
With well-separated singular values that is a staircase: a plateau, a
drop, a plateau.
How exact this is. Every word in the setup matters.
The solution is exact for gradient flow on a two-layer linear
network, with whitened inputs, aligned and balanced initialisation, and
one mode per singular value. Nothing in the derivation used a
nonlinearity, and nothing survives automatically once one is added. In a
nonlinear network the modes need not decouple, the “spectrum” being
learned is not the spectrum of any fixed matrix, and plateaus can have
other causes.
Reading it backwards. The forward direction is a
theorem in this model: a spectrum produces a staircase. The reverse
question asks whether you can take a staircase from a network you did
not design, read off its plateau times, and recover a spectrum. That is
P-027’s open question, together with what happens when two singular
values nearly coincide and their drops merge.
The Staircase panel of animations/learning-curves.html
draws
from the formula above with three adjustable singular values.
4.4 Past the threshold:
period doubling
Section 4.2 ended with the sharpest direction bouncing at
.
On a quadratic the bounce grows without limit. On a real loss the
curvature changes as the iterate moves away, and the bounce can settle
into a new attractor. The cleanest place to watch this happen is one
dimension.
The model. Take the double well
It has minima at
,
with curvature
,
and a maximum at
.
Gradient descent is the map
Stage 1: the threshold is exact. The minimum
is a fixed point, with multiplier
.
It is stable when
,
that is
.
Since the curvature there is 2, this is Theorem 4.1’s
exactly:
.
At
the multiplier passes through
,
and a period-2 orbit is born around
:
the iterate now alternates between two points on either side of the
minimum.
Stage 2: each side is a unimodal map. The map
is an odd cubic with two critical points, at
On the positive side,
has a single maximum at
,
and it is quadratic there, since
.
An orbit that starts on the positive side stays on it as long as the
map’s highest value does not reach past the positive zero of
,
which is
.
Compute that highest value:
Setting
gives
,
so
Below
the two sides never mix. An orbit on the positive side sees one hump
with one quadratic maximum, which is the setting of the logistic map.
Above it, orbits cross zero, the two symmetric attractors merge into
one, and both critical points matter at once.
Stage 3: the cascade happens before the merge. As
rises past 1, the period-2 orbit itself loses stability and doubles to
period 4, then 8, 16, and so on, at parameters crowding toward a limit.
This is the period-doubling cascade, and Feigenbaum
(“Quantitative universality for a class of nonlinear transformations,”
J. Stat. Phys. 19, 25–52, 1978, DOI 10.1007/BF01020332) showed
that for one-dimensional maps with a single quadratic maximum, the gaps
between successive doublings shrink by a universal ratio,
the same number for every map in the class.
Here the whole cascade finishes inside the one-sided regime. The
table below locates the superstable parameters, where
the periodic orbit passes through the critical point
.
They sit between successive doublings and give the cleanest estimate of
.
The values were computed for this chapter by root-finding on
at 50-digit precision. They are not a lab record, and Exercise 10 asks
you to reproduce them.
|
period
|
superstable
|
ratio
|
| 0 |
1 |
0.500000 |
none |
| 1 |
2 |
1.121320 |
4.3857 |
| 2 |
4 |
1.262990 |
4.5917 |
| 3 |
8 |
1.293843 |
4.6527 |
| 4 |
16 |
1.300475 |
4.6656 |
| 5 |
32 |
1.301896 |
4.6684 |
| 6 |
64 |
1.302200 |
4.6690 |
| 7 |
128 |
1.302266 |
none |
The ratios climb toward
,
and every doubling lands near
,
well below
.
Past the accumulation point the orbit is chaotic, still confined to one
side, until the merge at
lets it wander over both wells.
What this settles, and what it does not. In this
one-dimensional map, the edge of stability, the period-2 bounce and the
whole Feigenbaum cascade are settled mathematics. They are the known
result that any instrument claiming to measure “period doubling in
gradient descent” must reproduce before it is trusted on anything else.
What survives in high-dimensional training is open. There, many
directions share the curvature budget, several can cross
at once, and the sharpness itself moves as training proceeds. Whether a
doubling cascade appears at all, whether its spacing obeys
,
and in what dimension any of it stops holding is P-028’s question.
The Edge-of-stability panel of
animations/learning-curves.html draws the late iterates of
against
.
4.5 Weight decay
Every run in this course trains with weight decay. Here is what it
does to the dynamics.
On plain gradient descent. Add the penalty
to the loss. The step becomes
Every step first shrinks all weights by the factor
,
then follows the data gradient. With no data gradient the norm decays
geometrically,
.
With one, a fixed point needs
:
the data must push outward exactly as hard as the decay pulls in. Any
direction the data does not need drains to zero.
On a quadratic the penalty adds
to every eigenvalue of
.
So Theorem 4.1’s threshold becomes
,
and the flattest directions are no longer flat.
On AdamW, which run 037 uses. For plain gradient
descent, an
penalty and “shrink the weights each step” are the same thing. For
adaptive optimizers such as Adam they are not, because Adam rescales
each gradient coordinate. Loshchilov and Hutter (arXiv:1711.05101) made
this point and proposed applying the decay directly to the weights,
separately from the adaptive step. That is AdamW. The norm-shrinking
picture above describes AdamW’s decay term. The curvature picture does
not transfer directly, because Adam’s rescaling changes the effective
step size direction by direction.
The question passed to Chapter 5. Section 3.5 found
two ways for the one-layer transformer to lower its cross-entropy: grow
the scale
of its logits, or grow the margin
by using more frequencies. Weight decay charges for both. Scale costs
norm directly, and each extra frequency costs a cosine and a sine
direction in the unembedding, plus the neurons that feed them. Which
cost wins, and whether that competition is what sets the timing of
grokking, is Chapter 5’s question. This chapter does not answer it.
4.6 Labs: P-027 and P-028
Both labs are open probes. Neither has a run, so there are no lab
numbers here, only exercises. Neither probe has a prediction or a
falsifier filed. Both are the owner’s to write, before any compute.
P-027, staircase as spectral readout. Section 4.3
proves the forward direction in a model: singular values produce plateau
times
.
The probe asks about the reverse direction, and about near-coincident
singular values. Before anything nonlinear is tried, an instrument for
this probe should reproduce Section 4.3 on a network where the answer is
known: a two-layer linear network trained on data with a spectrum you
choose. Exercises 5 and 11 build that instrument.
P-028, edge of stability. Section 4.4 gives the
one-dimensional result exactly: the threshold at
,
the cascade accumulating near
with ratios converging to
,
and the merge at
.
An instrument for this probe must reproduce the table in Section 4.4
before it measures anything in higher dimension. Exercises 7, 10 and 12
build that instrument.
Open questions, with no prediction attached. - If a
nonlinear network’s loss curve shows plateaus, what, if anything, do the
plateau times measure? - When two singular values are close, do their
drops merge into one, and can the pair still be resolved from the curve?
- In a two- or three-dimensional loss, does gradient descent past
period-double, and if it does, with what spacing? - As the dimension
grows, at what point, if any, does the one-dimensional picture stop
describing training?
Exercises
- ★ Run gradient descent on
from
with
and with
,
for four steps each. Which component flips sign? Which run
diverges?
- ★ For the quadratic in Exercise 1, find the
that makes the slower component shrink fastest while the faster one
stays stable. What is the per-step factor of the slow component?
- ★ Verify the logistic solution of Section 4.3 by differentiating it.
Then check
and
.
- ★★ Show that
is conserved by the one-mode flow, and solve the flow for unbalanced
starts with
.
What changes about the plateau?
- ★★ Take
and singular values
,
and
.
Compute each
.
What fraction of the total loss has been removed at each one?
- ★★ With weight decay
added to the one-mode loss, the flow becomes
and likewise for
.
Find the balanced fixed point. Show that a mode with
is never learned.
- ★★ For
,
find the period-2 orbit around
just past
.
(Hint: solve
and remove the fixed points.)
- ★★ Rederive
from
.
- ★★★ The first superstable parameter in Section 4.4’s table is
,
which agrees with
to every digit computed. Prove it, or find where the agreement
breaks.
- Computational. Reproduce Section 4.4’s table: find the superstable
for
by root-finding on
,
starting each search just above the previous root. Report the ratios.
How many digits of precision do you need by
,
and why?
- Computational. Train a two-layer linear network by plain gradient
descent (finite
,
not flow) on whitened data whose correlation matrix has singular values
,
and
,
starting balanced at
.
Plot the loss and mark each predicted
,
converting time to steps by
.
Where does finite
make the formula drift?
- Computational, open. Run gradient descent on a two-dimensional loss
that has the double well in one coordinate and a quadratic in the other,
.
Sweep
and
.
Record what you see; this chapter makes no prediction about it.
Hints. - (1) The eigenvalues are 1 and 10, so the
threshold is
.
- (2) Balance the two factors:
.
- (6) The fixed point is
when that is positive. - (8) Square both sides before solving. - (10)
Near the accumulation point, neighbouring superstable parameters differ
in the sixth decimal place, and the orbit is iterated 128 times.
Summary and bridge to Chapter
5
Gradient descent is a map, and a loss curve is its trajectory.
- On a quadratic, gradient descent is stable exactly when
,
and it fails by a sign-flipping bounce in the sharpest direction.
- In a two-layer linear network, each singular mode follows a logistic
curve that reaches half height at
.
The loss is a staircase, exact only for linear networks.
- For gradient descent on the double well, the threshold at
is exact. The period-doubling cascade that follows finishes before the
two sides merge at
,
and inside that window each side is a unimodal map governed by
Feigenbaum’s
.
- Weight decay shrinks every weight each step, drains directions the
data does not need, and charges for both scale and margin.
- Reading staircases backwards (P-027) and carrying the cascade into
high dimensions (P-028) are open.
Next. Chapter 5 puts the pieces together on the
curve this course is named for. Grokking has a plateau, as in Section
4.3, but the plateau is not a saddle in a linear model. It is a
memorizing solution that fits the training set perfectly while the
generalizing circuit of Chapter 3 forms underneath, under the
weight-decay pressure of Section 4.5. Its lab reads run 037’s training
curves and grokking steps.
References for Chapter 4
- Saxe, A. M., McClelland, J. L. and Ganguli, S. Exact solutions to
the nonlinear dynamics of learning in deep linear neural networks.
arXiv:1312.6120 (v3, 19 Feb 2014).
- Cohen, J. M., Kaur, S., Li, Y., Kolter, J. Z. and Talwalkar, A.
Gradient descent on neural networks typically occurs at the edge of
stability. arXiv:2103.00065 (v3, 23 Nov 2022).
- Feigenbaum, M. J. Quantitative universality for a class of nonlinear
transformations. J. Stat. Phys. 19, 25–52 (1978). DOI
10.1007/BF01020332.
- Loshchilov, I. and Hutter, F. Decoupled weight decay regularization.
arXiv:1711.05101 (v3, 4 Jan 2019).
Chapter 5 —
Grokking: Memorization Against Circuit
A network that groks holds two solutions to the same training set,
one after the other. The first is a lookup table: fast to build, perfect
on the training pairs, useless on anything else. The second is the
character circuit of Chapters 1 and 3: slow to build, and correct
everywhere. Weight decay is what makes the second one win. This chapter
reads that story off run 037’s curves, and it says plainly which parts
of the story the curves can and cannot show.
5.1 The shape
Power, Burda, Edwards, Babuschkin and Misra (arXiv:2201.02177) named
the effect: on small algorithmic datasets, a network can reach perfect
training accuracy, sit there for a long time with test accuracy near
chance, and then generalize. In their words, the improvement “can happen
well past the point of overfitting.”
Three numbers describe the shape of one run: -
:
the first logged step at which training accuracy reaches 0.99. -
:
the first logged step at which test accuracy reaches 0.99. - The
gap
.
A network that simply learns has a gap near 1. A grokking network has
a gap in the tens or hundreds. Run 037’s six seeds have gaps from 40.5
to 116 (Section 5.6).
5.2 Two solutions to one
training set
The training set is 30% of the table: 3,831 pairs on the additive
task, 3,763 on the multiplicative one. Many functions fit it exactly.
Two matter here.
The lookup table. Store each training pair’s answer
somewhere in the weights. This fits the training set perfectly and says
nothing about the other 70%. Each stored pair needs its own capacity, so
the cost grows with the number of pairs.
The circuit. Embed each token as cosines and sines
at a few frequencies (Section 1.1), multiply them (Sections 2.4 and
3.2–3.3), and sum the interference (Sections 1.3 and 3.4). This fits the
training set and the rest of the table. Its cost depends on how
many frequencies it uses, not on how many pairs it has seen.
Why the lookup comes first. Gradient descent follows
the steepest local improvement. Storing one more training pair improves
the loss on that pair immediately and on its own. A single frequency of
the circuit helps only once the embedding, the products and the readout
for that frequency all line up, because a partial circuit mostly
produces noise. So early training fills the table. This is a heuristic,
not a theorem, and 037’s curves are consistent with it: training
accuracy hits 0.99 by step 200 or 300 in every seed.
5.3 What memorization
costs on the test set
A memorizer is not ignorant on unseen pairs. It is confidently wrong,
and cross-entropy (Section 3.5) punishes that harder than ignorance.
Guessing uniformly over 113 answers gives every test pair the
loss
Now suppose the network puts nearly all its probability on one wrong
answer and leaves only
for the remaining 112, spread evenly. The correct answer then gets
,
and the loss is
A test loss of 20 nats, which is what 037’s seeds show at the moment
they finish memorizing (Section 5.6), means
,
so the network leaves roughly
of its probability for everything but its chosen answer. It is about
four times worse than guessing. The test loss is therefore a much more
sensitive instrument than the test accuracy. Accuracy sits at chance
either way; the loss shows how confident the wrong answers are.
5.4 Weight decay picks the
winner
Once the training set is fitted, the data gradient on it is small.
What is left pushing on the weights is weight decay (Section 4.5), and
the two solutions pay for it differently.
The clock that decay runs on. Run 037 trains with
AdamW at a constant learning rate
and weight decay
(the record’s hyperparameters). PyTorch’s AdamW applies the decay as
every step, separately from the adaptive step (Section 4.5). So a weight
that no gradient defends shrinks by a factor of
per step:
Undefended weight halves in about
steps. That sets the natural unit of time for anything decay does in
037: about a thousand steps. (This number is computed for the chapter
from the record’s hyperparameters, not measured; Exercise 4 checks
it.)
The heuristic, stated as one. The lookup table’s
cost grows with the number of stored pairs; the circuit’s cost grows
with the number of frequencies. If the circuit fits the same training
set at lower weight norm, decay slowly favors it, and once the circuit
carries the training loss, decay can dismantle the table, which no
longer earns its norm. Section 3.5’s scale-against-margin trade-off is
one instance: under decay, every extra frequency and every extra unit of
logit scale has to pay for itself. Whether the circuit really is the
lower-norm solution for a given architecture is not proved here. Nanda
et al. (arXiv:2301.05217, §5.3 and Appendix D.1) report the empirical
side: with less weight decay their networks took much longer to grok,
and without weight decay or another regularizer they did not grok on
this task.
5.5 Three phases,
and measures that see them early
Nanda et al. (arXiv:2301.05217, §5.1–5.2) used the circuit of Chapter
3 to define two progress measures, both built from the
Fourier analysis of Chapter 1:
- Restricted loss. Take the logits, do a 2D DFT over
the inputs
,
keep only the constant term and the components at the key frequencies
that the circuit uses, and measure the loss of what is left. It asks how
well the circuit alone would do.
- Excluded loss. Remove exactly those key-frequency
components instead, and measure the training loss of the rest. It asks
how much of the training fit is carried by something other than the
circuit, which is to say by the lookup table.
With these, they split their mainline run (the same task and
hyperparameters as 037’s control arm) into three phases, measured in
epochs of full-batch training:
| Phase |
Epochs |
What the measures show (their §5.2) |
| Memorization |
0 to 1.4k |
train and excluded loss fall; test and restricted loss stay
high |
| Circuit formation |
1.4k to 9.4k |
excluded loss rises, restricted loss starts to fall, the sum of
squared weights falls; train and test loss stay flat |
| Cleanup |
9.4k to 14k |
test loss drops suddenly, the sum of squared weights drops
sharply |
Their headline reading: the circuit forms well before test
accuracy moves, and the jump happens during cleanup, when decay strips
out the memorizing components that were drowning the circuit’s answer on
test pairs. On the test set the circuit was already partly right; the
lookup table’s confident wrong answers, from Section 5.3, were outvoting
it. Grokking, on this reading, is the moment the circuit stops being
outvoted.
What 037 can and cannot show here. Both progress
measures need the full logits, or all the weights, at many points during
training. Run 037 logged train and test loss and accuracy every 100
steps and saved only the final embedding
(Chapter 3, Section 3.6). So 037 cannot compute restricted or excluded
loss, and it did not log the weight norm. The phase boundaries in the
table above are Nanda et al.’s, about their model; none of them is
measured in the repo.
Watch. After the phases table: Welch Labs, “The most
complex model we actually understand” (YouTube), which
animates the grokked modular-addition network and its phases. It adds a
moving picture of what this section can only tabulate. Any claim this
chapter takes from it would be marked as the video’s; none is.
5.6 Lab: run 037’s curves
The record is
lab/runs/037-grokking-multiplicative-characters.md, scored
confirmed. Every number below comes from
R26_training_dynamics/metrics/{add,mul}-transformer-s{0,1,2}.json:
the grokking block for the step numbers and the
per-100-step log for the rest. The first-crossing steps are
resolved to the logging interval, 100 steps. No spectral numbers appear
in this lab.
| Arm, seed |
(0.99) |
(0.99) |
gap |
test acc 0.1 → 0.9, steps |
test loss at
|
peak test loss (step) |
final test loss (step 25,000) |
| add, 0 |
200 |
21,900 |
109.5 |
17,600 → 21,400 |
20.11 |
38.59 (1,500) |
9.4 × 10⁻⁷ |
| add, 1 |
300 |
21,600 |
72.0 |
19,100 → 21,300 |
22.21 |
36.92 (1,500) |
6.3 × 10⁻⁵ |
| add, 2 |
200 |
8,100 |
40.5 |
5,900 → 7,700 |
20.23 |
37.11 (1,400) |
4.3 × 10⁻⁵ |
| mul, 0 |
200 |
10,900 |
54.5 |
9,000 → 10,700 |
20.48 |
36.04 (1,400) |
1.8 × 10⁻⁶ |
| mul, 1 |
200 |
12,300 |
61.5 |
9,700 → 12,000 |
19.22 |
37.19 (1,400) |
8.0 × 10⁻⁵ |
| mul, 2 |
200 |
23,200 |
116.0 |
17,700 → 22,400 |
19.09 |
37.82 (1,400) |
8.7 × 10⁻³ |
All losses are in nats; chance is
on the additive arm and
on the multiplicative one.
Reading it with the chapter. - Memorization
is fast and uniform. Every seed on both arms fits the training
set by step 300, and at that moment every test accuracy is between 0.4%
and 2.1%, near chance. - The confidence of the wrong answers
keeps growing after the fit. Test loss is about 20 nats at
,
already four times the guessing loss (Section 5.3). It keeps climbing
for another thousand steps and peaks at 36–39 nats, at step 1,400 or
1,500 in all six seeds. - The jump is short and late.
The climb from 10% to 90% test accuracy takes 1,700 to 4,700 steps,
against waits of 5,900 to 19,100 steps before it starts. On the decay
clock of Section 5.4 the jump spans a few thousand-step units; the wait
before it can be close to twenty. - Seeds disagree about when,
not whether. With identical hyperparameters,
ranges from 8,100 to 23,200, a factor of 2.9. Three seeds per arm is too
few to say whether the arms differ. - The latest groker is still
cleaning up. Multiplicative seed 2 reached 0.99 last (step
23,200), and at step 25,000 its test loss is
nats, two to four orders of magnitude above the others’. It is
consistent with a run caught before cleanup finished. The record cannot
confirm that, because 037 has no progress measure.
One coincidence, marked as one. The test-loss peak
lands at step 1,400–1,500 in all six seeds, and Nanda et al.’s
memorization phase ends at 1.4k epochs under the same learning rate and
weight decay. It is tempting to call these the same event. Nothing in
037’s record can test that: the peak is visible in the test loss, while
the phase boundary was defined with the excluded loss, which 037 did not
log.
Open questions, with no prediction attached. - What
sets a seed’s grokking step, when everything but the seed is fixed? - Is
the test-loss peak the end of memorization, in the sense of Nanda et
al.’s excluded loss? - Does the multiplicative arm, which must build its
own discrete-log coordinate (Section 1.6), grok at a different time from
the additive arm, once there are enough seeds to tell?
Answering any of them needs a retrain that logs the weight norm and
saves checkpoints through training, so the progress measures can be
computed. As Section 3.7 notes, that retrain must run on the device of
the original and check that it reproduces the saved
before anything is read (note N-031).
Exercises
- ★ Compute
and
.
Explain why the two arms have different chance losses.
- ★ A network puts probability 0.9 on one wrong answer and spreads 0.1
evenly over the other 112. What is its loss on that pair? Compare with
guessing.
- ★ From the lab table, compute the ratio of wait to jump for each
seed: the step at which test accuracy reaches 0.1, divided by the number
of steps from 0.1 to 0.9.
- ★★ Show that
when
is small. With
,
how far off is the approximation at
?
At
?
- ★★ Invert Section 5.3: what value of
gives a test loss of 37 nats? What does that say about the network at
the peak?
- ★★ A toy of the competition, labelled a toy. Let
be the strength of a “lookup” solution and
the strength of a “circuit” solution. Both fit the training set, so the
training loss is
.
Only the circuit helps on test pairs, so the test loss is
.
Add decay
,
where
makes the lookup more expensive. Find the minimum of the regularized
training loss. Which solution carries the fit, and how does the answer
depend on
?
- ★★★ In the toy of Exercise 6, start at
,
and run gradient flow. Show that the lookup is abandoned only as fast as
decay allows, and find the time constant. What in the toy plays the role
of the wait?
- Computational. From
add-transformer-s0.json, plot train
and test loss on a log scale against step. Mark
,
the peak, and
.
Then overlay the undefended decay curve
,
suitably scaled, and describe what it does and does not match.
- Computational, open. For all six seeds, find the step at which test
loss first falls back below its value at
.
Record it alongside
.
The chapter makes no prediction about how the two relate.
Hints. - (2) The correct answer gets
.
- (4) Take logarithms and use
.
- (6) At the minimum the gradients in
and
match the decay in each; the cheaper direction takes more of the fit. -
(7) Linearize about the minimum from Exercise 6; the slower of the two
decay rates is the time constant.
Summary and bridge to Chapter
6
Grokking is a competition between a lookup table and a circuit that
fit the same training set.
- The lookup table comes first, because each stored pair helps on its
own. In 037 it is done by step 300.
- A memorizer is confidently wrong: 037’s test loss reaches about four
times the guessing loss at the fit and eight times at its peak.
- Weight decay, on a clock of about a thousand steps in 037, favors
the solution that fits at lower norm. Nanda et al. report that without
it their networks did not grok.
- In their progress measures the circuit forms before the jump, which
happens during cleanup. Run 037 logged neither progress measure nor
weight norm, so the repo has not measured the phases.
- Across 037’s seeds the grokking step varies by a factor of 2.9 with
identical hyperparameters.
Next. Every group in this course so far has been
cyclic, so one discrete log was enough. Chapter 6 takes the units of
,
which form
,
with no single discrete log at all. The circuit then has to live on a
product of circles, and the Chinese Remainder Theorem becomes a
statement about an embedding. Its lab is P-031.
References for Chapter 5
- Power, A., Burda, Y., Edwards, H., Babuschkin, I. and Misra, V.
Grokking: generalization beyond overfitting on small algorithmic
datasets. arXiv:2201.02177 (v1, 6 Jan 2022).
- Nanda, N., Chan, L., Lieberum, T., Smith, J. and Steinhardt, J.
Progress measures for grokking via mechanistic interpretability.
arXiv:2301.05217 (v3, 19 Oct 2023). §5.1–5.3, Appendix D.1.
- Mix-in: Welch Labs, “The most complex model we actually understand,”
YouTube, https://www.youtube.com/watch?v=D8GOeCFFby4 (link checked
2026-09-29).