Research map / Architectures and efficiency / RNNs and normalization
Is Grokking a Computational Glass Relaxation?
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.
Frames grokking as computational glass relaxation by sampling the Boltzmann entropy landscape of one-layer transformers on modular arithmetic with Wang-Landau molecular dynamics, finds no entropy barrier, and proposes the WanD optimizer that eliminates grokking.
Problem
Grokking—delayed sudden generalization long after memorization—lacks a clear mechanistic explanation; prior theories (first-order phase transition, weight-norm Goldilocks zone) conflict on whether an entropy/free-energy barrier exists and what drives generalization.
Approach
Treats the network as a physical system with train loss as energy and uses Wang-Landau molecular dynamics to sample the Boltzmann entropy as a function of log training loss and smoothed test accuracy, interpreting memorization as a non-equilibrium glass and grokking as slow relaxation; a physics-inspired optimizer WanD built on the same dynamics avoids grokking.
Key results
- Entropy landscapes for the grokking tasks show no entropy barrier between memorization and generalization states, and the equilibrium curve is continuous, arguing against a first-order phase transition.
- WanD reaches generalization in about 2,000 steps, comparable to AdamW, at about 50 s per 1K steps versus about 72 s for AdamW on one RTX 4090.
- WanD trains to generalization at weight norm around 175 with grokking largely eliminated.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use WanD instead of AdamW | One-layer Transformer · Modular addition x+y mod 67 | running time per 1K steps: about 50 seconds vs roughly 72 seconds | Helped |
| Use WanD instead of AdamW | One-layer Transformer · Modular addition x+y mod 67 | optimization steps to generalization: ∼ 2,000 optimization steps for both methods | No clear effect |
| Use WanD instead of AdamW | One-layer Transformer · Modular addition x+y mod 67 | grokking | Helped |
| Use WanD instead of AdamW | One-layer Transformer · Modular addition x+y mod 67 | weight norm at generalization: about 175 | Helped |
| Constrain weight norm to 30 | AdamW-trained Transformers · Modular addition x+y mod 67 | grokking | Helped |
| Constrain weight norm to 30 | AdamW-trained Transformers · Modular addition x+y mod 67 | high-entropy advantage | Hurt |
| Increase the training dataset fraction | AdamW and WanD · Modular addition x+y mod 67 | grokking | Helped |
| Use WanD instead of AdamW | One-layer Transformer · Modular arithmetic tasks with different dataset fractions | training efficiency: nearly three times the efficiency of AdamW | Helped |
| Use WanD instead of AdamW | One-layer Transformer · Modular addition and x²+y mod 67, 50%–80% training fractions | training efficiency | Helped |
| Sample the entropy landscape using cross-entropy test loss instead of smoothed test accuracy (ablation) | WLMD entropy sampling · Modular arithmetic tasks | entropy barrier between memorization and generalization | No clear effect |
Surprising findings
- WanD generalizes at high weight norm (about 175), counter to the weight-norm Goldilocks-zone account of grokking.
- Under a fixed weight norm of 30 that suppresses grokking, a high-entropy advantage still persists, though weaker.
Limitations
- Experiments limited to one-layer transformers on modular arithmetic tasks with p=67.
- WanD is sensitive to hyperparameters (learning rate, temperature, friction) and unstable near generalization.
- No error bars; single fixed data split and random seed due to sampling cost.
Experimental setup
Datasets: modular arithmetic (x+y mod 67) modular arithmetic (x^2+y mod 67) modular arithmetic (x^3+xy^2+y mod 67)
Metrics: test accuracy training loss weight norm wall-clock time per 1K steps
Baselines: AdamW Entropy-SGD
Builds on: Wang-Landau molecular dynamics Entropy-SGD Yang et al. high-entropy advantage framework Nanda et al. grokking setup AdamW
Closest papers by meaning
- Grokking at the Edge of Numerical Stability (ICLR 2025 · similarity 0.89)
- Let Me Grok for You: Accelerating Grokking via Embedding Transfer from a Weaker Model (ICLR 2025 · similarity 0.88)
- Grokking as a First Order Phase Transition in Two Layer Networks (ICLR 2024 · similarity 0.88)
- Flatness is Necessary, Neural Collapse is Not: Rethinking Generalization via Grokking (NeurIPS 2025 · similarity 0.88)
- The Geometric Inductive Bias of Grokking: Bypassing Phase Transitions via Architectural Topology (NeurIPS 2026 · similarity 0.88)
- Grokking at the Edge of Linear Separability (ICML 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.