A Spectral Theory of Grokking: Weight Decay induces Feature Learning

TL;DR: Grokking — the strange delay between fitting the training set and suddenly generalizing — gets a mechanistic, quantitative explanation here: weight decay prevents a network from ever fully fitting the training data, and the small residual error it leaves behind is concentrated in exactly the directions the network's current features represent worst. That residual acts as a persistent teaching signal that slowly rotates the neural tangent kernel toward task-relevant structure, on a timescale set by 1/(\eta\lambda_W) — the inverse product of learning rate and weight decay. The theory predicts a full phase diagram in the (\eta, \lambda_W) plane, which the authors verify with a 7,560-run MLP sweep, and they back the causal story with a clean intervention: artificially aligning the kernel early makes grokking arrive 1.5× sooner.

Grokking as a clean laboratory for post-fit learning

In normal training, everything happens at once: the network memorizes, builds representations, and improves on held-out data simultaneously, so you can't tell which process drives which. Grokking (Power et al., 2022) pulls these apart in time. On tasks like modular addition — learn (a,b) \mapsto (a+b) \bmod p — a network hits 100% training accuracy early, then sits at chance-level test accuracy for thousands of updates, then abruptly generalizes. The delay is a gift: it lets you ask what training is doing during the plateau, when the loss says the job is done.

The natural lens is the neural tangent kernel (NTK). If \mathbf{f} is the vector of network outputs on the training set and J = \nabla_{\boldsymbol\theta}\mathbf{f} is the Jacobian with respect to parameters, then K = JJ^\top measures how easily gradient updates can move the outputs in each direction of function space. A frozen K means "lazy" training — the network is doing kernel regression with fixed features. A changing K means feature learning. Prior empirical work (Nanda et al., Lin, Sanguino Bautiste et al.) showed the NTK's eigenstructure keeps reorganizing toward Fourier-like task features during the grokking plateau. But why? What force keeps pushing on the features after the training labels are already predicted correctly? And what sets the length of the delay? This paper answers both, and the answer is surprisingly tidy.

Weight decay never lets you finish

The setup: a D-homogeneous network (e.g., a bias-free ReLU MLP, which is degree-2 homogeneous: scaling all weights by c scales outputs by c^D), trained with squared loss plus L_2 weight decay \frac{\lambda_W}{2}\|\boldsymbol\theta\|^2. Homogeneity buys one crucial identity (Euler's): J\boldsymbol\theta = D\mathbf{f}, which lets the parameter-space decay term be rewritten purely in function space. The residual \boldsymbol\sigma = \mathbf{f} - \mathbf{y} then obeys, for a fixed kernel,

where \lambda is the (normalized) decay coefficient. This is ridge regression dynamics, and it converges not to zero error but to a persistent residual. In the eigenbasis of K (eigenvalues \Lambda_j, target coefficients y_j), the equilibrium residual in each direction is

Read that fraction carefully: directions with large \Lambda_j — directions the current tangent features express easily — get fit almost perfectly. Directions with small \Lambda_j — directions the features barely span — retain nearly their entire target component as error. Weight decay effectively hands gradient descent a to-do list, and the list is sorted by exactly what the current features are worst at.

The feedback loop: residuals grow the kernel

Here's the aha. In a finite network, K is not fixed — its evolution is governed by the neural tangent hierarchy (Huang & Yau), and the residual appears directly in \dot K:

where K^{(2)} is a third-order tangent tensor contracted against the residual. So the very error that weight decay refuses to erase is the driving force reshaping the kernel. Weight decay plays a double game: it shrinks K isotropically (-2\lambda K), but by keeping \boldsymbol\sigma alive it also keeps the anisotropic, task-directed growth term switched on.

Modular addition supplies the perfect coordinate system for reducing this: the target decomposes exactly into Fourier characters \chi_k(a)\chi_k(b) of \mathbb{Z}_p, matched in frequency. Projecting everything onto one normalized Fourier direction \mathbf{u} — defining \Lambda = \mathbf{u}^\top K\mathbf{u} (kernel strength along the task mode) and \sigma = \mathbf{u}^\top\boldsymbol\sigma — collapses training to a two-variable system:

with y = \mathbf{u}^\top\mathbf{y} > 0 the target coefficient and a > 0 a local coupling constant (an assumption, verified qualitatively by experiment, not a consequence of the theory). After the fast fit, \sigma < 0, so -a\sigma > 0: the residual pumps kernel strength into the task direction; growing \Lambda in turn lets the network express more of the target (f = \frac{\Lambda}{\Lambda + 2\lambda}\,y), shrinking the residual. This negative feedback has a unique stable equilibrium at

which is positive precisely for target-bearing directions (y > 0) and zero for unsupported ones — the reduced dynamics performs spectral selection: it grows exactly the modes the task needs and lets the rest decay.

Test accuracy jumps sharply despite \Lambda(t) growing smoothly for a mundane reason: each held-out example flips its predicted class when its margin crosses zero, and if these per-example thresholds are narrowly distributed in \Lambda, a smooth internal coordinate produces a step in accuracy. No discontinuity in the learning dynamics is required — a nice general lesson for "emergent capability" debates.

A slow clock and a phase diagram

Because residual relaxation (rate \sim \Lambda) is much faster than kernel motion (rate \sim \lambda), an adiabatic reduction gives \dot\Lambda \propto \lambda \cdot (\text{drive}). Converting gradient-flow time to optimizer updates s, the natural slow coordinate is

so the number of updates before grokking scales as (\eta\lambda_W)^{-1}. Integrating the slow equation from initial strength \Lambda_0 to the generalization threshold \Lambda_g within a training budget T carves the (\eta, \lambda_W) plane into predicted phases:

• Memorization when \eta\lambda_W is too small: the kernel doesn't reach \Lambda_g in T updates. Boundary: \eta_* \propto 1/\lambda_W.
• A critical decay \lambda_W^{\mathrm{mode}}: since \Lambda_* = -\lambda + \sqrt{\lambda^2 + ay} shrinks with decay, above this cutoff the equilibrium falls below \Lambda_g and generalization becomes unreachable — with the crossing time diverging logarithmically as you approach the cutoff, so the boundary curls upward.
• A fitting cutoff at even larger \lambda_W: the expressed gain \Lambda/(\Lambda + 2\lambda) drops so far that even training accuracy fails.
• Instability at large \eta: the discrete-time edge-of-stability condition \eta < 2/(\mathcal{N}\Lambda_{\max} + 2\lambda_W).

The evidence

Three lines, in increasing order of persuasiveness.

Post-fit NTK organization. On addition mod 97 with a bias-free one-hidden-layer ReLU MLP, the authors project the empirical NTK through true-label logits and average by output label, giving a 97×97 label-space kernel. Its periodic structure sharpens between updates 3,000 and 20,000 — long after training accuracy saturates — and its leading eigenvectors develop Fourier standing-wave profiles matching the task's character basis (Figure 2 of the paper). This is the predicted task-aligned kernel growth, seen directly.

Phase geometry and timing. An 84×90 grid — 7,560 independently trained MLPs on addition mod 23, 24,000 updates each — reproduces the predicted plane: broad memorization region, grokking band, upward-curling boundary near a decay cutoff, near-vertical fitting cutoff, and a large-\eta instability edge. The empirical memorization–grokking boundary has log–log slope −1.17 against the theory's calibrated −0.90; the weak-decay asymptote fits \eta = A/\lambda_W cleanly. Within the grokking region, the first update at which held-out accuracy reaches 10% scales linearly with (\eta\lambda_W)^{-1}, with the predicted slope not fitted to the data. A one-block Transformer (42×45 grid, 1,890 runs, 60,000 updates) — which violates the homogeneity assumption via attention, LayerNorm, biases, and residual paths — shows the same macroscopic phase structure and the same (\eta\lambda_W)^{-1} transition-time scaling (median \eta\lambda_W s_{80} = 0.92 across the grokking cohort).

The causal intervention (Appendix F — read this). The theory predicts \partial s_g / \partial \Lambda_0 < 0: start with more task-aligned kernel and grokking should arrive earlier, even if you stop helping before generalization begins. The authors train 40 paired students (shared initialization and data within each pair) on mod-97 addition; the intervention member gets an auxiliary penalty aligning its label-space NTK to a grokked teacher's — but only for the first 3,000 updates, then it's removed and both members train identically. The mean first 50%-accuracy crossing moves from 20,200 updates (baseline) to 13,528 (intervention), a 1.49× paired speedup. Since the objectives are identical from update 3,000 onward, only early tangent geometry can explain the shift.

An important honesty check the paper itself makes: (\eta\lambda_W)^{-1} timing alone proves nothing about feature learning — Xu et al. (2026) derive the same scaling for pure ridge regression with fixed features, and Kim (2026) for linear models. What distinguishes this account is the combination: the timing, the phase geometry including the critical-decay upturn, the directly observed post-fit Fourier organization of the NTK, and the intervention.

What to make of it

If the picture holds, it reframes weight decay: not merely a norm-based regularizer that eventually favors "simpler" solutions, but an active generator of feature learning. It manufactures a persistent supervision signal whose spectral profile is inversely matched to the current features' competence, and the tangent hierarchy converts that signal into kernel growth precisely where it's needed. Grokking's delay is then just the ratio of two timescales — fast ridge relaxation versus \lambda-paced kernel motion — and the practical corollary is unambiguous: the product \eta\lambda_W is the knob, up to a critical decay beyond which generalization becomes impossible rather than merely slow. The intervention result also hints at something actionable: seeding tangent geometry (e.g., from a teacher) buys real wall-clock generalization speedups.

Reasons for skepticism, in rough order of severity. The theory requires squared loss, coupled L_2 decay, full-batch vanilla SGD, and exact homogeneity — whereas the canonical grokking phenomenon lives in Transformers trained with AdamW and cross-entropy, where adaptive-optimizer effects and logit-scaling dynamics (Thilak et al., Prieto et al.) are known to matter and the relevant slow variable may differ. The Transformer evidence here is only macroscopic: same phase shape and timing scaling, but with an MSE-on-softmax loss chosen to stay close to the theory, and no microscopic verification of the residual→kernel mechanism. The one-mode reduction assumes the task Fourier direction stays an approximate NTK eigendirection with positive coupling a — assumed, not derived. And the fitted phase boundary is a genuine but weaker test than it first appears: \Lambda_0, \Lambda_g, and ay are inferred from the boundary itself rather than measured from NTK trajectories, so Figure 3a validates the functional form of the crossing time, not the absolute magnitude of kernel growth. Finally, modular arithmetic hands the theory its spectral basis for free; for general tasks there is no obvious analogue of the Fourier direction, and interacting task modes are explicitly left to future work.

Where to spend your time: Section 2.2–2.3 for the reduced system and phase-boundary derivation (the intellectual core, and short), then Appendix F for the intervention — the strongest evidence in the paper, oddly relegated to the appendix. Appendix E is worth a skim too: it shows the residual-driven \dot K term is exact for any differentiable model, isolating homogeneity as the only ingredient that doesn't generalize — a concrete roadmap for extending this to realistic architectures.