Why Neural Networks Grok: A Spectral Theory of Delayed Generalization
Neural networks exhibit a peculiar behavior known as grokking: they fit the training data early on, then appear to plateau, and only much later suddenly generalize to unseen examples. The phenomenon has been documented in modular arithmetic tasks, but its mechanism has remained poorly understood. Is the network still learning during the apparent plateau? If so, what is it learning, and what triggers the delayed generalization? A new theoretical paper by Lenz Pracher, Pascal de Jong, and Oskar Lieshaus of LMU Munich, Alan Jeffares of Cambridge, and Steffen Rulands of Stanford provides a rigorous mathematical answer: after memorization, a finite residual remains that drives continued evolution of the network's tangent features, and the timescale of this process is controlled by a single simple quantity—the product of learning rate and weight decay.
The NTK Lens: Lazy Learning Versus Rich Learning
To understand the paper's approach, it helps to understand the Neural Tangent Kernel (NTK). The NTK quantifies how much a neural network's output changes when its parameters are perturbed. During training, if the NTK remains approximately constant, the network operates in the "lazy" regime: it is essentially a fixed feature extractor, and learning happens only by adjusting the output layer. If the NTK evolves substantially, the network is in the "rich" regime: the internal features themselves are being reshaped, which is what most people intuitively think of as learning.
The authors argue that grokking is precisely the transition between these two regimes. A network first trains in the lazy regime, memorizing the training labels through output-layer adjustments while the NTK remains nearly fixed. During this phase, a finite residual error persists—the network cannot perfectly fit the data with its current feature representation. That residual, it turns out, is the engine of everything that follows.
The Key Mechanism: Residuals Drive Feature Growth
The paper's central theoretical result, stated as Theorem 1, proves that for homogeneous networks trained with squared loss and L2 weight decay, a finite residual remains after memorization. More precisely, the fraction of each target component that remains as residual error is larger for NTK eigendirections with smaller eigenvalues. These are the directions that the current features represent poorly—the "weak target modes."
This is the crucial observation. The network has memorized the training data, but it has done so unevenly. Directions in output space where the current features are strong are fit well. Directions where the features are weak retain substantial error. And because the NTK is not truly fixed—it evolves in response to the parameters—those large residuals feed back into the NTK dynamics themselves.
The paper's Theorem 2 formalizes this feedback. When the neural tangent hierarchy is projected onto a task-aligned Fourier direction, the dynamics reduce to a two-dimensional system: the residual drives growth of the NTK strength along that direction, while weight decay simultaneously shrinks it. The equilibrium of this system determines how much of the target component is ultimately expressed in the network's output. For a direction with nonzero target projection, the equilibrium NTK strength is positive, meaning the residual-driven growth can sustain task-aligned feature learning indefinitely—as long as the coupling between residual and NTK growth is positive.
The Slow Clock: η·λ_W Controls the Grokking Timescale
Perhaps the most practically striking prediction is that the grokking timescale is controlled by a single quantity: the product of the learning rate η and the weight decay λ_W. The paper derives this through an adiabatic reduction of the two-mode dynamics. After the fast initial fit, the residual relaxes quickly to its ridge value, and the remaining evolution of NTK strength occurs on a slow timescale proportional to 1/(η·λ_W).
Concretely, if the network has fitted the training data after s_fit updates, the slow-time increment during grokking is Δτ ≈ η·λ_W·(s − s_fit). This means that changing the learning rate and weight decay while keeping their product constant produces essentially the same grokking behavior. The paper confirms this prediction empirically with a remarkable experimental grid: an 84-by-90 grid of trained homogeneous MLPs across varying learning rates and weight decays, showing that the transition-time contours collapse onto a single curve when parameterized by the product η·λ_W.
Phase Boundaries: When Grokking Fails
The theory also predicts the boundaries of the grokking phase in the (η, λ_W) plane. These boundaries arise from several constraints. First, the task-aligned NTK strength must be able to reach a threshold sufficient to move held-out predictions across their decision margins. If the weight decay is too large, the equilibrium NTK strength falls below this threshold, and generalization cannot occur at all. Near this critical decay, the grokking time diverges logarithmically—the network needs an exponentially longer training run to reach the required feature strength.
Second, the weight decay cannot be so large that the network fails to fit the training data in the first place. The paper derives an approximately vertical boundary in the (η, λ_W) plane beyond which the training accuracy criterion is never reached. Third, sufficiently large learning rates encounter a discrete-time stability edge analogous to the familiar edge-of-stability phenomenon in gradient descent.
Together, these constraints partition the optimizer space into regions where the network memorizes without generalizing, where it groks after a predictable delay, and where it either fails to fit or becomes unstable. The paper's Corollary 1 gives the explicit finite-time boundary formula.
Empirical Confirmation: Fourier Structure After Saturation
The paper tests its predictions on modular addition, a task where the target function has a natural Fourier structure. The network must learn to compute (a + b) mod p for a prime p, which requires learning specific Fourier modes. This task provides an ideal testbed because the relevant spectral directions are known in advance.
The experiments reveal that after training accuracy has saturated, the NTK continues to reorganize. The leading eigenvectors of the label-space NTK develop Fourier-like standing-wave profiles that sharpen over thousands of additional updates. This post-fit Fourier organization is precisely the task-aligned feature learning that the theory predicts. The empirical NTK is not static during grokking; it is actively being reshaped by the residual into a representation that better supports the task.
The Transformer experiments extend the findings beyond the idealized homogeneous setting. A one-block Transformer trained on modular addition shows the same macroscopic phase structure in a 42-by-45 grid, and the same (η·λ_W) scaling of the generalization transition time, despite violating the exact homogeneity assumption. This suggests that the mechanism is robust and not an artifact of the simplified network architecture.
Why This Matters
The paper connects several previously disparate observations into a unified mechanistic picture. Prior work had shown that learned representations and the empirical NTK can continue to evolve during grokking, and that leading NTK eigenfunctions become increasingly task-relevant. But what drives this evolution? The answer, according to this paper, is the residual itself: the error that remains after memorization feeds back into the NTK dynamics, gradually reshaping the network's tangent features in task-aligned directions.
This resolves a long-standing question in the grokking literature. Is the delay just fixed-feature dynamics playing out slowly, or is genuine feature learning occurring? The paper shows it is the latter: the NTK evolves, the residual drives that evolution, and the timescale is controlled by the optimizer hyperparameters in a specific, testable way. The spectral theory provides a quantitative framework that predicts not just when grokking occurs but where it can occur at all.
Limitations and Open Directions
The theoretical results are derived under specific assumptions: homogeneous networks, squared loss, gradient flow, and a one-mode spectral reduction. The paper acknowledges that the one-mode approximation holds when the task-aligned Fourier direction remains approximately an eigendirection of the NTK, with weak mixing into other directions. The modular arithmetic task is particularly well-suited to this approximation because its symmetry naturally isolates clean Fourier modes.
The paper also notes that the (η·λ_W) scaling, while observed empirically in both MLPs and Transformers, was derived in the adiabatic limit where residual relaxation is fast compared to mode motion. Whether the scaling holds more precisely in architectures where this separation of timescales is less clean remains an open question. The paper's appendices develop the full mathematical machinery—including the neural tangent hierarchy closure, the adiabatic reduction, the discrete-time stability analysis, and Jacobian-mediated feature dynamics beyond homogeneity—for readers who want to trace every derivation.
Finally, the paper's framework addresses the specific case of delayed generalization driven by post-fit feature learning. Other mechanisms for grokking—fixed-feature dynamics in linear models, representation competition, circuit-level changes—can produce similar phenomenology through different routes. The spectral theory provides one rigorous account, not the only one, and the authors are careful to situate their work among these alternatives.
The practical implications are worth noting. The (η·λ_W) scaling means that practitioners who observe grokking in their models can estimate the expected delay by tuning these hyperparameters. More fundamentally, the paper shows that weight decay is not merely a regularizer that prevents overfitting—it is an active driver of feature learning that shapes the network's tangent space long after the training labels have been memorized. The residual is not noise to be eliminated; it is the signal that guides the network toward better generalization.
Read the paper on arXiv