The Geometry of Grokking: Norm Minimization on the Zero-Loss Manifold

TL;DR

This paper explains the grokking phenomenon in neural networks as a process of norm minimization on the zero-loss manifold, providing theoretical proofs and a simplified model that reproduces delayed generalization.

Contribution

It introduces a formal framework linking grokking to constrained optimization and derives a closed-form expression for post-memorization dynamics, supported by experimental validation.

Findings

Gradient descent minimizes weight norm on the zero-loss manifold.

The derived model reproduces delayed generalization and representation learning.

Theoretical proofs are provided in the limit of small learning rates and weight decay.

Abstract

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.