The Geometry of Grokking: Norm Minimization on the Zero-Loss Manifold (2511.01938v1)

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.

• The paper demonstrates that after perfect memorization, weight decay drives gradient descent toward the minimum-norm solution on the zero-loss manifold.
• It establishes a rigorous constrained optimization framework validated empirically using modular addition tasks and isolated dynamics of network components.
• Results reveal that delayed generalization (grokking) naturally emerges from norm minimization, offering insights for improved deep network training.

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

Introduction and Motivation

This paper provides a rigorous theoretical and empirical analysis of the grokking phenomenon in neural networks, where generalization emerges only after a significant delay following perfect memorization of the training data. The authors focus on the modular addition task, a canonical setting for grokking, and establish a geometric and optimization-based framework for understanding post-memorization learning dynamics. The central thesis is that, after memorization, gradient descent with weight decay drives the network to minimize the weight norm while remaining on the zero-loss manifold. This constrained optimization perspective is formalized and extended to the isolated dynamics of network components, particularly the embedding layer in two-layer architectures.

Grokking and Norm Minimization: Toy Model Intuition

The paper begins with illustrative toy models to build intuition for the main theoretical claims. In the simplest case—a two-parameter linear model trained on a single addition sample—the authors show that after rapid memorization, further learning is governed entirely by weight decay, which pushes the parameters toward the minimum-norm solution on the zero-loss set. This two-phase behavior is empirically demonstrated: initial loss minimization is followed by slow norm minimization, with generalization emerging only in the latter phase.

Figure 1: A two-parameter linear model $\hat{y} = w_1 x_1 + w_2 x_2$ exhibits grokking, with norm minimization driving generalization after memorization.

Extending to a three-parameter linear model, the authors visualize training trajectories from different initializations. All trajectories converge to the zero-loss plane and then move directly toward the minimum-norm solution, regardless of initialization. This supports the claim that post-memorization dynamics are governed by constrained norm minimization.

Figure 2: A three-parameter linear model $\hat{y} = w_1 x_1 + w_2 x_2 + w_3 x_3$ demonstrates that all trajectories converge to the minimum-norm solution on the zero-loss plane.

The authors further discuss the geometric structure of the zero-loss set, noting that it is generally a manifold but may contain singularities. They show that, except for a null set of singular points, the training trajectory remains on a smooth zero-loss manifold, even in the presence of non-smooth activations such as leaky ReLU.

Figure 3: Training trajectories in various architectures, illustrating curved zero-loss manifolds and the effect of singularities and non-smooth activations.

Theoretical Framework: Constrained Optimization on the Zero-Loss Manifold

The main theoretical contribution is a formal proof that, in the limit of infinitesimal learning rate and weight decay, gradient descent dynamics after memorization are equivalent to minimizing the weight norm subject to the constraint of zero training loss. The authors define the zero-loss set $\mathcal{Z}$ and show that, under overparameterization and perfect memorization, the trajectory remains arbitrarily close to $\mathcal{Z}$ for sufficiently small weight decay.

A key result is the orthogonality of the loss gradient to the tangent space of the zero-loss manifold. This implies that, once on $\mathcal{Z}$, the loss gradient does not induce movement along the manifold, and weight decay becomes the sole driver of learning. The authors provide a rigorous proof using Taylor expansion and properties of the Hessian and Jacobian of the network output.

Isolated Dynamics of Network Components

To analyze representation learning within specific network components, the authors introduce an effective theory that isolates the dynamics of a parameter subset (e.g., the embedding layer). By assuming that the remaining parameters are always optimal for the current value of the subset, they derive an approximate cost function for the isolated dynamics. This leads to a tractable optimization problem for the component of interest, enabling closed-form analysis in certain cases.

Application to Two-Layer Networks: Closed-Form Dynamics

Specializing to two-layer networks, the authors derive a closed-form expression for the post-memorization dynamics of the first layer (embedding matrix) under the isolated dynamics approximation. The second layer is solved optimally via ridge regression, yielding a cost function for the first layer that is differentiable and amenable to analysis. In the zero-loss limit, the cost reduces to a function of the embedding norm and a trace term involving the Moore-Penrose pseudo-inverse of the hidden activations.

The resulting gradient flow equation for the embedding matrix is:

$$\dot{W}_1 \approx X^{T}\Bigl(\bigl(A Y Y^{T} A H\bigr) \odot \sigma'\bigl(X W_{1}\bigr)\Bigr) - W_{1}$$

where $H = \sigma(X W_1)$, $A = (H H^T)^{-1}$, and $\odot$ denotes the Hadamard product.

Empirical Validation: Simulated Dynamics and Grokking Phenomena

The authors empirically validate their theoretical framework by simulating the isolated dynamics of the embedding matrix in a modular addition task. They show that, even with zero training loss throughout, generalization on the test set emerges only after a substantial delay, reproducing the grokking phenomenon. Furthermore, the learned embeddings exhibit circular structure and orthogonality in their Fourier features, consistent with mechanistic interpretability results in prior work.

Figure 4: Simulated dynamics according to the derived gradient flow reproduce delayed generalization and circular representation learning in modular addition.

Implications, Limitations, and Future Directions

The paper provides a rigorous geometric and optimization-based explanation for grokking, linking delayed generalization to norm minimization on the zero-loss manifold. This framework clarifies the role of weight decay and enables effective theories for isolated network components. The results have practical implications for designing training regimes and understanding representation learning in overparameterized models.

Limitations include the focus on mean-squared error loss and two-layer architectures; extension to cross-entropy loss and deeper networks remains an open challenge. The theoretical machinery developed here may inform future work on the dynamics of more complex models and tasks, as well as the development of more efficient and interpretable learning algorithms.

Conclusion

This work establishes that, in the grokking regime, neural network learning dynamics are governed by norm minimization constrained to the zero-loss manifold. The authors provide both theoretical proofs and empirical validation, demonstrating that this perspective explains delayed generalization and representation learning in modular addition. The framework for isolated dynamics offers a powerful tool for analyzing specific network components, with potential applications to broader settings in deep learning.