Research map / Architectures and efficiency / RNNs and normalization
Grokking as a First Order Phase Transition in Two Layer Networks
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.
Applies the adaptive kernel approach with a Gaussian-mixture extension to two teacher-student models (cubic polynomial and modular addition) to show Grokking corresponds to a first-order phase transition with a mixed GMFL-I phase.
Problem
Explaining Grokking and feature learning in two-layer networks within a unifying analytical framework rather than through handcrafted toy solutions or perceptron models.
Approach
Extends the adaptive kernel approach with a variational Gaussian mixture approximation so that posterior weight distributions of a two-layer network under Langevin training become saddle-point mixtures; Grokking then appears as a first-order phase transition between a GFL phase and a mixed Gaussian-mixture phase.
Key results
- Theory predicts a first-order phase transition at effective interaction u_c approx 30.2 for the cubic teacher with epsilon=-0.3; experiments with a 200-network ensemble show a sharp transition near sigma^2=0.18.
- For modular addition, the mixed phase GMFL-I begins near sigma^2 approx 0.227 and GMFL-II near sigma^2 approx 0.175, with test RMSE dropping from -1 to -0.7 (about 30% reduction) across the mixed phase.
- The trained cubic-teacher network learns the cubic component with n=3000 samples, versus an estimated 0.5e+6 needed by the NNGP limit.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Reduce noise strength to around σ²=0.18 | Erf student network · Cubic teacher-student model | cubic teacher component learned | Helped |
| Increase α=β under the paper's scaling | Erf student network · Cubic teacher-student model | feature learning of the cubic component | Helped |
| Use finite-DNN feature learning rather than the FCN NNGP prior | Cubic teacher-student model · Cubic teacher-student model | samples to learn cubic target component: n=0.5e+6 -> n=3000 | Helped |
| Decrease noise strength σ² through the modular-algebra transition | Modular addition network · Modular addition | test discrepancy (a): -1 to -0.7; 30% reduction in the magnitude of the discrepancy | Helped |
Surprising findings
- The finite-width network learns the cubic teacher component with about two orders of magnitude fewer samples (n=3000) than the NNGP limit requires (about 0.5 million).
Limitations
- Analysis relies on the equivalent kernel approximation, which washes out generalization phenomena associated with Grokking.
- Scaling limits and saddle-point/Gaussian-mixture approximations are only exact in the limit; finite-size effects appear near transitions.
- Sample complexity advantage is conjectured and not rigorously established.
Experimental setup
Datasets: synthetic Gaussian teacher-student data modular addition two-hot dataset
Metrics: negative log posterior discrepancy of network output with teacher components test RMSE
Baselines: GP limit (NNGP)
Builds on: adaptive kernel approach equivalent kernel approximation Gaussian feature learning
Closest papers by meaning
- Grokking as the transition from lazy to rich training dynamics (ICLR 2024 · similarity 0.92)
- Deep Networks Always Grok and Here is Why (ICML 2024 · similarity 0.89)
- Dichotomy of Early and Late Phase Implicit Biases Can Provably Induce Grokking (ICLR 2024 · similarity 0.88)
- Three Mechanisms of Feature Learning in a Linear Network (ICLR 2025 · similarity 0.88)
- Is Grokking a Computational Glass Relaxation? (NeurIPS 2025 · similarity 0.88)
- Let Me Grok for You: Accelerating Grokking via Embedding Transfer from a Weaker Model (ICLR 2025 · similarity 0.87)
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.