Research map / Architectures and efficiency / RNNs and normalization
Let Me Grok for You: Accelerating Grokking via Embedding Transfer from a Weaker Model
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.
Proposes GrokTransfer, which trains a small weak model to partial generalization and transfers its learned input embedding (via a low-rank A·B factorization) to a larger target model, eliminating grokking delay on XOR theory and algorithmic tasks.
Problem
Grokking (delayed generalization after memorization) makes neural network training unpredictable and inefficient; the goal is to make models generalize without delay.
Approach
Data embedding determines whether generalization is delayed; a small weak model trained only to non-trivial test accuracy yields a task-informative embedding that, factorized as A·B and used to initialize the target model's embedding, lets the target generalize almost immediately.
Key results
- On the XOR task, after one gradient step the target model with transferred embedding classifies all training points and achieves test error at most exp(-Omega(log^2 n)), whereas normal training needs many epochs.
- On the Transformer target with FNN weak model, total wall-clock time was 71079 ms for GrokTransfer versus 392667 ms from scratch (about 5x faster).
- On modular addition with a 2-layer Transformer target, the time gap was 46 for GrokTransfer versus 1119 for GrokFast.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Transfer a weak FNN embedding to the target FNN | Target FNN · Modular addition | generalization delay | Helped |
| Transfer a weak FNN embedding to the target FNN | Target FNN · Modular multiplication | generalization delay | Helped |
| Transfer a weak FNN embedding to the target FNN | Target FNN · (40, 3)-parity | generalization delay | Helped |
| Transfer a weak FNN embedding to the target Transformer | 8-layer Transformer · Modular addition | generalization delay | Helped |
| Transfer a weak FNN embedding instead of using GrokFast | 2-layer Transformer · Modular addition | Time Gap: 1119 -> 46 | Helped |
| Use GrokTransfer rather than scratch training | 8-layer target Transformer · Target training computation | wall-clock time (ms): 392667 -> 71079 | Helped |
| Train the target model for one step after transferring the embedding | Target model · XOR cluster data | test error | Helped |
| Use task-specific Fourier embeddings instead of general one-hot embeddings | FNN · Modular addition, p=113 | generalization delay | Helped |
| Use binary embeddings instead of general one-hot embeddings | FNN · Modular addition, p=113 | generalization delay | Helped |
| Use binary embeddings instead of general one-hot embeddings | FNN · Modular multiplication | generalization delay | No clear effect |
| Use embeddings extracted at different weak-model training epochs (ablation) | GrokTransfer target FNN · Modular addition | target test performance | Mixed |
| Train the weak model until it has grokked to nontrivial test performance (ablation) | GrokTransfer target FNN · Modular addition | generalization delay | Helped |
Surprising findings
- Grokking persists even after carefully tuning the initialization scale, suggesting the kernel regime is not the sole cause of grokking.
- An FNN-learned embedding transfers to a Transformer target and still removes most generalization delay.
Limitations
- Theory covers only a simplified two-layer XOR cluster task with fixed first layer and one gradient step.
- Method is only studied on problems where grokking already occurs; no broader tasks.
- Training instabilities and oscillations at large learning rates appear in Transformer experiments.
Experimental setup
Datasets: modular addition (p=113) modular multiplication (p=113) (40,3)-parity XOR cluster data (synthetic)
Benchmarks: modular addition modular multiplication (40,3)-parity XOR cluster data
Metrics: test accuracy time gap between 95% train and 95% test accuracy wall-clock time
Baselines: training from scratch GrokFast Goldilocks sphere-constrained training one-hot embedding binary embedding Fourier embedding GPT embedding
Builds on: grokking GrokFast Goldilocks zone / weight norm sphere empirical NTK weak-to-strong generalization
Closest papers by meaning
- Deep Networks Always Grok and Here is Why (ICML 2024 · similarity 0.92)
- Grokking at the Edge of Numerical Stability (ICLR 2025 · similarity 0.91)
- Grokking as the transition from lazy to rich training dynamics (ICLR 2024 · similarity 0.91)
- Grokking in Linear Estimators -- A Solvable Model that Groks without Understanding (ICLR 2024 · similarity 0.91)
- Grokking Beyond the Euclidean Norm of Model Parameters (ICML 2025 · similarity 0.90)
- To Grok or not to Grok: Disentangling Generalization and Memorization on Corrupted Algorithmic Datasets (ICLR 2024 · similarity 0.90)
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.