Chapter 12 of 0–21

The Fourier multiplication algorithm

How a network can add modulo 113 using rotations — derived step by step.

Sources for this chapter

[1] Sec. 3.1, App. B
[3] “Summary of the algorithm”

We now have every ingredient to state the paper's central claim. Nanda et al. write [1]:

“We claim that the learned networks use the following algorithm”— Nanda et al. 2023, Sec. 3.1, p. 3 · source

— and the algorithm, which they call Fourier multiplication, is the one below. This chapter derives it step by step and runs it by hand on our example; chapters 13–16 then check it against the trained network.

The algorithm, step by step

Fix one frequency k and write w = w_k = 2\pi k/113.

Step 1 — embed: numbers become angles. Map a to \cos(w a) and \sin(w a), and b to \cos(w b) and \sin(w b). In the network this is done by the embedding matrix W_E (chapter 13 shows that its rows, read as functions of the number, are mixtures of such waves).

Step 2 — combine: add the angles with the identities of chapter 9.

\begin{aligned}\cos(w(a+b)) &= \cos(wa)\cos(wb) - \sin(wa)\sin(wb),\\ \sin(w(a+b)) &= \sin(wa)\cos(wb) + \cos(wa)\sin(wb).\end{aligned}

These need products of two numbers; in the network they are computed by the attention layer and the MLP (chapter 14 shows how a ReLU multiplies).

What these products look like over all inputs: each product of two waves is a checkerboard, and their difference depends only on a + b, so it forms stripes along the lines a + b = constant. Press ▶ to run through the frequencies.

Step 3 — score each candidate answer c with one more identity (the subtraction form of chapter 9):

\cos(w(a+b-c)) = \cos(w(a+b))\cos(wc) + \sin(w(a+b))\sin(wc).

For fixed (a, b) this is a linear function of the two numbers \cos(w(a+b)) and \sin(w(a+b)), with weights \cos(wc) and \sin(wc) that depend only on c — exactly what one matrix can do. In the network it is the neuron–logit map W_L (chapter 15). The score is largest, equal to 1, when a + b - c is a multiple of 113, i.e. when c = (a + b) mod 113.

The steps so far on a clock: the hands for a and b add up to the hand for a + b, and the score of a candidate c is the cosine of the angle between c's hand and a + b's. Press ▶ next to c to sweep it around.

Step 4 — add up several frequencies. The network does steps 1–3 for a few key frequencies k ∈ \mathcal{K} at once and adds the results with weights \alpha_k:

\Logits(a,b)_c \approx \sum_{k\in\mathcal K} \alpha_k \cos\big(w_k(a+b-c)\big).

Running it by hand

Worked example

(a, b) = (100, 50), frequency k = 14, so w ≈ 0.7785.

Step 1: \cos(100w) = -0.7680, \sin(100w) = 0.6404, \cos(50w) = 0.3406, \sin(50w) = 0.9402.

Step 2: \cos(150w) = (-0.7680)(0.3406) - (0.6404)(0.9402) = -0.2616 - 0.6021 = -0.8637, and \sin(150w) = (0.6404)(0.3406) + (-0.7680)(0.9402) = 0.2181 - 0.7221 = -0.5040. (Directly: cos(150w) = −0.8637, sin(150w) = −0.5040. ✓)

Step 3: for c = 37, \cos(37w) = -0.8637 and \sin(37w) = -0.5040 (37 sits on the same clock position as 150, since 150 − 113 = 37), so the score is (-0.8637)^2 + (-0.5040)^2 = 1.0000. For c = 38 the score is 0.7120, for c = 0 it is −0.8637.

But for c = 45 the score is 0.9985 — almost as high as the right answer! At frequency 14, a step of 8 numbers turns the clock by 14\cdot8 = 112 \equiv -1 steps out of 113, nearly a full turn, so 45 looks almost like 37.

Why several frequencies: constructive interference

One wave alone cannot single out the answer reliably: it has near-peaks elsewhere (above, c = 45). Different frequencies have their near-peaks in different places, but all of them peak exactly at c = (a + b) mod 113. Summing them is constructive interference at the right answer and destructive interference (partial cancellation) everywhere else [1] (App. B, Fig. 8).

Worked example

Interference with the paper's key frequencies {14, 35, 41, 42, 52} (all weights 1): at c = c* every cosine equals 1, so the sum is 5.0; at every other c the five waves disagree and the largest sum anywhere else is only 2.59. Softmax then turns this gap into a confident answer.

Press ▶ next to a to move the input: the peak of the sum follows a + b.

On the real network, for the selected input and step (try different a, b in the dashboard), the logits and their best fit by a sum of cosines over the run's key frequencies look like this:

The model's answer for the selected input

Top: the 113 logits (bars; the correct answer highlighted) and their fit by a sum of cosines. Bottom: the softmax probabilities.

Loading…

How to read it · the math

\Logits(a,b)_c \approx \beta + \sum_{k\in\mathcal K} \alpha_k \cos\big(w_k(a+b-c)\big)

<p class="fig-meta">
  <a href="../chapters/12-fourier-multiplication.html">Learn step by step →</a>
   · Source: <a class="cite" href="../references.html#nanda2023" title="Nanda et al. 2023">[1]</a> Sec. 3.1, eq. (1)
  · Computed by <code>browser (GrokModel) + least squares</code> (<a href="https://github.com/BurnyCoder/visualizing-grokking-modular-addition">code</a>)
</p>

Constructive interference

Each key frequency's cosine wave α_k cos(w_k(a+b−c)) over all candidate answers c, and their sum; switch frequencies off and sweep c to see where they agree.

Loading…

How to read it

<div class="fig-read"><p>The inset chains one hand per switched-on frequency, tip to tail. Each hand has length α_k and angle w_k(a+b−c), so its horizontal part is that frequency's contribution to the logit of c, and the chain's end is the sum. At c = a + b every angle is 0: the hands line up and the sum is largest. At other c they point in different directions and largely cancel.</p>

Step	Mathematical operation	Network component	Evidence
1	a ↦ cos(w_k a), sin(w_k a)	embedding W_E	ch. 13: W_E sparse in the Fourier basis
2	products → cos/sin(w_k(a+b))	attention + MLP (ReLU)	ch. 14: neurons are degree-2 polynomials of one frequency
3	cos(w_k(a+b−c)) for every c	neuron–logit map W_L	ch. 15: W_L has rank ≈ 10 and reads cos/sin(w_k(a+b))
4	sum over key frequencies	unembedding (adding into the logits)	ch. 15–16: logits ≈ Σ α_k cos(w_k(a+b−c)); ablations