The Geometry of Grokking: Norm Minimization on the Zero-Loss Manifold
This provides a theoretical explanation for a puzzling behavior in neural networks, which is incremental as it builds on prior links to weight decay.
The paper tackles the grokking phenomenon in neural networks, where generalization occurs long after memorization, by proving that gradient descent minimizes weight norm on the zero-loss manifold and deriving closed-form dynamics for a two-layer network, with experiments confirming the predicted gradients reproduce delayed generalization.
Grokking is a puzzling phenomenon in neural networks where full generalization occurs only after a substantial delay following the complete memorization of the training data. Previous research has linked this delayed generalization to representation learning driven by weight decay, but the precise underlying dynamics remain elusive. In this paper, we argue that post-memorization learning can be understood through the lens of constrained optimization: gradient descent effectively minimizes the weight norm on the zero-loss manifold. We formally prove this in the limit of infinitesimally small learning rates and weight decay coefficients. To further dissect this regime, we introduce an approximation that decouples the learning dynamics of a subset of parameters from the rest of the network. Applying this framework, we derive a closed-form expression for the post-memorization dynamics of the first layer in a two-layer network. Experiments confirm that simulating the training process using our predicted gradients reproduces both the delayed generalization and representation learning characteristic of grokking.