🤖 machine learning

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

This paper explains the delayed generalization phenomenon known as grokking by proving that, in the limit of small learning rates and weight decay, gradient descent minimizes weight norm on the zero-loss manifold, a mechanism validated through a derived closed-form expression for post-memorization dynamics and experimental simulations.

Original authors: Tiberiu Musat

Published 2026-06-02

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Tiberiu Musat

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Mystery: What is "Grokking"?

Imagine you are teaching a robot to do math. You show it one example: $1 + 1 = 2$ .

Phase 1 (Memorization): The robot quickly learns to say "2" whenever it sees "1 + 1". It has memorized the answer. If you ask it to solve a new problem, like $2 + 2$ , it fails. It's just a parrot repeating what it heard.
The Long Wait: You keep training the robot. For a long time, nothing seems to change. It still only knows the one example.
Phase 2 (Grokking): Suddenly, after hundreds or thousands of extra steps, the robot has a "lightbulb moment." It figures out the rule of addition. Now, it can solve $2+2$ , $5+3$ , and any other addition problem perfectly.

This strange delay—where the robot goes from "memorizing" to "understanding" long after it has already mastered the training data—is called Grokking.

The Paper's Main Idea: The "Zero-Loss" Manifold

The authors of this paper want to explain why this happens. They propose a new way of looking at how the robot learns during that long waiting period.

Think of the robot's brain as a giant, multi-dimensional landscape.

The Goal: The robot wants to reach a valley where the "error" (how wrong it is) is zero.
The Zero-Loss Valley: Once the robot memorizes the single example ( $1+1=2$ ), it has reached the bottom of a very specific, flat valley. In this valley, the robot makes zero mistakes on the training data.
The Problem: This valley is huge. There are millions of different ways to arrange the robot's internal knobs (weights) to get zero error. Some of these arrangements are "smart" (they generalize to new math), and some are "dumb" (they only work for the one example).

The Secret Mechanism: Norm Minimization

The paper argues that once the robot is in this "Zero-Loss Valley," the training process changes its goal. It stops trying to lower the error (because the error is already zero) and starts trying to simplify itself.

The Analogy: The Tightrope Walker
Imagine the robot is a tightrope walker on a very long, winding wire (the Zero-Loss Valley).

The Wire: The wire represents all the possible ways the robot can get the answer right.
The Push: The training process includes a "weight decay" (a small force that constantly pushes the robot to use less energy).
The Result: Because the robot is already on the wire (zero error), the only thing it can do is slide along the wire. The "energy-saving" force pushes it toward the part of the wire that is the shortest and simplest.

The authors prove mathematically that the robot is essentially sliding along this wire, looking for the simplest possible solution. Eventually, it finds the "simplest" path, which happens to be the one that understands the general rule of addition. This is why generalization happens after memorization: the robot has to finish its long, slow slide along the wire to find the simple solution.

The "Toy Model" Proof

To show this isn't just a fluke, the authors built a tiny, simple robot (a linear model) with just two knobs.

They trained it on $1+1=2$ .
Observation: The robot quickly found a solution that worked for the training data but was weird (e.g., using a huge positive number and a huge negative number that canceled out).
The Slide: Then, slowly, the "weight decay" pushed the knobs toward a much simpler, balanced solution (1 and 1).
The Result: As soon as the knobs reached that simple, balanced spot, the robot suddenly became great at solving any addition problem, not just $1+1$ .

The "Isolated" View: Focusing on the Embedding Layer

The paper also tackles a second question: Can we understand just part of the robot's brain without modeling the whole thing?

In the specific case of modular addition (math with a clock, like 11 + 2 = 1), previous research showed the robot arranges numbers in a circle.

The Analogy: Imagine the robot's first layer is a mapmaker. It takes numbers and places them on a map.
The Discovery: The authors created a mathematical shortcut. They showed that if you assume the second layer of the robot is always "perfectly adjusted" to the first layer, you can write a simple formula that predicts exactly how the mapmaker (the first layer) will move.
The Simulation: When they ran this formula on a computer, it perfectly recreated the "Grokking" effect. The mapmaker slowly rearranged the numbers from a messy blob into a perfect circle, and then the robot started understanding the math.

Summary of Findings

Grokking is Geometry: The delay in learning isn't a bug; it's a feature of geometry. The robot has to travel a long distance along a "zero-error" path to find the simplest solution.
Weight Decay is the Engine: The small force that tries to make the robot's numbers smaller (weight decay) is what pushes it along this path. Without it, the robot would get stuck in the "memorization" spot forever.
Simplification Works: You can predict how a complex network learns by looking at just one part of it (the embedding layer) and assuming the rest of the network instantly adapts to it.

What the Paper Does Not Claim

It does not claim this works for every type of AI or every kind of data (it focuses on specific math problems and simple networks).
It does not suggest using this to build better medical AI or self-driving cars yet.
It does not claim to have solved the mystery of all neural networks, only the specific "grokking" phenomenon in these settings.

In short, the paper tells us that Grokking is the robot taking a long, slow walk down a straight line to find the simplest, most elegant way to solve a problem.

Technical Summary: The Geometry of Grokking

Problem Statement
The paper addresses the phenomenon of "grokking" in neural networks, where a model achieves perfect generalization on test data only after a substantial delay following the complete memorization of the training set. While previous research has linked this delayed generalization to representation learning driven by weight decay, the precise underlying dynamics remain unclear. Specifically, the authors seek to answer: (1) What is the exact role of weight decay in post-memorization learning? and (2) Can the learning dynamics of specific network components (like embedding layers) be isolated from the rest of the network?

Methodology and Theoretical Framework
The authors approach the problem through the lens of constrained optimization and gradient flow, analyzing the system in the limit of infinitesimally small learning rates and weight decay coefficients ( $\lambda \to 0$ ).

Post-Memorization Dynamics: The paper formally proves that once a network achieves zero training loss (memorization), the subsequent learning dynamics are constrained to the zero-loss manifold ( $Z$ ). In this regime, gradient descent effectively minimizes the weight norm ( $\|\theta\|^2$ ) while remaining on $Z$ . The authors establish that the loss gradient becomes orthogonal to the tangent space of the zero-loss manifold as the parameters approach $Z$ , meaning the loss function no longer drives parameter updates; instead, weight decay drives the parameters toward the minimum-norm solution within the zero-loss set.
Isolated Component Dynamics: To analyze specific layers, the authors propose an approximation where the dynamics of a subset of parameters ( $\theta_1$ ) are decoupled from the rest ( $\theta_2$ ). By assuming $\theta_2$ adapts instantaneously to the optimal value for any given $\theta_1$ (effectively treating $\theta_2$ as a fast variable), they derive an effective cost function $R(\theta_1) = \min_{\theta_2} L_\lambda(\theta_1, \theta_2)$ .
Two-Layer Network Derivation: Applying this framework to a two-layer network with mean-squared error loss, the authors derive a closed-form expression for the learning dynamics of the first layer (the embedding layer). In the over-parameterized, zero-loss limit, the dynamics of the first layer weights $W_1$ are shown to follow the gradient flow of a specific cost function that balances weight decay against the fit of the second layer.

Key Results

Theoretical Proof: The paper proves that for almost every dataset, the zero-loss set is a smooth manifold free of singularities during standard training. It further proves that the loss gradient is orthogonal to this manifold near the solution, confirming that post-memorization learning is driven entirely by norm minimization.
Toy Model Validation: Using linear models on addition tasks, the authors demonstrate that models initially converge to a zero-loss solution (memorization) and then slowly drift along the zero-loss manifold toward the minimum-norm solution (generalization). The speed of this drift is inversely proportional to the weight decay coefficient.
Simulated Grokking: The authors simulate the derived closed-form dynamics for the first layer of a network trained on modular addition. The simulation successfully reproduces the characteristic delayed generalization: the model maintains zero training loss while test accuracy remains at random chance for hundreds of steps, eventually jumping to perfect generalization.
Representation Learning: The simulation reveals that during the generalization phase, the embedding vectors organize into circular structures in the Fourier domain. Specifically, the norms of Fourier features equalize, and the features become orthogonal, mirroring the "circular representations" observed in prior interpretability studies of grokking.

Significance and Claims
The paper claims to provide a formal geometric explanation for grokking, identifying it as a two-phase process: a fast phase of memorization followed by a slow phase of norm minimization on the zero-loss manifold. By deriving a closed-form expression for the dynamics of the embedding layer, the authors offer a theoretical basis for understanding how specific network components learn representations.

The authors position this work as a step toward demystifying the inner workings of neural networks, which is crucial for developing safer and more reliable AI systems. They note that their findings align with concurrent work by Boursier et al. (2026) regarding regularized gradient flow. The paper modestly acknowledges limitations, noting that it focuses on mean-squared error loss rather than cross-entropy and is currently restricted to two-layer networks, leaving the dynamics of deeper architectures as an open challenge.