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.

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 $\mathbb{Z}/n$ is a way of placing the clock onto the unit circle that respects addition. There are exactly $n$ of them, and that one fact is the whole engine of this course.

Start with the clock $\mathbb{Z}/n = \{0, 1, \dots, n-1\}$, added mod $n$. Write $\omega = e^{2\pi i/n}$ for the first $n$-th root of unity.

Definition. A character of $\mathbb{Z}/n$ is a map $\chi$ from $\mathbb{Z}/n$ to the nonzero complex numbers with

$$\chi(a+b) = \chi(a)\,\chi(b) \quad \text{for all } a, b \in \mathbb{Z}/n.$$

Theorem 1.1. The characters of $\mathbb{Z}/n$ are exactly the $n$ maps

$$\chi_k(a) = \omega^{ka} = e^{2\pi i k a / n}, \qquad k = 0, 1, \dots, n-1.$$

Proof. A character is fixed by its value on 1, since $\chi(a) = \chi(1)^a$. Also $\chi(1)^n = \chi(n) = \chi(0) = 1$, so $\chi(1)$ is an $n$-th root of unity, one of $\omega^0, \dots, \omega^{n-1}$. Each choice gives a character (Exercise 4). ∎

Picture $\chi_k$ as a hand that moves $k$ notches around the circle for every tick of the clock. $\chi_0$ never moves. $\chi_1$ goes round once as $a$ runs from $0$ to $n-1$, and $\chi_k$ goes round $k$ times. The number $k$ is the frequency, the same word as in music.

The real form. A network stores real numbers, not complex ones. Splitting $\omega^{ka}$ into its real and imaginary parts gives the pair

$$\bigl(\cos(2\pi k a/n),\ \sin(2\pi k a/n)\bigr).$$

So “the network learned frequency $k$” means that two coordinates somewhere in its embedding of the token $a$ hold this pair. Section 1.6 puts that sentence to work.

1.2 Orthogonality and the Fourier transform

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).

$$\frac{1}{n}\sum_{a=0}^{n-1} \chi_j(a)\,\overline{\chi_k(a)} = \begin{cases} 1 & j = k \\ 0 & j \neq k \end{cases}$$

Proof. The sum is $\sum_a z^a$ with $z = \omega^{j-k}$. If $j = k$ then $z = 1$ and the sum is $n$. Otherwise $z \neq 1$ but $z^n = 1$, and the geometric series gives $(z^n - 1)/(z - 1) = 0$. ∎

So the $n$ characters form an orthonormal basis for functions on $\mathbb{Z}/n$. The coordinates of $f$ in that basis are its Fourier coefficients:

$$\hat f(k) = \frac{1}{n}\sum_{a} f(a)\,\overline{\chi_k(a)}, \qquad f(a) = \sum_{k} \hat f(k)\,\chi_k(a).$$

The energy is shared out exactly (Parseval):

$$\sum_k |\hat f(k)|^2 = \frac{1}{n}\sum_a |f(a)|^2.$$

Worked by hand: n = 5

Take $f = (1, 1, 0, 0, 0)$, a bump on the first two hours. Then $\hat f(k) = (1 + \omega^{-k})/5$, and since $|1 + e^{i\theta}| = 2|\cos(\theta/2)|$,

$$|\hat f(k)| = \tfrac{2}{5}\,\bigl|\cos(\pi k/5)\bigr|.$$

The squares sum to 0.4000, and the right side of Parseval is $2/5 = 0.4$. It balances.

Notice two things: - The mass sits at low frequency, because a bump two hours wide is smooth on a five-hour clock. - $k$ and $n - k$ have the same size, because $f$ is real. Frequency $k$ and frequency $-k$ 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, $\{k, n-k\}$, 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 $\mathbb{Z}/n$, define $(f * g)(c) = \sum_a f(a)\, g(c - a)$. With the normalization of Section 1.2,

$$\widehat{f * g}(k) = n\,\hat f(k)\,\hat g(k).$$

Proof. Expand the sum, substitute $b = c - a$, and use $\chi_k(a + b) = \chi_k(a)\chi_k(b)$ to split the double sum into a product. ∎

Here is why this matters. The answer to “$a + b$” is the point mass at $a + b$, which is the convolution of the point masses at $a$ and at $b$. 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 $2\pi ka/n$ for its input $a$, and the same for $b$. It can then form

$$\cos\frac{2\pi k(a+b)}{n} = \cos\frac{2\pi k a}{n}\cos\frac{2\pi k b}{n} - \sin\frac{2\pi k a}{n}\sin\frac{2\pi k b}{n},$$

and likewise the sine. This is $\chi_k(a)\chi_k(b) = \chi_k(a + b)$ 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 $c$, score every candidate by how well it lines up with $a + b$, summed over a few frequencies $k_1, \dots, k_K$:

$$\text{logit}(c) = \sum_{i=1}^{K} \cos\frac{2\pi k_i\,(a+b-c)}{n}.$$

At $c = a + b$ every term is 1, so the score is $K$. 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 $p$ form a second clock, with $p - 1$ 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 $p$ prime, the group $(\mathbb{Z}/p)^\times$ of nonzero residues under multiplication is cyclic of order $p - 1$.

Proof sketch. A polynomial of degree $d$ has at most $d$ roots mod $p$, so for every $d$ dividing $p - 1$, at most $d$ elements satisfy $x^d = 1$. A counting argument over element orders then forces an element of order exactly $p - 1$. ∎

An element $g$ of order $p - 1$ is a primitive root. Every nonzero $a$ is a power of $g$, $a = g^j$. The exponent $j = \operatorname{ind}_g(a)$ is the discrete log of $a$, defined mod $p - 1$. The key identity is

$$\operatorname{ind}_g(ab) = \operatorname{ind}_g(a) + \operatorname{ind}_g(b) \;(\mathrm{mod}\;p-1).$$

The discrete log turns the multiplicative clock into the additive clock $\mathbb{Z}/(p - 1)$, 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 $3^0 = 1$, $3^1 = 3$, $3^2 = 2$, $3^3 = 6$, $3^4 = 4$, $3^5 = 5$, and $3^6 = 1$ closes the loop.

Check: $2 \times 4 = 8 = 1$ mod 7, and $\operatorname{ind} 2 + \operatorname{ind} 4 = 6 = 0$ mod 6, which is $\operatorname{ind} 1$. 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 $a = 1, 2, 3, \dots$

Dirichlet characters. The characters of $(\mathbb{Z}/p)^\times$ are

$$\chi_k(a) = e^{2\pi i\,k\,\operatorname{ind}_g(a)/(p-1)}, \qquad k = 0, \dots, p-2.$$

The order of $\chi_k$ is $(p - 1)/\gcd(k, p - 1)$. Changing the primitive root multiplies every frequency $k$ by a unit mod $p - 1$. That relabels the characters but leaves each order unchanged.

The Legendre symbol lives here. Take $k = (p - 1)/2$. Then $\chi(a) = (-1)^{\operatorname{ind} a}$, which is $+1$ exactly when $a$ is a square mod $p$. For $p = 7$ 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 $g$. 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 $\mathbb{Z}/p$, 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 $p$ carry both clocks at once: addition on all of $\mathbb{Z}/p$, and multiplication on the nonzero part. The Gauss sum of a multiplicative character $\chi$ pairs it against the first additive character:

$$\tau(\chi) = \sum_{a=1}^{p-1} \chi(a)\, e^{2\pi i a/p}.$$

Theorem 1.4. If $\chi$ is not the trivial character, then $|\tau(\chi)| = \sqrt p$.

Proof. Expand $|\tau|^2$ as a double sum over $a$ and $b$, and substitute $a = bc$ (as $a$ runs over the nonzero residues, so does $c$):

$$|\tau|^2 = \sum_{b}\sum_{c} \chi(c)\, e^{2\pi i\, b(c-1)/p} = \sum_{c} \chi(c) \sum_{b=1}^{p-1} e^{2\pi i\, b(c-1)/p}.$$

The inner sum is $p - 1$ when $c = 1$ and $-1$ otherwise, because the full sum over all $b$, including $0$, vanishes. So $|\tau|^2 = (p - 1) - \sum_{c\neq1} \chi(c)$. The values of a nontrivial character sum to zero over the group (Theorem 1.2 with $j = 0$), so $\sum_{c\neq1} \chi(c) = -\chi(1) = -1$. That leaves $|\tau|^2 = p$. ∎

What it means. Replace $e^{2\pi ia/p}$ by $e^{2\pi ika/p}$ for any $k \neq 0$ and you get $\overline{\chi(k)}\,\tau(\chi)$, which has the same size $\sqrt p$. So extend $\chi$ to all of $\mathbb{Z}/p$ by setting $\chi(0) = 0$. Its Fourier transform, taken at length $p$, 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 $\chi = (+1, -1, -1, +1)$ on $a = 1, 2, 3, 4$. With $\zeta = e^{2\pi i/5}$,

$$\tau = \zeta - \zeta^2 - \zeta^3 + \zeta^4 = 2\cos 72^\circ - 2\cos 144^\circ.$$

Now $2\cos 72^\circ = 0.618$ and $-2\cos 144^\circ = 1.618$, which are $\varphi - 1$ and $\varphi$, where $\varphi$ is the golden ratio. The sum is $2\varphi - 1 = 2.236 = \sqrt 5$. 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-$p$ transform over $\mathbb{Z}/p$, with the zero slot included. A network trained on $a \times b$ has tokens only for $1, \dots, p-1$. Run 037 therefore read those tokens with a length-$(p-1)$ 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 $a \times b$ 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 $a + b$ mod 113, which should reproduce the published additive result (Nanda et al. 2023, as cited in the record). The main task is $a \times b$ on the 112 nonzero residues. - The network and training. One layer, 4 heads, $d_\text{model} = 128$. 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 ($a = 1, 2, 3, \dots$) and discrete log (primitive root $g = 3$).

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 $10/56 = 0.179$. 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 $\chi(-1) = -1$, by growing margins (odd share 0.536, 0.661, 0.763; noise gives 0.505). Since $-1 = g^{56}$, a character is odd exactly when its frequency $k$ 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.

1. ★ Write the 4 × 4 character table of $\mathbb{Z}/4$ (its entries are $1, i, -1, -i$). Check by hand that rows 1 and 3 are orthogonal.
2. ★ Redo the $n = 5$ computation of Section 1.2 for $f = (1, 0, 1, 0, 0)$. Where does the mass sit now, and why? Check Parseval.
3. ★ 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.
4. ★★ Fill in the “each choice works” step of Theorem 1.1: show that $\chi(a) = \zeta^a$ is well defined on $\mathbb{Z}/n$ when $\zeta^n = 1$.
5. ★★ The perfect-interference exercise. On $\mathbb{Z}/7$, compute $\text{logit}(d) = \sum \cos(2\pi k d/7)$ for $d = 0, \dots, 6$ with the frequencies $\{1\}$, then $\{1, 3\}$, then $\{1, 2, 3\}$. Record the margin between $d = 0$ and the best wrong answer each time. The last case has a surprise; explain it with Theorem 1.2.
6. ★★ Compute the Gauss sum of the Legendre symbol mod 3 exactly, then mod 7. In the second case, show $\tau^2 = -7$ without using decimals.
7. ★★ Prove that the order of a Dirichlet character does not depend on the primitive root chosen, and that the Legendre symbol’s frequency $(p - 1)/2$ never moves.
8. ★★★ Rebuild the proof of Theorem 1.4 from memory.
9. ★★★ Using $\chi(-1)$, show that the quadratic Gauss sum is real when $p \equiv 1$ mod 4 and purely imaginary when $p \equiv 3$ mod 4. (Which sign it has is a different, famous story: it took Gauss years.)
10. Computational. Plot the Legendre symbol mod 113 twice: in ordinary order, then in discrete-log order with $g = 3$. 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.)
11. Computational, open. Run 037’s seeds leaned toward odd characters. Propose one reason why a network trained on $a \times b$ mod 113 might prefer $\chi(-1) = -1$, 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 $\{1, 2, 3\}$, add the conjugates 4, 5 and 6 and use orthogonality. - (6) Mod 3 the answer is $i\sqrt3$. - (9) Conjugate $\tau$ and substitute $a \to -a$. - (10) The length-113 reading of the Legendre symbol should come out at exactly $10/56$. 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.

• $\mathbb{Z}/n$ has $n$ characters, $\chi_k(a) = \omega^{ka}$; 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 $\mathbb{Z}/p$, the two clocks are maximally incompatible: $|\tau| = \sqrt p$. 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 $\mathbb{Z}/7$ addition, small enough to read its weights directly, and asks whether the character table of $\mathbb{Z}/7$ 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 $x \in \mathbb{R}^d$, weights $w \in \mathbb{R}^d$ and a bias $b$, and outputs

$$y = \begin{cases} 1 & w \cdot x + b \ge 0 \\ 0 & \text{otherwise.} \end{cases}$$

The set $w \cdot x + b = 0$ is a hyperplane, and the perceptron answers one question: which side of it is $x$ 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 $(x, t)$. If it answers correctly, do nothing. If it is wrong, move the hyperplane toward the right answer:

$$w \leftarrow w + (t - y)\,x, \qquad b \leftarrow b + (t - y).$$

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 $\{0,1\}^2$.

Proof. Suppose one did. Then $(0,0) \mapsto 0$ gives $b < 0$. The inputs $(1,0)$ and $(0,1)$ give $w_1 + b \ge 0$ and $w_2 + b \ge 0$; adding them, $w_1 + w_2 + b \ge -b > 0$. But $(1,1) \mapsto 0$ needs $w_1 + w_2 + b < 0$. 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 $\mathbb{Z}/2$: $a \oplus b = a + b \bmod 2$.

Now read it in characters, as in Chapter 1. $\mathbb{Z}/2$ has two characters, and the nontrivial one is $s(a) = (-1)^a$, sending $0 \mapsto +1$ and $1 \mapsto -1$. Section 1.3 says the group law becomes multiplication:

$$s(a + b) = s(a)\, s(b).$$

So in the $\pm 1$ encoding, XOR is the product $s_a s_b$. Theorem 2.1 says a linear unit cannot compute it, and here is a sharper reason.

Theorem 2.2. On the four inputs $(s_a, s_b) \in \{\pm1\}^2$, the functions $1$, $s_a$, $s_b$ and $s_a s_b$ are orthogonal. So the best linear fit to XOR, in least squares, is the zero function.

Proof. These four functions are the characters of $\mathbb{Z}/2 \times \mathbb{Z}/2$, and characters of a finite abelian group are orthogonal (Theorem 1.2 applied coordinate by coordinate). Least squares projects $s_a s_b$ onto the span of $1, s_a, s_b$, 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 $z = w \cdot x + b$ and passing it through a nonlinear activation $\sigma$. The output is a linear combination of the hidden values. This chapter uses the ReLU,

$$\operatorname{ReLU}(z) = \max(0, z).$$

XOR with two hidden units, on $\{0,1\}^2$:

$$h_1 = \operatorname{ReLU}(x_1 + x_2), \quad h_2 = \operatorname{ReLU}(x_1 + x_2 - 1), \quad y = h_1 - 2h_2.$$

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

$$xy = \tfrac{1}{4}\bigl((x+y)^2 - (x-y)^2\bigr).$$

So two hidden units with activation $\sigma(z) = z^2$, one fed $x + y$ and one fed $x - y$, compute the product exactly. A ReLU is not a square, but it can make one well enough where it is needed.

In the $\pm1$ encoding it is exact. Since $s^2 = 1$, $(s_a + s_b)^2 = 2 + 2 s_a s_b$, and $|s_a + s_b|$ is 2 or 0. So

$$s_a s_b = |s_a + s_b| - 1, \qquad |u| = \operatorname{ReLU}(u) + \operatorname{ReLU}(-u).$$

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 $\mathbb{Z}/n$, the angle-addition trick of Section 1.3 needs products like $\cos\alpha \cos\beta$. The product-to-sum identity

$$\cos\alpha\cos\beta = \tfrac{1}{2}\bigl(\cos(\alpha - \beta) + \cos(\alpha + \beta)\bigr)$$

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 $\theta$ produces an output $y(x; \theta)$. A loss $L$ measures how wrong that output is. This chapter uses the squared loss for one example with target $t$:

$$L = \tfrac{1}{2}(y - t)^2.$$

(Classification over $n$ answers uses softmax and cross-entropy instead; Chapter 3 introduces them.)

Gradient descent moves every parameter a small step downhill:

$$\theta \leftarrow \theta - \eta\, \frac{\partial L}{\partial \theta},$$

where $\eta$ is the learning rate. Near a point, the loss changes by about $-\eta \lVert \nabla L \rVert^2$ 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 $\partial L / \partial \theta$ 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.

$$z_1 = w_{11}x_1 + w_{12}x_2 + b_1, \quad z_2 = w_{21}x_1 + w_{22}x_2 + b_2,$$ $$h_i = \operatorname{ReLU}(z_i), \quad y = v_1 h_1 + v_2 h_2 + c, \quad L = \tfrac12 (y - t)^2.$$

Starting weights: $w_{11} = 0.5$, $w_{12} = -0.5$, $b_1 = 0$; $w_{21} = 1$, $w_{22} = 1$, $b_2 = -0.5$; $v_1 = 1$, $v_2 = -1$, $c = 0$.

The example: the XOR input $x = (1, 0)$ with target $t = 1$.

Forward pass

Backward pass

Write $\delta_q = \partial L / \partial q$. Start at the loss and walk back, reusing each $\delta$ once it is known.

Notice that $\delta_{w_{12}} = \delta_{w_{22}} = 0$: an input that is off carries no blame.

One step, then check

Take $\eta = 0.1$. The new weights are $v_1 = 1.05$, $v_2 = -0.95$, $c = 0.1$, $w_{11} = 0.6$, $b_1 = 0.1$, $w_{21} = 0.9$ and $b_2 = -0.6$; $w_{12}$ and $w_{22}$ are unchanged. Run the forward pass again:

$$z_1 = 0.7,\ h_1 = 0.7; \quad z_2 = 0.3,\ h_2 = 0.3; \quad y = 1.05(0.7) - 0.95(0.3) + 0.1 = 0.55.$$

The loss falls from $0.5$ to $\tfrac12(0.45)^2 = 0.10125$.

Compare the first-order prediction. The squared gradient sums to $0.25 + 0.25 + 1 + 1 + 0 + 1 + 1 + 0 + 1 = 5.5$, so the prediction is a drop of $\eta \cdot 5.5 = 0.55$. 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 $z^{(\ell)}$, weights $W^{(\ell)}$ and activation $\sigma$,

$$\delta_{z^{(\ell)}} = \bigl(W^{(\ell+1)\top}\, \delta_{z^{(\ell+1)}}\bigr) \odot \sigma'\bigl(z^{(\ell)}\bigr), \qquad \frac{\partial L}{\partial W^{(\ell)}} = \delta_{z^{(\ell)}}\, h^{(\ell-1)\top}.$$

The table above is this rule with one hidden layer and every sum written out.

2.7 Lab: a tiny network on $\mathbb{Z}/7$ 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 $a$ and $b$ as one-hot vectors in $\mathbb{R}^7$ 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 $a$. Each hidden unit has a row of 7 numbers, one per token. Take that row’s length-7 DFT over the token index $a = 0, \dots, 6$. Because the row is real, fold the frequencies into the DC term and three conjugate-pair bins, $\{1, 6\}$, $\{2, 5\}$ and $\{3, 4\}$, 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 $a$ and for input $b$ 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

1. ★ Run the perceptron rule by hand on AND, starting from $w = (0, 0)$, $b = 0$, cycling through the four inputs in order. How many mistakes does it make before it stops?
2. ★ Redo the proof of Theorem 2.1 for the $\pm1$ encoding: no $w, b$ give $\operatorname{sign}(w \cdot s + b) = s_a s_b$ on all four inputs.
3. ★ Verify $s_a s_b = \operatorname{ReLU}(s_a + s_b) + \operatorname{ReLU}(-s_a - s_b) - 1$ on all four inputs. Draw it as a network with two hidden units.
4. ★★ Using polarization, build a network with square activations that computes $\cos\alpha\cos\beta$ from inputs $\cos\alpha$ and $\cos\beta$. How many hidden units does it need?
5. ★★ Redo the backpropagation of Section 2.6 on the input $(1, 1)$ with target $t = 0$, from the same starting weights. Which weights get no gradient, and why?
6. ★★ The dead unit. Change $b_2$ to $-2$ and redo the forward and backward passes for $x = (1, 0)$. Show that no gradient reaches $w_{21}$, $w_{22}$ or $b_2$. What does that mean for training?
7. ★★ Derive the general rule of Section 2.6 from the chain rule, for one hidden layer written in matrices.
8. ★★★ Show that every continuous piecewise-linear function $\mathbb{R} \to \mathbb{R}$ with finitely many breakpoints is $c_0 + c_1 x + \sum_i a_i \operatorname{ReLU}(x - p_i)$. Deduce that one hidden layer of ReLUs represents it exactly.
9. ★★ Frequency by hand. A hidden unit’s row over the tokens $a = 0, \dots, 6$ is $(1.000, -0.223, -0.901, 0.623, 0.623, -0.901, -0.223)$. Take its inner product with $\cos(2\pi k a/7)$ for $k = 1, 2, 3$, and name the frequency. Why is no sine term needed?
10. Computational. Check the gradients of Section 2.6 by finite differences: nudge each weight by $\pm 10^{-4}$ and compare $(L_+ - L_-)/(2 \cdot 10^{-4})$ with the table.
11. 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 $(1, 1)$ has target 1, so the bias must end up negative. - (5) With $x = (1, 1)$ both inputs are on, so no input carries zero blame this time. Look at the signs. - (6) $\operatorname{ReLU}'(z) = 0$ for $z < 0$. - (9) Orthogonality (Theorem 1.2): only one inner product survives, and it equals $7/2$.

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 $\mathbb{Z}/2$, and in the $\pm1$ encoding it is the character product $s_a s_b$.
• The characters of $\mathbb{Z}/2 \times \mathbb{Z}/2$ are orthogonal, so a linear unit has zero correlation with XOR.
• Two ReLU units compute $s_a s_b$ exactly; squares and polarization give products in general.
• Backpropagation is the chain rule, with each $\delta$ 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, $a + b \bmod 113$, 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.

Chapter 3 — Embeddings and Attention: The One-Layer Transformer, Piece by Piece

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, $a$, $b$ and $=$, each a one-hot vector. Position 3 is where the answer is read.

The embedding. A matrix $W_E$ with one column per token value sends the one-hot vector of $a$ to that column, a vector in $\mathbb{R}^d$. 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 $p_i$ is added, so the network can tell the first slot from the second:

$$x_i^{(0)} = W_E\, t_i + p_i.$$

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 $a$ must contain the pairs $\bigl(\cos(2\pi k a/p), \sin(2\pi k a/p)\bigr)$ for a few frequencies $k$. Nanda et al. report exactly this for their mainline model: a Fourier transform along the token index of $W_E$ shows significant norm at 6 frequencies, of which 5, $k \in \{14, 35, 41, 42, 52\}$, 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 $W_Q$, a key $W_K$ and a value $W_V$, plus an output map $W_O$. The position reading (here, $=$) forms a query; every position offers a key and a value.

$$A = \operatorname{softmax}_i\!\Bigl(\bigl(W_K x_i^{(0)}\bigr)^{\!\top} W_Q\, x_3^{(0)}\Bigr), \qquad \text{head output} = W_O \sum_i A_i\, W_V x_i^{(0)}.$$

The softmax (Section 3.5) turns scores into weights $A_i \ge 0$ 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 $A_i$, which depends on the tokens, times a value $W_V x_i$, 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, $a$ and $b$. A softmax over two items is a sigmoid of their score difference, $\sigma(C[a] - C[b])$. The score $C$ was close to a single wave, $C[a] \approx F\cos(w_k a) + G\sin(w_k a)$, and the sigmoid operated in a nearly linear range. So each head’s weight is roughly $0.5$ plus a single-frequency wave in $a$ minus the same wave in $b$. 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.

$$\text{MLP} = \operatorname{ReLU}\bigl(W_{\text{in}}\, x^{(1)}\bigr), \qquad x^{(2)} = W_{\text{out}}\, \text{MLP} + x^{(1)}.$$

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 $\cos\bigl(w_k(a+b)\bigr)$ and $\sin\bigl(w_k(a+b)\bigr)$ 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 $\cos\bigl(w_k(a+b)\bigr)$ or $\sin\bigl(w_k(a+b)\bigr)$. The network had computed Chapter 1’s angle-addition formula in its neurons.

3.4 The unembedding: the interference sum, learned

A final matrix $W_U$ maps the residual stream to one logit per candidate answer $c$. In their model the skip connection around the MLP barely mattered, so the logits are approximately $W_L\,\text{MLP}$ with $W_L = W_U W_{\text{out}}$.

Now the payoff. To score candidate $c$, the network needs

$$\cos\bigl(w_k(a+b-c)\bigr) = \cos\bigl(w_k(a+b)\bigr)\cos(w_k c) + \sin\bigl(w_k(a+b)\bigr)\sin(w_k c).$$

This is linear in the two quantities the MLP produced, with coefficients $\cos(w_k c)$ and $\sin(w_k c)$ that depend only on $c$. So a fixed matrix can do it: the rows of $W_L$ just need to be those cosines and sines. Nanda et al. found $W_L$ 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

$$\text{logit}(c) \approx \sum_{k} \alpha_k \cos\bigl(w_k(a + b - c)\bigr),$$

with this form explaining 95% of the logits’ variance.

This is the interference sum of Section 1.3, with learned weights $\alpha_k$ in place of 1. Every term peaks at $c = a + b$ and the terms disagree elsewhere. Their Appendix B shows why one frequency is not enough: $\cos(2\pi \cdot 14 \cdot 8/113) \approx 0.998$, so frequency 14 alone nearly confuses $c$ with $c \pm 8$. But $\cos(2\pi \cdot 35 \cdot 8/113) \approx -0.990$, 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 $\ell_c$, one per answer. Softmax turns them into probabilities:

$$p_c = \frac{e^{\ell_c}}{\sum_{c'} e^{\ell_{c'}}}.$$

Adding the same constant to every logit changes nothing, so only differences between logits matter. Cross-entropy for the correct answer $c^\ast$ is

$$L = -\log p_{c^\ast} = \log\Bigl(1 + \sum_{c \neq c^\ast} e^{\ell_c - \ell_{c^\ast}}\Bigr).$$

Its gradient is simple: $\partial L / \partial \ell_c = p_c - [c = c^\ast]$. 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 $\ell_c = A \sum_k \cos\bigl(w_k(a+b-c)\bigr)$, a scale $A$ times the Section 1.3 sum, and let $M$ be the margin between the correct answer and the best wrong one. Every wrong logit sits at least $AM$ below the right one, so

$$L \le \log\bigl(1 + (p - 1)\, e^{-AM}\bigr).$$

Two ways to lower the loss follow. Grow the scale $A$, or grow the margin $M$ 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 $W_L$, 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:

$$x_i^{(0)} = W_E t_i + p_i, \quad x^{(1)} = x_3^{(0)} + \sum_{\text{heads}} W_O W_V \Bigl(\sum_i A_i x_i^{(0)}\Bigr), \quad x^{(2)} = x^{(1)} + W_{\text{out}} \operatorname{ReLU}\bigl(W_{\text{in}} x^{(1)}\bigr), \quad \ell = W_U x^{(2)}.$$

Read the last column carefully. Run 037 saved only $W_E$ — 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 $a + b \bmod 113$ 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.

What 037 measured on this arm. For each seed, the final $W_E$ was Fourier-transformed along the token index with a length-113 DFT over $a = 0, \dots, 112$. Non-DC power was folded into 56 conjugate-pair bins $\{k, 113-k\}$, and the share in the top 10 bins was recorded.

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 $W_E$, 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 $\{14, 35, 41, 42, 52\}$. 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 $k \in \{2, 9, 19, 31\}$, $\{40, 44, 53\}$, $\{31, 45, 49, 52\}$ and $\{17, 32, 42, 51\}$ — 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 $W_E$ is unmeasured here, and cannot be recovered from what was saved. Only $W_E$ was written out (R26_training_dynamics/emb/*.npy), so whether 037’s heads, neurons and $W_L$ 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:

1. A retrain that saves every weight, not just $W_E$ — the model’s whole state_dict at the end, and ideally at intervals, since Chapter 5’s question is when the circuit forms.
2. A check that the retrained $W_E$ reproduces the saved one, seed for seed, before any head or $W_L$ 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 $10^{-6}$ at step 0 and have already diverged to $2\times10^{-3}$ 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).
3. Unpacking MultiheadAttention into per-head $W_Q, W_K, W_V, W_O$, 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 $W_L$ match the frequencies in its $W_E$, 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

1. ★ Write out the softmax probabilities for logits $(2, 0, 0)$ and for $(12, 10, 10)$. Why are they the same?
2. ★ Show that a softmax over two positions with scores $s_0, s_1$ equals $\sigma(s_0 - s_1)$ and $1 - \sigma(s_0 - s_1)$, where $\sigma(u) = 1/(1 + e^{-u})$.
3. ★ Derive $\partial L / \partial \ell_c = p_c - [c = c^\ast]$ for cross-entropy after softmax.
4. ★★ Checking Appendix B by hand. Compute $\cos(2\pi \cdot 14 \cdot 8/113)$ and $\cos(2\pi \cdot 35 \cdot 8/113)$. Hint: reduce $14 \cdot 8$ and $35 \cdot 8$ mod 113 first, and look at how close each is to $0$ or to $113/2$.
5. ★★ Prove the bound $L \le \log\bigl(1 + (p-1)e^{-AM}\bigr)$ of Section 3.5. When is it an equality?
6. ★★ Expand $\bigl(0.5 + \gamma\cos(w a) - \gamma\cos(w b)\bigr)\bigl(u\cos(wa) + v\cos(wb)\bigr)$. Which terms are degree 2? Which of them are pieces of $\cos\bigl(w(a+b)\bigr)$, via the product-to-sum identity?
7. ★★ Show that if the rows of a matrix $W$ are $\cos(w_k c)$ and $\sin(w_k c)$ for $c = 0, \dots, p-1$, then $W$ applied to $\bigl(\cos(w_k s), \sin(w_k s)\bigr)$ gives the vector $\cos\bigl(w_k(s - c)\bigr)$ over $c$. This is the unembedding of Section 3.4.
8. ★★★ Using Theorem 1.2, show that the logits $\sum_{k \in K} \cos\bigl(w_k(a+b-c)\bigr)$, summed over $K$ together with the conjugates $p - k$, have margin exactly $2|K| + 1$ when $K \cup (p - K)$ covers every nonzero frequency. What is that margin as a fraction of the peak?
9. Computational. For the frequency set $\{14, 35, 41, 42, 52\}$ and $p = 113$, compute the interference sum of Section 3.4 with all $\alpha_k = 1$ over $d = a + b - c = 0, \dots, 112$. Report the margin, then drop one frequency at a time and report which one costs the most margin.
10. 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) $14 \cdot 8 = 112 \equiv -1$, and $35 \cdot 8 = 280 \equiv 54$, which is close to $113/2$. - (8) Chapter 1, Exercise 5 is the case $p = 7$, $K = \{1, 2, 3\}$.

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 $\cos(w_k(a+b))$ and $\sin(w_k(a+b))$ 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, $\theta \mapsto \theta - \eta \nabla L(\theta)$, 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:

$$\theta_{t+1} = G(\theta_t), \qquad G(\theta) = \theta - \eta\,\nabla L(\theta).$$

This is a discrete dynamical system, and the vocabulary of dynamics applies directly.

• A fixed point of $G$ is a point with $\nabla L = 0$: 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 $DG = I - \eta H$, where $H$ is the Hessian of $L$: the fixed point is stable when every eigenvalue of $DG$ 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 $\eta$ makes $G$ track the continuous gradient flow $\dot\theta = -\nabla L(\theta)$, with $t$ steps of size $\eta$ corresponding to time $\eta t$. Section 4.3 works in that limit. Section 4.4 is about what happens when $\eta$ is not small.

4.2 The quadratic: stable exactly when $\eta < 2/\lambda_{\max}$

Take a quadratic loss with its minimum at the origin:

$$L(\theta) = \tfrac12\, \theta^{\top} H \theta, \qquad H \text{ symmetric, eigenvalues } \lambda_1 \ge \dots \ge \lambda_d > 0.$$

Then $\nabla L = H\theta$, and gradient descent is linear: $\theta_{t+1} = (I - \eta H)\,\theta_t$.

Theorem 4.1. Gradient descent on this quadratic converges to 0 from every starting point if and only if $\eta < 2/\lambda_1$.

Proof. Expand $\theta_0$ in the eigenvectors $v_i$ of $H$: $\theta_0 = \sum_i c_i v_i$. Each component evolves on its own, $c_i \mapsto (1 - \eta\lambda_i)\,c_i$, so after $t$ steps it is $(1 - \eta\lambda_i)^t c_i$. This goes to 0 for every $c_i$ exactly when $|1 - \eta\lambda_i| < 1$ for every $i$, that is $0 < \eta\lambda_i < 2$. The binding case is the largest eigenvalue. ∎

Read the factor $1 - \eta\lambda_i$ component by component.

The instability arrives through the sign-flipping row. At $\eta = 2/\lambda_1$ 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 $\lambda_1$ is called the sharpness.

Two consequences: - The slowest direction sets the speed. With $\eta$ just under $2/\lambda_1$, the flattest direction shrinks by $1 - 2\lambda_d/\lambda_1$ per step. The ratio $\lambda_1/\lambda_d$, the condition number, is the price of a curved valley. - On a quadratic the threshold is exact. On a real loss, $H$ changes as $\theta$ moves, so $2/\lambda_1$ 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 $2/\eta$, 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 $x \mapsto W_2 W_1 x$. 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, $\Sigma_{yx} = U S V^{\top}$, with singular values $\sigma_1 \ge \sigma_2 \ge \dots$. If the weights start aligned with the singular vectors, they stay aligned, and each mode $i$ is carried by one scalar strength per layer, $a_i$ in $W_1$ and $b_i$ in $W_2$. The network’s strength on mode $i$ is the product $u_i = a_i b_i$. The modes then decouple, and the loss is, up to a constant,

$$L = \tfrac12 \sum_i (\sigma_i - u_i)^2 = \tfrac12 \sum_i (\sigma_i - a_i b_i)^2.$$

One mode, by hand. Drop the index. Gradient flow on $\tfrac12(\sigma - ab)^2$ gives

$$\dot a = b\,(\sigma - ab), \qquad \dot b = a\,(\sigma - ab).$$