Grokking with 3000× fewer parameters

The standard story of grokking is thus: a small neural network trained on addition modulo 97 memorizes its training pairs \(a\) and \(b\), and then, many epochs later, is induced by weight decay to abandon the strategy of memorizing individual pairs in favor of learning a rule which generalizes to all, including unseen ones. Nanda et al. (2023) showed that the generalizing solution it learns is an instantiation of a fairly simple trigonometric solution. If you embed each number such that it sits at an angle on a circle, then a candidate answer \(c\) can be scored by \[\cos(\theta_a + \theta_b - \theta_c),\] which is largest when \(c\)’s angle is the sum of \(a\)’s and \(b\)’s. This equation minimizes the model’s L2 weight norm, so weight decay selects for it, and it works equally well for seen and unseen examples, so it generalizes.

The story of grokking, as traditionally told, stops here. The network has generalized, and anyone who examines its internal representations using interpretability techniques can find the trigonometric features that solve modular addition. However, there are two desiderata for an ideal learner which remain elusive:

In this essay, we show that it is possible to reduce the grokked network to a simple 18-parameter net which directly expresses the known solution, while retaining full generalization. The 18 numbers are the multiplication table of the complex numbers, and nothing about 97 is in them. The reduction requires exploiting a simple fact about the grokked network: grokking finds symbolic features which have extremely strong explanatory power — but in order to directly unveil them, we need to re-express the problem in the natural basis of those features rather than retaining the encoding that a model naïvely starts with. We read the found features off the grokked network’s own behavior rather than introducing external domain knowledge, and show that this allows for the compact solution to be identified by a secondary simplicity prior similar to but distinct from weight decay.

I argue that this suggests models trained to convergence are, in some sense, smarter than they know — they just are not typically given a way of knowing.

Performing addition with respect to a certain modulus can be expressed as rotation on a circle, since adding “wraps around” back to 0 the same way that tracing the circumference of a circle does. Under this framing, number \(t\) sits at angle \(\theta_t\) on a circle of radius one, at the point \(x_t = (\cos\theta_t, \sin\theta_t)\) of the plane; we call that point the number’s position. Adding \(b\) is turning through \(b\)’s angle: the answer to \(a + b\) sits at angle \(\theta_a + \theta_b\). A candidate answer \(c\) is scored by how close its own position is to that one, as the cosine of the angle between them, \(\cos(\theta_a + \theta_b - \theta_c)\): 1 when \(c\) sits exactly where \(a + b\) does, and falling off the farther away it sits.

This is the circle of the introduction, and the grokked network holds it in its weights: as Nanda et al. found, the first layer embeds each number as its position on such a circle, with the cosine and the sine of the number’s angle as the two coordinates. We did not supply the circle, and we did not take it from the weights either. We read it from the network’s answers. With the second input held fixed, the network sends each first input to one answer, and the 97 answers carry the numbers around a single loop; a loop has natural coordinates, and they come in cosine and sine pairs (the next section says how they are found). Each number \(t\) receives its position \(x_t\) this way, with \(\theta_t = 2\pi \cdot 18t/97\) up to where the circle starts. The recovered positions agree with this formula to five parts in ten thousand and lie on a unit circle to within 0.05%. The numbers do not appear in counting order around the circle. Going around once meets 0, 27, 54, 81, 11, …, and each step from \(t\) to \(t + 1\) is a jump of 18/97 of a turn. The grokked network uses thirteen such circles at once, at different frequencies, and any of the 48 possible windings of the 97 numbers around a circle would serve the rule; this is the one that survived the procedure described below.

To score candidates this way, a machine has to turn \(a\)’s position through \(b\)’s angle. Angles are how we have described this, but a network has no operation that adds angles. What it has are the two coordinates of each position, and in coordinates, turning \(a\)’s position through \(b\)’s angle is the angle-addition formula, \[\begin{aligned}\cos(\theta_a + \theta_b) &= \cos\theta_a\cos\theta_b - \sin\theta_a\sin\theta_b \\ \sin(\theta_a + \theta_b) &= \sin\theta_a\cos\theta_b + \cos\theta_a\sin\theta_b,\end{aligned}\] four real products, each of a coordinate of \(a\) with a coordinate of \(b\). This is the rule for multiplying two complex numbers. A position on the unit circle is a complex number of size one, and the product of two such numbers sits at the sum of their angles, which is the content of Euler’s formula \(e^{i\theta_a}e^{i\theta_b} = e^{i(\theta_a+\theta_b)}\). So a machine that multiplies positions adds angles with no further step, and the only remaining question is how cheaply it can multiply. In coordinates the product takes the four multiplications above, and a rearrangement going back to Gauss does it with three. Each of Gauss’s products multiplies a weighted sum of \(a\)’s two coordinates by a weighted sum of \(b\)’s two coordinates, and the two coordinates of the result are weighted sums of the three products. Three multiplying units can carry out that rearrangement, one product each, and three multiplying units are the network. Each unit holds three pairs of numbers: \(U\) weights \(a\)’s two coordinates, \(V\) weights \(b\)’s, the unit multiplies the two weighted sums, and \(W\) writes the product into the plane of the circle. The three units’ contributions add to one output point, and a candidate answer \(c\) scores by the projection of that point onto its own position: \[s(c \mid a, b) \;=\; \sum_{i=1}^{3} \,(U_i \cdot x_a)\,(V_i \cdot x_b)\,(W_i \cdot x_c).\]

That table is the entire network, and it is the multiplication table of the complex numbers, written as three products. Each row is one product: \(U\) and \(V\) weight the two sums it multiplies, and \(W\) carries the product to the output’s two coordinates. The procedure that produced it had no target number of units, and it stopped at three, the smallest number of products that can compute the rule exactly.

Written out in full, the program has three parts. The first is the position table, 97 rows and two columns holding the cosine and the sine of each number’s angle. Everything specific to this task lives in the table: that there are 97 numbers, and that each step of one is 18/97 of a turn. The second part is the 18 numbers, which say how two positions combine, and they hold nothing about 97, as the next section shows. The third is the readout, which picks the row of the table nearest the combined position, and it has no parameters of its own. We count the 18 and not the 194 because the network built the table itself. It is the recovered circle shown above, read from the network’s own answers, so we treat it as something the network knows and factor it out of the count. The 3,000× of the title counts the 18 against the original 53,985 on those terms. Counting the table as well makes 212, still 250× fewer. The network answers all 9,409 pairs correctly.

It adds by turning. A pair of weights applied to a position on the circle gives a sinusoid of that position’s angle, \(U \cdot x_a = |U| \cos(\theta_a - \alpha)\), with the pair setting the amplitude and the phase, so the three units compute nine sinusoids and multiply them in threes. Together they write an output point whose direction is \(a\)’s position turned through \(b\)’s angle, which is the position of \(a + b\). The network’s answer is the candidate whose own position lies closest to that point. On every one of the 9,409 pairs the point lands within 0.5° of the right number’s position, and the numbers on either side of it around the circle are 3.7° away.

Every score it can produce lies on one cosine. Across 9,409 pairs and 97 candidates there are 912,673 scores, and all of them fall on one curve, with a relative residual of 0.6%: \[\sum_{i=1}^{3} \,(U_i \cdot x_a)\,(V_i \cdot x_b)\,(W_i \cdot x_c) \;\approx\; 146.2\,\cos(\theta_a + \theta_b - \theta_c - \theta_0),\] where \(\theta_0\) is the angle at which the recovered circle happens to place the number 0. The left side is the network. The right side is the rule, and the rule adds: the angle inside the cosine is the angle from \(a + b\)’s position to \(c\)’s, so the score is highest for the candidate sitting where \(a + b\) sits, which is \(a + b\) itself. The runner-up is the candidate one step away around the circle, where the cosine has fallen by 0.2% of its peak, a margin of 0.3 at an amplitude of 146.2. Each unit’s product of three sinusoids expands into four cosine terms of equal size, one in \(\theta_a + \theta_b - \theta_c\) and three in other combinations of the three angles; across the three units the three unwanted terms cancel while the wanted one adds up, and the 0.6% is what the cancellation leaves.

The three units sum to a single table of eight numbers, one for each way of choosing a coordinate of \(a\), a coordinate of \(b\) and a coordinate of the output, and that table is a rule for multiplying two points of the plane. The rule the network found is the multiplication of complex numbers. Changing the order of the factors does not change the product, nor does the way three factors are grouped, and there is a point that leaves every other point unchanged when multiplied by it, an identity. That identity points at the number 0 on the circle, to within 0.15°, and the point that undoes multiplication by \(t\) sits where \(-t\) does, to within 0.6°. So the algebra’s identity is the group’s zero and its reciprocals are the group’s negatives. The 97 positions are then the 97th roots of unity of this product, the points that come back to the identity after 97 multiplications, and adding two numbers modulo 97 is multiplying the corresponding roots. That is the textbook way of placing the integers modulo 97 inside the complex numbers, and the network arrived at it without anyone having mentioned complex numbers.

Nothing in the 18 numbers refers to 97 or to the winding of 18. Both live in the position table. To see this, keep the 18 numbers exactly as they are and swap in a different table, placing the numbers 0 to \(n-1\) around the circle by steps of \(1/n\) of a turn from the same starting angle. The unchanged network then adds modulo \(n\) with every pair correct, for every \(n\) from 2 to 347. Hand it two arbitrary angles with no modulus at all, and the output arrow lands within 0.52° of their sum. A wrong answer needs the arrow to miss by more than half the gap between neighbouring positions, and at \(n = 348\) half the gap first drops below 0.52°, which is exactly where the first wrong answer appears. The 18 numbers are a rule for composing rotations. Which rotations, and how many of them, is the table’s business.

Nor do the 18 numbers hold any fitted constants. A rule for multiplying points of the plane can be written as a sum of products in many ways, as 12 is 3 times 4 and also 2 times 6, and the three units are one way of writing the complex product. Gauss’s version is another, with entries of 0, 1 and −1. A third has three units of equal strength whose angles step evenly around the circle. All three give the same eight-number table. Within that table two things are free: an overall size, which does not change which candidate wins, and a starting angle, which is pinned to wherever the recovered circle happened to place 0. Everything else is fixed by the requirement that every pair come out right. The particular 18 numbers are one of many spellings of the one table the task allows.

There is a one-line way to say all of this. The position table is the Fourier transform of a number, taken at one frequency: the number \(t\), as a spike at position \(t\), has transform \(e^{2\pi i\cdot 18t/97}\), which, up to a fixed turn, is its row of the table. The 18 numbers multiply two such transforms, and the readout is the inverse transform evaluated at the candidate. Multiplying two transforms convolves the two originals, and convolving two spikes gives one spike at the sum of their positions. The network is the convolution theorem at one frequency and nothing else. Nanda et al. described the grokked network as a transform, a multiplication and an inverse transform, and the squeeze has reduced it to those three steps. The same three steps, with matrices in place of complex numbers, are how networks have been found to compose elements of non-commutative groups as well (Chughtai, Chan and Nanda, 2023), so the object here is the smallest case of a general construction.

The number three has a reason of the same kind. With complex-valued weights, two multiplying units would do, since in the right complex coordinates the product of two complex numbers separates into two independent products. With real weights it takes three, and the third unit is the price of doing complex arithmetic with real numbers. A different two-dimensional number system, the split-complex numbers, whose unit elements compose hyperbolic rotations rather than circular ones, splits into two real products and would have squeezed to two units. The count of three is the signature of a circle.

We start with a typical grokking setup: we take an MLP with two hidden ReLU layers and 53,985 parameters, which receives \(a\) and \(b\) as one-hot vectors and returns 97 scores. We train on 2,822 of the 9,409 pairs; the model fits the training set by epoch 300, and reaches 100% on the remaining test set at epoch 18,550. It runs the rule on thirteen circles at once, one per frequency, each carried by several neurons. At this point, we have a grokked model as traditionally thought of: a model of the same dimensions as it started, but which went from memorizing its training set to generalizing to the test data over the course of training, under weight norm compactness pressure from weight decay.

From here, we apply two distinct moves to reduce the model from 54k to 18 parameters:

Move 1: rotate the features into the network’s own coordinates. Fix the second input at some number \(g\) and ask the network for the answer to every \(a\). The answers permute the 97 numbers, sending each \(a\) to \(a + g\), and because 97 is prime the permutation is a single loop through all of them. The natural coordinates of a loop are the sinusoids that go around it a whole number of times. The eigenvectors of the permutation are these sinusoids, in cosine and sine pairs, one pair for each of the 48 possible frequencies, and each pair places the 97 numbers evenly around a circle. The network’s own answers therefore define the circles. Nothing about the task is supplied: not the modulus, not the order of the numbers, and not that cosines are involved.

The network’s first layer is then rewritten to read each number’s positions on these circles in place of its one-hot vector, and its output layer to write a point in the plane of each circle, with a candidate’s score given by how well those points line up with the candidate’s own positions. This is a rotation of the input space. The information is the same and the coordinates differ. Keeping the thirteen circles the network already uses reproduces its answers, and the rewritten network begins the next move at 100% held-out accuracy.

Move 2: squeeze the weights out. The rewritten network is then asked to keep the answers it gives on its 2,822 training pairs while giving up as many weights as it can. The procedure is iterative pruning: remove the smallest weights, retrain briefly, and accept the removal only if every training answer is unchanged, repeating until no removal is accepted. Held-out accuracy is recorded throughout and never consulted.

The squeeze is a simplicity prior, as weight decay is, and the two measure simplicity differently. Think of the network as a program whose length is its number of nonzero weights. Weight decay prefers programs with small constants: it penalizes the sum of squared weights and leaves their number alone. The squeeze prefers short programs: it penalizes the number of weights and leaves their size alone, subject to the program still giving every training answer. It finds the shortest program, and the length of a program depends on the language it is written in.

What the squeeze leaves behind depends on the coordinates. In the original one-hot basis it produces a memorizer. The network keeps every training answer down to 7,085 weights while its held-out accuracy falls from 100% to 8%. In the rotated basis the same squeeze reaches 230 parameters at 99.8% held-out accuracy and, once whole circles may be removed at a time, 76 parameters on two circles at 99.94%. At 76 the surviving wiring is the rule’s block structure: each circle has its own first-layer units, its own second-layer units and its own output, and the two circles share no connection.

The reason is that the same feature has a different length in each language. With one-hot inputs, a unit responds to one particular number with a single weight, while a sinusoidal response takes 97 weights, one per number. A lookup table is short in this language and the rule is long, so the squeeze keeps the single weights and the sinusoids erode. On the circles a sinusoidal response takes two weights, one per coordinate, and singling out one number is the long thing. The same prior selects a lookup table in one basis and the rule in the other, because in each it selects the shortest program. The rotation changes the language, and that is what makes the squeeze work.

From 76 to 18. The grokked network’s own units are ReLUs, and a ReLU can only approximate a product. The rule needs exact ones, because the angle-addition identity reaches \(\theta_a + \theta_b\) only through products of terms in \(\theta_a\) and terms in \(\theta_b\). On a single circle the right answer beats its neighbor by 0.2% of the peak score, too thin a margin for approximate products, which is why the grokked network runs the rule on thirteen circles and why its own weights stop squeezing at two. A second network whose units multiply, built on the same circles, has no such limit. The same squeeze takes it from 7,845 parameters to three units on one circle with no biases: the 18 numbers above.

The usual story of grokking ends when the general solution wins. At that point the network computes the rule, an interpretability effort can find the trigonometric features inside it, and weight decay has finished its work: the solution it settles on is 53,985 weights in which no angle appears. We think this ending is premature. The network has found the solution without being able to state it or to run it cheaply, which is to say it has the function and not the formula. These two moves allow for discovery of the formula itself. The first supplies coordinates in which the formula is short, and the second writes it down. Neither is another round of learning modular addition. Both act on what the network already knows.

The squeeze is a minimum description length step: it finds the shortest program in the language of the features the network discovered. It counts terms, and a count of terms is only defined relative to a vocabulary. Here the vocabulary is the thirteen circles, each a cosine-and-sine pair of eigenvectors of the permutation \(a \mapsto a + g\) that the network itself computes when its second input is held at \(g\). The constraint is the network’s own answers on its 2,822 training pairs. Nothing comes from outside except the instruction to use fewer terms and the definition of a term. Weight decay is the kind of prior a learner can be given before it knows anything, and that is what makes it blind afterward. Once the network knows something, units fixed before learning began are the wrong units to measure it in.

A step like this can only be taken by a network with something to stand behind, and the grokked network has two such things. It has a body of answers it will not give up, the 2,822 training answers. And its behavior is coherent enough to define a language: with one input fixed, it carries the 97 numbers around a single loop, and the loop has eigenvectors. The memorizer at epoch 300 has the first and not the second. Its training answers are already at 100%, but most of its answers on the other 6,587 pairs are wrong, so fixing one input does not carry the numbers around a loop, and there are no circles to read. The vocabulary the squeeze counts in does not exist until the function generalizes. The two pressures are in sequence rather than in competition. Weight decay finds a function that has a short expression, regardless of what that expression is qualitatively. The squeeze finds the expression in its own terms. The first makes the second possible, and the second makes the first legible. The separation also shows afterward in what each part is good for: the 18 numbers, having nothing of 97 in them, run unchanged on addition modulo any other number up to 347 once that number’s positions are supplied.

What was learned once as a function is available thereafter as a vocabulary.

We present this as a proof of concept for abstraction creation: forming new terms from a learner’s own activity, and re-expressing what the learner knows in them. The terms here are the circles. They were read from the network’s behavior alone, without labels or any knowledge of the arithmetic, by an operation applied from outside and suited to this task: take the eigenvectors of the network’s own permutation. They were minted for one program. The case for a learner that mints such terms for itself is that its next program is written in the same terms, so that what was learned once as a function is available thereafter as a vocabulary. That learner is the aim of the wider research program.