Research map / Optimization / Implicit bias and generalization
A Theoretical Framework for Grokking: Interpolation followed by Riemannian Norm Minimisation
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Shows that gradient flow with small weight decay on a general smooth loss follows a two-phase trajectory: fast convergence to a critical manifold, then slow Riemannian descent of the l2 norm on that manifold, offered as a theoretical explanation of grokking.
Problem
Explaining why grokking occurs, i.e. why test loss plateaus long after training loss reaches zero and then improves, in terms of optimisation dynamics of gradient flow with weight decay.
Approach
Under smoothness, o-minimality, boundedness and a Morse-Bott assumption, the weight-decayed gradient flow as lambda->0 splits into a fast unregularised phase converging to the critical manifold and a slow phase at time order 1/lambda that follows the Riemannian gradient flow of the squared l2 norm on that manifold, with a self-contained proof avoiding heavy machinery from prior work.
Key results
- Theorem: as lambda->0, the weight-decayed trajectory converges uniformly to the unregularised flow on finite horizons and, after time order 1/lambda, follows a Riemannian gradient flow minimising the l2 norm on the critical manifold.
- On matrix completion (20x20, rank 3), the reconstruction converges to the low-rank ground truth as weight norm decreases during the slow phase, with grokking starting around t≈10^2 and the trajectory reaching the solution by t≈10^4.
- On a two-layer ReLU network with lambda=1e-3, training loss reaches near zero by t=1 and the weight norm then steadily decreases around t=1/lambda, with generalisation improving by t≈10^5.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| add small weight decay (λ = 10^{-3}) | gradient descent for low-rank matrix completion · synthetic matrix completion | test loss | Helped |
| add small weight decay (λ = 10^{-3}) | two-layer ReLU network · synthetic ReLU regression | test loss | Helped |
| allow weight decay to drive the slow-phase drift | gradient descent for low-rank matrix completion · synthetic matrix completion | reconstruction rank | Helped |
Surprising findings
- The grokking drop only appears sudden on a logarithmic time axis; in linear time it spans a duration comparable to the preceding plateau.
Limitations
- Analysis restricted to smooth losses; ReLU networks are not covered by the theory, only illustrated empirically.
- Requires global Morse-Bott assumption and bounded unregularised trajectories, excluding classification with separable data.
- Asymptotic analysis as lambda->0 only; fixed-lambda case is harder and only heuristically treated.
Experimental setup
Datasets: synthetic linear regression synthetic low-rank matrix completion (20x20, rank 3, 50% observed) synthetic 1D regression with 10 points for two-layer ReLU network
Metrics: training loss test loss weight norm singular values of reconstruction
Builds on: singularly perturbed systems theory Morse-Bott / Polyak-Lojasiewicz analysis Lyapunov stability on interpolation manifolds
Closest papers by meaning
- Grokking Beyond the Euclidean Norm of Model Parameters (ICML 2025 · similarity 0.90)
- Grokking at the Edge of Numerical Stability (ICLR 2025 · similarity 0.89)
- Grokking at the Edge of Linear Separability (ICML 2025 · similarity 0.88)
- Grokking as the transition from lazy to rich training dynamics (ICLR 2024 · similarity 0.87)
- Let Me Grok for You: Accelerating Grokking via Embedding Transfer from a Weaker Model (ICLR 2025 · similarity 0.86)
- Why Do You Grok? A Theoretical Analysis on Grokking Modular Addition (ICML 2024 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.