Grokking Math

Grokking stated precisely

Fix a small algorithmic task — modular addition (a + b) mod p is the canonical one — and train with a regularized objective. Grokking is the regime where the training and test curves decouple in time: training accuracy reaches 100% at some step t_fit, test accuracy stays near chance until a much later step t_gen, then rises sharply to near 100%. The interesting quantity is the gap t_gen - t_fit, which can span three or four orders of magnitude in optimizer steps.

The puzzle is that classical learning theory says nothing about when generalization arrives — it bounds the final gap, not its timing. So grokking is really a claim about optimization dynamics on a fixed landscape: the generalizing solution exists the whole time, and the optimizer takes a long, seemingly stalled path to reach it.

Weight decay is the engine, not a detail

The single most important experimental fact is that grokking is fragile without regularization. Remove weight decay and, on these tasks, the network typically memorizes and never groks within any practical budget; add it and the delayed transition appears reliably. Weight decay is not a minor hyperparameter here — it is the force that makes generalization eventually happen.

Recall what the regularizer does to the objective. With loss L(θ) and coefficient λ, the training target is L(θ) + (λ/2)·||θ||^2. The gradient step gains a term -ηλθ, a constant pull of every weight toward zero. Once L is essentially zero — the plateau after t_fit — the data gradient nearly vanishes and this shrinkage term dominates. The network is no longer learning to fit; it is being slowly compressed among all the ways it can already fit, and that compression is what eventually selects a generalizing solution.

The norm-growth versus generalization tradeoff

To see why compression helps, separate two things a solution has: its function (the input-output map) and its norm (the size of its weights). Memorization is cheap to find but expensive to store: to carve out an isolated correct answer for every training pair without any shared structure, a network leans on large, finely tuned weights — sharp bumps in function space cost large ||θ||.

A solution that has discovered the underlying rule reuses the same structure across all inputs, realizing the same labels with much smaller weights. This sets up a tension the regularizer resolves: early training drives the norm up, because the fastest way to cut training loss is to amplify whatever features separate the data (memorization), while weight decay pushes it down. The plateau is the period in which those pressures roughly balance and the optimizer searches, at a shrinking norm budget, for a function that still fits.

Two competing solutions

Varma et al. sharpen this into a picture of two circuits that coexist inside the network throughout training. The memorizing circuit fits the training set by rote; it does not generalize and needs a large weight norm to store all those independent facts. The generalizing circuit computes the actual algorithm; it generalizes perfectly and needs a smaller norm to produce confident outputs.

Crucially, both drive training loss to zero, so on the training set the data term cannot tell them apart. What separates them is the regularizer. Think of the state as a mixture of the two circuits, with the total correct-answer logit a sum of a memorizing and a generalizing contribution. Weight decay cares only about the norm cost of producing that logit — and it is that cost, not accuracy, that ultimately decides which circuit the optimizer keeps.

Effective complexity and circuit efficiency

The precise lever is efficiency: how much correct-answer logit a circuit produces per unit of weight norm, roughly efficiency = (logit on correct class) / ||θ||. Because the generalizing circuit shares structure across all inputs, it is the more efficient of the two — it buys the same confident output more cheaply in norm.

This reframes grokking as an implicit-bias argument. Among all functions with zero training loss, the regularized objective prefers the one of least norm — the lowest effective complexity. Early on the memorizing circuit wins because it is easier to find by gradient descent, but it is not the cheapest zero-loss solution; so once training loss is exhausted the norm penalty keeps reallocating weight toward the cheap generalizing circuit.

A worked example of the efficiency argument

Make it concrete with deliberately illustrative numbers (a toy calculation, not measured values). Suppose reaching a target logit of 10 on the correct class costs a memorizing circuit a weight norm of 100, while a generalizing circuit reaches the same logit of 10 at a norm of 25. Their efficiencies are 10/100 = 0.1 and 10/25 = 0.4.

circuit        logit   norm    logit/norm   penalty ~ norm^2
memorizing      10     100        0.10          10000
generalizing    10      25        0.40            625   ← ~16x cheaper

Both give identical training predictions, so the data loss is equal; the only term that differs is the penalty, which favors the generalizing circuit by a wide margin. The optimizer does not jump — it descends this penalty gradient slowly, which is why the payoff is delayed.

Why the transition is late and abrupt

Two features of the dynamics explain the shape. It is late because during the plateau the data gradient is tiny; only the small -ηλθ term is doing work, so the migration between circuits proceeds at a crawl set by ηλ. It is abrupt because test accuracy is a threshold read-out of a continuous quantity: as weight shifts to the generalizing circuit its logits grow smoothly, but test accuracy only flips once they overtake the memorizing ones on held-out inputs.

So the sharp jump in the accuracy curve is a nonlinearity of the metric, not a discontinuity in the weights: the plateau is a slow, monotone reallocation that accuracy only reports once it crosses a decision boundary.

Liu et al.: generalization on the constrained-norm manifold

Liu et al. give a complementary geometric account. Split a weight vector into its direction and magnitude, θ = r · u with ||u|| = 1, and note that generalization depends far more on the direction than the scale. Fix the norm r and the solutions of that norm form a sphere — a manifold on which some regions generalize and others memorize.

In this view grokking is slow travel along that manifold toward its generalizing region. Weight decay sets the radius the optimizer is pulled toward; the delay is the time spent drifting across the sphere from a memorizing initialization into a generalizing basin. Their ‘Omnigrok’ analysis shows the delay can be compressed or removed by initializing at a well-chosen norm — strong evidence that the norm is the controlling variable rather than an accident of the task.

The learning-rate and weight-decay dependence

Both pictures agree on a prediction: the delay is governed by the product ηλ, the rate at which the shrinkage term acts. Increase the weight-decay coefficient λ and the norm is driven down faster, so the generalizing circuit is selected sooner and t_gen moves earlier; increase the learning rate η (within stability) and the same migration happens in fewer steps.

Push either too far and grokking disappears: with very aggressive regularization the network may fail to fit, or generalize so promptly that there is no visible plateau. This is why grokking shows up as a band in a phase diagram over λ, data fraction, and model size: the regime where memorization is reached first but regularization is strong enough to eventually overturn it. The gap is a symptom of a particular balance, not a fixed property of the architecture.

Data fraction and the memorize-first regime

The amount of training data sets which circuit is easier to find first. With very little data, memorizing is trivially cheap and there is little to migrate to, so the network overfits; with abundant data, the generalizing circuit is the easiest way to fit even early, and the network generalizes immediately with no plateau. Grokking lives between these. As you sweep the data fraction upward, the delay t_gen - t_fit shrinks and collapses to zero — a clean confirmation that the delay is a competition between two reachable solutions, not a fixed feature of the task.

Implications for real training

Modular arithmetic is a laboratory, but the mechanism is general. Weight decay — the ubiquitous AdamW default — is usually justified as vague ‘regularization,’ but grokking shows it also acts as a selection pressure that, given enough steps, prefers the lower-complexity solution among many that fit. The unusually clean gap on toy tasks is likely smeared out in large models into the ordinary observation that test loss keeps improving after training loss has stalled.

For a small CPU-trained model the readings are concrete: do not equate a flat training-loss plateau with a finished model; treat the weight-decay coefficient as a real dial on generalization timing, not a throwaway constant; and remember a well-chosen initialization scale can buy much of the benefit without a long plateau.

What the theory does and does not settle

The efficiency and manifold accounts explain a great deal: why weight decay is necessary, why the transition is both delayed and sharp, why the delay depends on η, λ, and data fraction, and why a norm-aware initialization can erase it — a coherent story where the objective, not luck, decides the outcome.

What they do not fully pin down is the exact trajectory — the precise route the optimizer takes between the two circuits, and its sensitivity to architecture and optimizer choice, remain active research. And this formal view is deliberately silent on the content of the generalizing circuit: the structure the network settles into is the subject of the mechanistic companion. The theory here tells you a low-complexity solution will be selected and roughly when; it does not, by itself, tell you what that solution computes.