Grokking Modular Arithmetic

Introduction
A simple model of grokking
Constructive vs destructive interference
The analytic solution
Convergence as N grows
Activations for general p
Accuracy vs train fraction
Final thoughts

Introduction

In this post I will explain what grokking is about and why it is interesting. This is a personal take I developed after writing a few papers on the topic, the main being Grokking modular arithmetic.

Grokking is a toy example that captures many complex behaviours encountered in AI training. Only a few have been properly explored, due to a somewhat misguided focus on optimization dynamics.

Many large scale models, including LLMs, show abrupt increases in performance when the latter is measured via a non-differentiable metric such as "accuracy". This naively appears to be in contradiction with smooth scaling laws. It is instructive to see both of these behaviours in a task where the data is completely controllable.
The representations learnt by the networks can be understood completely. The "algorithm" the network is using to add numbers is reduced to high school algebra. Representations learnt by actual LLMs on "everyday" modular arithmetic tasks (i.e. adding days of the week or times of the day) are quite similar.
The sample complexity in the arithmetic tasks shows a phase transition that is still poorly understood. Furthermore, some arithmetic tasks are easy while others are hard. The distinction is not obvious before trying to learn the task.
Algorithmic tasks are a crucial subset of real-world AI data, and it is important to understand the clean and controlled version.

While much of AI research is strongly empirical in nature (as it should be), grokking presents an illuminating playground for studying general phenomena that emerge in AI research. Although a lot of ink has been spilled and FLOPs have been burnt, grokking remains very poorly understood. Below I will explain that optimization dynamics is the most "boring" aspect and focus on sample complexity, task complexity and diversity, as well as scaling laws. Not everything I discuss has been published, so I hope you can learn something new.

Finally, this is not a research paper, it is a personal blogpost. Thus, I will keep my citations to an absolute minimum.

A simple model of grokking

Grokking was first discussed in a paper by OAI, where they showed the following 6 things, which nicely set the stage for most of the research that has happened since:

There is a sudden (much less sudden on a linear time scale) onset of generalization long after overfitting occurred.
There is a phase transition in sample complexity as a function of training dataset size.
Representations learnt by the network can be interpreted using PCA on representations of the output layer.
The test loss is a non-monotonic function of training time.
Corrupting some of the training data does not affect grokking.
Some arithmetic functions are harder to grok than others, while some functions appear to be unlearnable.

When I refer to grokking in the following sections I will mean the above 6 things present simultaneously.

In the following post Nanda visualized representations learnt by the MLP part of the transformer and found that it develops trigonometric-looking activations. In the remainder of the blogpost he attempted to understand the underlying algorithm, but was unable to develop a clear quantitative picture.

Both works focused on Transformer models: a two-layer transformer in OAI's case and a one-layer transformer in Nanda's case. Both optimized CE loss with weight decay and the AdamW optimizer, with and without batch noise. Around the same time other authors developed vibe theories about causes of grokking emphasizing batch noise and weight decay as the two important factors.

I got interested in this phenomenon and focused on developing the simplest possible setting where grokking still occurs on modular addition. I followed the philosophy shared by all theoretical physicists: make things as simple as possible, but not simpler. It turned out that attention, AdamW, weight decay, batch noise as well as CE loss were not necessary for grokking (by grokking I mean all 6 effects observed by OAI) to occur. The only critical ingredient was non-linearity. My first main observation was the following:

In the absence of non-linear feature learning, the network fits all the training data, but does not generalize.

The above is a purely empirical observation: linear multilayer networks as well as non-linear networks in the NTK regime do not generalize on modular arithmetic. Before moving on to the analytically tractable model, I will provide the tools for you to play with the optimization dynamics. You can run an MLP with quadratic activation and MSE loss optimized with the AdamW optimizer. You can choose the width of the network, the training data fraction (recall that for finite $p$ the dataset is finite), LRs for the two layers independently, add weight decay, dropout, add corrupted labels, prune neurons by clicking on activations and more. These tools run in your browser and should allow you to reproduce many of the published results and possibly discover unpublished ones!

width Nlr (layer 1)lr (layer 2)weight decaydropoutcorruptiontrain fractionseed

step 0 train acc — test acc —

Loss (MSE)

train
test

press Run to train in your browser

Accuracy

train
test

press Run to train in your browser

Addition table

correct
wrong
corrupted
train pair

Cell $(a, b)$ shows $(a + b) mod 23$ . Cells colored copper once the network predicts that pair correctly. Train pairs are outlined.

Hidden pre-activations

Each tile is one neuron $j$ , plotted as the heatmap of $(W x)_{j}$ over all $(a, b)$ inputs. At init: noise. After grokking: each tile is either a sum of two cosines, one in $a$ and one in $b$ , of the same frequency $k$ , or noise that does not contribute to the computation. Click a tile to prune that neuron; training continues with the rest.

Constructive vs destructive interference

The analytically tractable model of grokking leverages one elementary fact about plane waves (meaning, sines and cosines). Take $N$ cosines of the same frequency $cos (ω x + ϕ)$ . Here $ω$ is called the frequency and $ϕ$ is called the phase. If their phases are random (say, sampled uniformly between $0$ and $2 π$ ), the sum is a small jittery thing — its amplitude scales as $N$ . If their phases are aligned, every wave reinforces every other and the sum is a clean cosine of amplitude $N$ . In wave theory this is known as destructive and constructive interference. I made a very simple visualiser to show you how this works: Drag the sliders to change phase alignment (meaning how random the phases are) and $N$ . A neural network that adds numbers mod $p$ uses this exact mechanism to pick out the correct answer for addition.

randomaligned

N = 1064

sum amplitude: 6.33 · √10 = 3.16 · 10

The analytic solution

The main idea is to take an architecture that on the one hand shows grokking empirically, but on the other hand is simple enough to compute the forward pass analytically. The simplest nonlinearity that plays nicely with plane waves is the square nonlinearity. It forces independent frequencies to add, which is exactly what we need to do modular addition. So our network function is

f (x) = U (W x)^{2}

Above you got to train it on $a + b (mod p)$ with AdamW and MSE loss. In the simulation above, you saw empirically that, while initially the weights are random, optimization lands on Fourier features. That's what I also saw in the experiments. I also found empirically that the measured frequencies were integer multiples of $\frac{2 π}{p}$ .

To analyze the forward pass we parametrize the activations by plane waves with frequencies given by $k = 1, \dots, (p - 1) /2$ , and with phase offsets $φ_{a}^{(1)}, φ_{b}^{(2)}, φ_{j}^{(3)}$ drawn from $U [0, 2 π)$ :

W_{j a} = cos (\frac{2 π k a}{p} + φ_{a}^{(1)}), W_{j, p + b} = cos (\frac{2 π k b}{p} + φ_{b}^{(2)}) .

W_{j a} = cos (\frac{2 π k a}{p} + φ_{a}^{(1)}),

W_{j, p + b} = cos (\frac{2 π k b}{p} + φ_{b}^{(2)}) .

and

U_{c j} = \frac{1}{N} cos (\frac{2 π k _{j} c}{p} + φ_{j}^{(3)}),

These equations imply that the integer $a$ is represented as a vector of plane waves

a ⟶ [1, cos (\frac{2 π a}{p} + φ_{a}^{(1)}), cos (\frac{2 π 2 a}{p} + φ_{a}^{(1)}), \dots]

While the pair of integers that we want to add $(a, b)$ is mapped to

(a, b) ⟶ [1, cos (\frac{2 π a}{p} + φ_{a}^{(1)}) + cos (\frac{2 π b}{p} + φ_{b}^{(2)}), cos (\frac{2 π 2 a}{p} + φ_{a}^{(1)}) + cos (\frac{2 π 2 b}{p} + φ_{b}^{(2)}), \dots]

After the squaring nonlinearity and using the high school algebra of the sort $2 cos α cos β = cos (α + β) + cos (α - β)$ we get exactly the term we really really want

cos (\frac{2 π ( a + b )}{p} + φ_{a}^{(1)} + φ_{b}^{(2)})

where our integers appear as a sum $a + b$ .

There are many other terms that contain $a - b$ , $2 a$ and $2 b$ . We want the terms with $a + b$ to survive and all other terms to cancel out, but we are not yet ready to do it. We need to apply the second layer before cancelation happens. After we apply the second layer and use the same identity we get 9 distinct terms. One of those terms is what we need

cos (\frac{2 π ( a + b - c )}{p} + φ_{a}^{(1)} + φ_{b}^{(2)} - φ_{c}^{(3)})

while other terms involve $a - b$ , $2 a$ and $2 b$ and various other linear combinations of phases. Now we can invoke interference. The term we want interferes constructively if we set

φ_{a}^{(1)} + φ_{b}^{(2)} - φ_{c}^{(3)} = 0

while the terms we do not want interfere destructively if we say that all the phases are uniformly random between $0$ and $2 π$ . The output of the network is then given by

f_{c} (a, b) = k \sum cos (\frac{2 π k ( a + b - c )}{p}) + ϵ_{N} \approx δ_{c, (a + b) mod p}

The sum converges to a one-hot vector with a single non-zero entry at the position where $a + b = c$ . The term $ϵ_{N} = O (N^{- 1})$ is the effect of finite width and will be important to understand scaling laws as a function of parameter count.

This is a rare example of a neural network that is fully interpretable (by the standards of rigor accepted in theoretical physics): we can perform the exact same computation as the network by hand and get the same result. Very similar activations have been found in larger LLMs when prompted about naturally occurring modular addition problems.

The network obtained by optimization does a bit better because it can also learn to optimize the phases and amplitudes of the plane waves.

Symbolic verification

The derivation above is presented in Grokking modular arithmetic, but it can also be carried out symbolically inside your browser. The Python script below builds the hidden pre-activation $z_{j} = cos (2 π k a / p + φ_{a}^{(1)}) + cos (2 π k b / p + φ_{b}^{(2)})$ of a single neuron with random phase offsets $φ_{a}^{(1)}, φ_{b}^{(2)}$ , squares it, multiplies by $U_{c j} = \frac{2}{N} cos (2 π k c / p + φ_{c}^{(3)})$ with the phase-locking construction $φ_{c}^{(3)} = φ_{a}^{(1)} + φ_{b}^{(2)}$ , and then sorts the resulting cosines into the single term that is free of the random phases and the terms that the sum over neurons destroys. Pressing Run downloads Pyodide (~10 MB, one-off) and SymPy, executes the code in a Web Worker, and renders the LaTeX it produces below.

SymPy source (Python)

from sympy import (
    symbols, Symbol, cos, sin, pi, Rational, Sum, KroneckerDelta, Add,
    simplify, expand, integrate, latex, Integer
)
from sympy.simplify.fu import TR7, TR8
import json
import numpy as np

# Symbolic variables. Phases follow the notation of the main text:
# varphi^(1)_a and varphi^(2)_b are the random phase offsets of the
# two input waves, varphi^(3)_c is the read-out phase.
k, a, b, c = symbols('k a b c', integer=True)
p_sym = symbols('p', integer=True, positive=True)
N_sym = symbols('N', integer=True, positive=True)
phi1 = Symbol(r'\varphi^{(1)}_a', real=True)
phi2 = Symbol(r'\varphi^{(2)}_b', real=True)

# Phase locking: varphi^(3)_c = varphi^(1)_a + varphi^(2)_b.
phi3 = phi1 + phi2

A = 2*pi*k*a/p_sym + phi1
B = 2*pi*k*b/p_sym + phi2
C = 2*pi*k*c/p_sym + phi3

# 1) Single hidden pre-activation
z = cos(A) + cos(B)

# 2) Square; reduce the powers and flatten products into single
# cosines (cos^2 x = (1+cos 2x)/2 and 2 cos x cos y = cos(x-y) +
# cos(x+y) -- the same high-school identities as in the prose).
z_sq = TR8(TR7(expand(z**2)))

# 3) Multiply by U_cj = (2/N) cos(C) and flatten again. A single TR8
# pass can leave products of cosines behind, so iterate expand + TR8
# until every term is a single cosine.
zU = z_sq * 2 * cos(C) / N_sym
for _ in range(4):
    zU = TR8(expand(zU))

# 4) Sort the expansion into the terms that still carry the random
# phases and the single one that does not.
terms = Add.make_args(zU)
fully_flat = all(len(t.atoms(cos)) == 1 for t in terms)
locked = [t for t in terms if not (t.has(phi1) or t.has(phi2))]
phased = [t for t in terms if t.has(phi1) or t.has(phi2)]
surviving = simplify(Add(*locked))

# The phase signature (n1, n2) of each remaining term: the integer
# coefficients of varphi^(1)_a and varphi^(2)_b inside its cosine
# argument. Every one must be nonzero for the construction to work.
def signature(t):
    cos_arg = list(t.atoms(cos))[0].args[0]
    return (int(cos_arg.coeff(phi1)), int(cos_arg.coeff(phi2)))

signatures = sorted({signature(t) for t in phased})
nonzero_ok = fully_flat and all(s != (0, 0) for s in signatures)

# Check the surviving term is exactly cos(2 pi k (a+b-c)/p) / N.
expected_surviving = cos(2*pi*k*(a + b - c)/p_sym) / N_sym
surviving_ok = simplify(surviving - expected_surviving) == 0

# 4b) Large-N limit of the phased part. Let g_j = N * (phased terms
# of neuron j) -- bounded and zero-mean in the random phases. The
# phased part of the logit is (1/N) sum_j g_j, so by the strong law
# of large numbers it converges almost surely to the per-neuron phase
# average E[g], which evaluates to 0 exactly (every cosine carries a
# nonzero integer multiple of varphi^(1)_a or varphi^(2)_b). The CLT
# fixes the approach rate: fluctuations of size O(N^{-1/2}).
g = N_sym * Add(*phased)
phase_avg = simplify(
    integrate(integrate(g, (phi1, 0, 2*pi)), (phi2, 0, 2*pi))
    / (2*pi)**2
)
phase_avg_ok = phase_avg == 0

# 5) Sum over N neurons, one frequency per neuron (k_j = j = 1..N).
# The phase-free term adds coherently; the phased part vanishes in the
# limit by 4b (and is kept, in full, in the numerics below).
f_c = Sum(expected_surviving, (k, 1, N_sym))

# At target (a + b - c == 0): every cosine is 1, so the sum gives N,
# divided by N gives 1.
f_c_at_target = simplify(f_c.subs({a: 0, b: 0, c: 0}).doit())

# 6) Closed form via the Dirichlet kernel identity.
#    sum_{k=1}^{N} cos(k x) = -1/2 + sin((2N+1) x/2) / (2 sin(x/2)).
# With x = 2 pi m / p (m = a + b - c) and divided by N:
m = symbols('m', real=True)  # generic so sin(pi*m) does not collapse yet
half = Rational(1, 2)
dirichlet_closed = (
    -half
    + sin((2*N_sym + 1) * pi * m / p_sym) / (2 * sin(pi * m / p_sym))
)
f_c_closed = dirichlet_closed / N_sym

# 7) Large-N limit, off-target (m an integer, m != 0 mod p).
# The closed form is -1/(2N) + [bounded numerator] / (2 N sin(pi*m/p)).
# Both pieces vanish: -1/(2N) -> 0 trivially, and the bounded
# oscillatory numerator divided by 2 N sin(pi*m/p) -> 0 by the squeeze
# theorem (|sin| <= 1, denominator grows linearly, and sin(pi*m/p) is
# a nonzero constant for m an integer not divisible by p).
# SymPy's Limit() returns AccumBounds for the oscillating term and
# leaves the expression unevaluated; the squeeze bound is the
# justification, and the numerical sequence below shows the decay.
limit_off_target = Integer(0)

# 8) Final boxed identity. In the N -> infinity limit:
#     - target  (a + b == c mod p): f_c -> 1 (every term is 1)
#     - off     (a + b != c mod p): f_c -> 0 (Dirichlet kernel / N -> 0)
# Both cases combined as a Kronecker delta:
final_expr = KroneckerDelta(c, a + b)  # equality understood mod p

# ---- Numerics: the FULL forward pass with real random phases. ----
# No expectation is taken anywhere and no term is dropped. Every
# neuron j gets fresh phases psi1_j, psi2_j ~ U[0, 2pi) and frequency
# k_j = j; the phased terms must die by the sum over neurons alone.
rng = np.random.default_rng(0)

def forward_all(p_val, N_val):
    j = np.arange(1, N_val + 1)
    psi1 = rng.uniform(0.0, 2.0*np.pi, N_val)
    psi2 = rng.uniform(0.0, 2.0*np.pi, N_val)
    grid = np.arange(p_val)
    Za = np.cos(2.0*np.pi*np.outer(grid, j)/p_val + psi1)
    Zb = np.cos(2.0*np.pi*np.outer(grid, j)/p_val + psi2)
    Z2 = (Za[:, None, :] + Zb[None, :, :])**2
    U = (2.0/N_val)*np.cos(
        2.0*np.pi*np.outer(grid, j)/p_val + (psi1 + psi2)
    )
    # logits f_c(a,b) for all pairs at once: (p*p, p)
    return Z2.reshape(p_val*p_val, N_val) @ U.T

# Convergence at p = 7, (a, b) = (2, 3): the target logit (c = 5) vs
# the largest |logit| among the six wrong classes.
p7, a_demo, b_demo = 7, 2, 3
tgt_c = (a_demo + b_demo) % p7
N_demo = [10, 100, 1000, 10000]
target_seq = []
off_seq = []
for N_v in N_demo:
    F = forward_all(p7, N_v).reshape(p7, p7, p7)
    logits = F[a_demo, b_demo]
    target_seq.append(round(float(logits[tgt_c]), 4))
    wrong = np.delete(logits, tgt_c)
    off_seq.append(round(float(np.abs(wrong).max()), 4))

# Per-class logits and argmax check at p = 7, N = 10000.
F7 = forward_all(p7, 10000).reshape(p7, p7, p7)
expected_c = (np.arange(p7)[:, None] + np.arange(p7)[None, :]) % p7
argmax_c = F7.argmax(axis=2)
all_match = bool((argmax_c == expected_c).all())
ex_logits = [round(float(v), 3) for v in F7[a_demo, b_demo]]
ex_a, ex_b = a_demo, b_demo

# Verify argmax for EVERY input pair across primes, full forward pass.
test_primes = [5, 7, 11, 13, 17, 19, 23, 29, 31]
N_verify = 4000
all_verified_at_N = True
for p_chk in test_primes:
    F = forward_all(p_chk, N_verify)
    am = F.argmax(axis=1).reshape(p_chk, p_chk)
    exp = (np.arange(p_chk)[:, None] + np.arange(p_chk)[None, :]) % p_chk
    if not bool((am == exp).all()):
        all_verified_at_N = False
        break

result = {
    "step_z": latex(z),
    "step_z2": latex(z_sq),
    "step_zU": latex(zU),
    "step_locked": latex(expected_surviving),
    "surviving_ok": bool(surviving_ok),
    "phased_count": len(phased),
    "phased_signatures": [list(s) for s in signatures],
    "signatures_nonzero": bool(nonzero_ok),
    "phased_mean": latex(phase_avg),
    "phased_mean_ok": bool(phase_avg_ok),
    "step_sum": latex(f_c),
    "step_closed_form": latex(f_c_closed),
    "limit_off_target": latex(limit_off_target),
    "f_c_at_target": latex(f_c_at_target),
    "final_identity": latex(final_expr),
    "N_demo": N_demo,
    "target_seq": target_seq,
    "off_seq": off_seq,
    "verified_primes": test_primes,
    "N_verify": N_verify,
    "all_verified_at_N": bool(all_verified_at_N),
    "p_value": p7,
    "example_a": ex_a,
    "example_b": ex_b,
    "example_correct_c": int(expected_c[ex_a, ex_b]),
    "example_argmax_c": int(argmax_c[ex_a, ex_b]),
    "example_logits": ex_logits,
    "all_match": all_match,
}
json.dumps(result)

First run downloads Pyodide (~10 MB) + SymPy. Cached afterwards.

The cosines in

z_{j}^{2} U_{c j}

contain linear combinations of

φ_{a}^{(1)}

and

φ_{b}^{(2)}

with integer coefficients. Exactly one combination is identically zero — the

(a + b - c)

term, by construction. The rest are zero-mean over the random phases, so by the law of large numbers their average over

N

neurons vanishes as

N \to \infty

, approached with

O (N^{- 1/2})

fluctuations: destructive interference by the sum alone. The numerical check runs the full forward pass — every term kept, phases drawn fresh — and confirms the argmax over

c

equals

(a + b) mod p

for every input pair.

Convergence as N grows

Next I would like you to get a feel for how quickly the analytic solution converges to 100% accuracy as a function of width $N$ (that is, parameter count $3 N p$ ). For each class $c \in {0, \dots, p - 1}$ we plot the network output logit $f_{c} (a, b)$ computed from the analytic solution. As $N$ grows, the target logit $(c = (a + b) mod p)$ grows toward its limit while the other logits average out to zero. Move the slider to increase the width and see how accuracy jumps. Observe that in log-log form the loss is a linear function of $N$ and has the form $lo g L = - lo g N + const$ . There is no entropy term because the dataset is fully deterministic.

target = (1 + 3) mod 5 = 4

Loss (MSE)

0.183

fit:

4.23 \cdot N^{- 0.827}

Accuracy1.000

N = 64256

Activations for general p

The same random-Fourier construction works for any modulus. Pick a $p$ and a width $N$ ; each tile below is one neuron's pre-activation $z_{j} (a, b) = cos (2 π k_{j} a / p + φ_{a}) + cos (2 π k_{j} b / p + φ_{b})$ drawn as a $p \times p$ heatmap over the input grid. The loss / accuracy plots show the same $N$ -convergence story, this time for the chosen $p$ .

Loss (MSE)

0.478

fit:

32.3 \cdot N^{- 0.887}

Accuracy1.000

p = 2331

N = 128256

Here are a few things to play with. Two sliders that increase $p$ up to $31$ and the width $N$ up to $256$ are available. As $p$ grows the frequency resolution gets finer because the frequencies are multiples of $2 π / p$ .

By re-rolling the neurons you should be able to see that measuring the exponent in the loss is quite tricky: finite- $N$ fluctuations change the value by up to $20%$ , even though we know the value analytically to be $- 1$ . You should also note that the intercept of the log-log loss plot depends on the data: as you increase $p$ the intercept grows.

A note on the math. The construction is exact for $p$ prime — every non-target logit converges to $0$ and the target to $1$ in expectation. For non-prime $p$ there's a small residual bias (the discrete Fourier basis on $Z / p Z$ has slightly different structure), but it sits well below the finite- $N$ fluctuations, so the picture stays the same.

Accuracy vs train fraction

In my opinion, the most interesting unresolved question about grokking (that may have some ramifications for AI more broadly) is the shape of sample complexity. Namely, how does accuracy at the end of training depend on the fraction of training data? Below is a simulation that lets you see explicitly that there is a very sharp transition in accuracy as a function of data fraction. Below the critical value accuracy is random, while above the critical value accuracy is $100%$ . The transition (most likely sharp in the limit of finite $p$ and large $N$ ) indicates a misalignment between training and test sets.

max stepsseed 7idle

test acc — converged (rounded to 1.0)test acc — didn't converge in 20,000 steps final test loss

The transition typically sits around train fraction $\approx 0.4 - 0.5$ and the exact number depends on the details of optimization. I found that momentum makes the biggest difference and can reduce the critical value. The fraction also depends on the modulus $p$ , but does not seem to depend on the width $N$ very much. The dependence on $p$ is interesting: as $p$ grows the total number of training points grows, but the training fraction decreases approximately as $1/ p$ . Thus, I do not expect the transition in the limit of large $p$ to be present (or, perhaps it will sit at zero training fraction).

Final thoughts

I tried to convince you that even in deep learning we can hope to have solvable models akin to the Ising or XYZ models in physics that can be understood completely, yet are useful to get a clear picture of how real-world phenomena play out. My personal hope is that interpretability research will evolve beyond descriptive explorations of alien botany and zoology towards quantitative science akin to biophysics. Analytically tractable models are one place where this transformation can begin.