A Spectral Theory of Grokking: Weight Decay induces Feature Learning

A mechanistic derivation linking post-fit feature learning to delayed generalization and phase structure in the learning rate-weight decay plane

Lenz Pracher; Pascal de Jong; Oskar Lieshaus; Alan Jeffares; Steffen Rulands· 2026· DOI 10.48550/arXiv.2609.26679

The core problem

Grokking refers to a phenomenon in neural network training where an initial fit to the training data is followed by a much later improvement in generalization. During this delay, the network transitions from a lazy learning regime, characterized by a fixed neural tangent kernel (NTK), to a rich learning regime where task-relevant kernel eigendirections continue to evolve. This paper aims to provide a quantitative theory for how this transition from lazy to rich learning can produce delayed generalization. The authors focus on homogeneous networks trained with squared loss and L_{2} weight decay, and they derive a reduced dynamical system that captures the interplay between residual-driven kernel growth and weight decay. The theory predicts specific scaling laws and phase boundaries that are then tested in modular addition tasks using both MLPs and Transformers.

The theoretical analysis yields several predictions. First, the grokking timescale is controlled by the product of learning rate (η) and weight decay (λ). Specifically, the generalization time t_{g} scales inversely with η λ : t_{g} \propto \frac{1}{η λ} . Second, feature learning slows logarithmically near a critical decay λ_{c}, above which task-aligned NTK structure can no longer support generalization. Third, stronger decay can prevent fitting altogether. These predictions are tested in modular addition. In the homogeneous MLP, task-aligned Fourier structure continues to emerge in the NTK after training accuracy has saturated. The 84 \times 90 grid of trained networks recovers the predicted phase geometry and the inverse-product scaling of the generalization time with learning rate and weight decay. The one-block Transformer shows similar macroscopic phase structure in the 42 \times 45 grid, as well as the same transition-time scaling, despite violating exact homogeneity. These results confirm the key predictions of the theory.

Grokking refers to a phenomenon in neural network training where an initial fit to the training data is followed by a much later improvement in generalization. During this delay, the network transitions from a lazy learning regime, characterized by a fixed neural tangent kernel (NTK), to a rich learning regime where task-relevant kernel eigendirections continue to evolve. This paper aims to provide a quantitative theory for how this transition from lazy to rich learning can produce delayed generalization. The authors focus on homogeneous networks trained with squared loss and L_{2} weight decay, and they derive a reduced dynamical system that captures the interplay between residual-driven kernel growth and weight decay. The theory predicts specific scaling laws and phase boundaries that are then tested in modular addition tasks using both MLPs and Transformers.

The authors consider homogeneous networks trained with squared loss and weight decay. They show that after memorization, a finite residual remains, with larger residual fractions in target components associated with smaller NTK eigenvalues. These residuals feed back into the dynamics of the NTK itself. By projecting the resulting dynamics onto task-relevant spectral directions, they obtain a reduced system in which residual-driven kernel growth competes with weight decay. The reduced system is analyzed to derive predictions about the grokking timescale and phase structure. The theoretical predictions are then tested in modular addition tasks. They train a homogeneous MLP and a one-block Transformer across grids of learning rates and weight decay values. For the MLP, an grid is used, and for the Transformer, a grid. They monitor the emergence of task-aligned Fourier structure in the NTK and measure generalization times. The experimental setup allows them to verify the predicted phase geometry and scaling laws.

The 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. The theory explains why grokking occurs: after memorization, residuals in directions with small NTK eigenvalues drive further kernel evolution, which competes with weight decay. The product η λ sets the timescale for this competition, leading to the inverse scaling of generalization time. The logarithmic slowdown near critical decay indicates a phase transition where task-aligned structure cannot be maintained. The fact that a Transformer, which is not exactly homogeneous, exhibits similar behavior suggests the theory may have broader applicability. The authors suggest that this framework could be extended to other architectures and tasks. The findings also have practical implications for hyperparameter tuning: to avoid grokking, one might increase weight decay or learning rate, but too much decay can prevent fitting. The paper opens avenues for understanding delayed generalization in more complex settings.

CS practitioners and researchers

Opening member content…