A Spectral Theory of Grokking: Weight Decay induces Feature Learning

Lenz Pracher; Pascal de Jong; Oskar Lieshaus; Alan Jeffares; Steffen Rulands

cs.LGSep 22, 2026

A Spectral Theory of Grokking: Weight Decay induces Feature Learning

Authors: Lenz Pracher, Pascal de Jong, Oskar Lieshaus, Alan Jeffares, Steffen Rulands

Abstract

In grokking an early fit to the training data separates from a much later improvement in generalization. During this delay, training can move from a fixed neural tangent kernel (NTK) regime to one in which task-relevant kernel eigendirections continue to evolve. We provide a quantitative theory for how this transition from lazy to rich learning can produce delayed generalization. For homogeneous networks trained with squared loss and $L_2$ weight decay, we show that a finite residual remains after memorization, with larger residual fractions in target components associated with smaller NTK eigenvalues. These residuals feed back into the dynamics of the NTK itself, and projecting the resulting dynamics onto task-relevant spectral directions yields a reduced system in which residual-driven kernel growth competes with weight decay. This system predicts that the grokking timescale is controlled by the product of learning rate and weight decay, that feature learning slows logarithmically near a critical decay above which task-aligned NTK structure can no longer support generalization, and that stronger decay can prevent fitting altogether. We test these predictions in modular addition. In a homogeneous MLP, task-aligned Fourier structure continues to emerge in the NTK after training accuracy has saturated, and an 84 $\times$ 90-grid of trained networks across varying learning rate and weight decay recovers the predicted phase geometry and inverse-product scaling of the generalization time with learning rate and weight decay. A one-block Transformer shows similar macroscopic phase structure in a 42 $\times$ 45-grid, as well as the same transition-time scaling despite violating exact homogeneity. Together, these results provide a mechanistic derivation connecting post-fit feature learning to both the onset of generalization and its phase structure in the learning rate and weight decay plane.

Explore similar work

Jul 27, 2026cs.AI

Grokking on the Weight-Decay Clock: A Rate Hierarchy from Softly Broken Symmetries

Delayed generalization, or grokking, remains poorly understood despite extensive empirical study. We identify an exactly solvable late-time relaxation mechanism for grokking in linear models trained with full-batch heavy-ball optimization and weight decay, together with a locally quadratic extension to nonlinear neural networks. Our analysis reveals a distinguished population-active component of the empirical null space, which we call the grokking subspace. Along this subspace, the training predictions remain unchanged, leaving weight decay as the sole restoring force and giving rise to a slow dissipative relaxation governed by an exact discrete-time and continuous-time law. We show that only this subspace contributes to the slow asymptotic decay of the population risk and derive explicit iteration-scale predictions for the grokking time, recovering the familiar

(1-β)/(ηλ)

scaling in the weak-regularization regime. The theory further predicts distinct effects of optimizer choice, distinguishing coupled

L_2

regularization from decoupled weight decay, and yields causal predictions for interventions that modify the grokking component. We verify all theoretical identities without fitted parameters in a synthetic model where every subspace and relaxation rate is computable in closed form. We further observe genuine delayed generalization in modular addition, where the measured delay follows the predicted scaling and the late-time relaxation agrees closely with the theoretical clock.

Taeyoung Kim

Grokking Parallel Liquid-Inspired Temporal Relaxation Regularization Techniques Clocks Symmetry Hierarchical

Ask PDF

Sep 12, 2026cs.AI

Quantifying the Memorization-to-Generalization Transition: Scaling Laws and Phase Structure in Grokking

Neural networks trained past memorization frequently undergo a delayed transition to generalization, a phenomenon known as grokking. Despite theoretical progress on \emph{why} this transition occurs, the quantitative structure of \emph{when} it occurs in hyperparameter space remains uncharacterized. We map the memorization-to-generalization boundary across 384 configurations of two-hidden-layer MLPs on modular arithmetic, fitting a power-law scaling relation for generalization onset time:

T_{\mathrm{grok}} \propto H^{-0.27}\, D^{-2.04}\, \eta^{-0.50}\, \lambda^{-0.64}

(

R^2 = 0.732

;

0.821

with interactions). The exponent hierarchy reveals that data complexity (

D^{-2.04}

) is the dominant driver of regime transition, not model capacity (

H^{-0.27}

): doubling data accelerates generalization by

{\sim}4\times

, while doubling width yields only

{\sim}1.2\times

. A sharp phase boundary at weight decay

\lambda \gtrsim 1.0

separates grokking from non-grokking configurations, and weight norm trajectories show monotonic compression during the transition, consistent with implicit regularization selecting low-complexity solutions. These results provide a quantitative foundation for predicting and controlling regime transitions in overparameterized networks.

Anish Kataria

Grokking Over-Parameterized Neural Networks Two-Layer Neural Networks Memorization Scaling Laws Structural Generalization+2

Ask PDF

Sep 7, 2026cs.LG

A Theoretical Analysis of Generalization Dynamics in Neural Networks under Gradient Descent with Weight Decay

Understanding generalization remains a central challenge in machine learning because it requires jointly considering data, architecture, and training dynamics. In this paper, we develop a theoretical framework that characterizes how these factors jointly shape generalization performance throughout training. More precisely, we study a broad class of neural networks trained under the

\ell^2

loss by gradient descent (GD) with weight decay, and prove the convergence of GD to a neighbourhood of the global minimizers of the empirical loss. By partitioning the space based on the input data, we then decompose the population error into data error, optimization error, and prediction variation error, and bound them separately. In particular, for the prediction variation error, which measures the oscillations of the learned function, we propose (local) approximate homogeneity and derive explicit cellwise and layerwise bounds for its evolution along the training trajectory. These bounds yield two important implications: a necessary condition of improved generalization explains differences in layerwise generalization behavior; a sufficient condition describes delayed generalization and provides a theoretical characterization of grokking.

Yuqing Wang, Ioannis G. Kevrekidis, Mikhail Belkin

Generalization Bounds Two-Layer Neural Networks Gradient Descent Neural Network

Ask PDF