Beyond Embedding Transfer: Component Roles in Grokking Transfer and Stability
cs.LG
Quick verdict
Pre-registered factorial study: grokking-transfer speed comes from internal blocks, stability from frozen carriers; strong stats but toy-scale, n=1 mechanism evidence.
Summary
The paper dissects warm-start transfer in grokking on modular arithmetic (p=113), decomposing transferred parameters into token embeddings (E), readout head (U), and internal attention/MLP blocks (B), each Frobenius-norm-matched to cold init. A pre-specified 108-run factorial battery over 12 paired seed blocks shows Full (E+U+B) beats representation-only (E+U) by +5.46 pp early accuracy and 558 steps faster confirmation (Holm-corrected exact permutation tests), while U+B is TOST-equivalent to Full within a +/-500-step margin though Full wins 11/12 seeds. The core contribution dissociates acquisition from stability: continued training triggers severe post-grokking relapse; freezing E+U nearly eliminates offline relapse (19.40% to 0.07%), and online validation-triggered freezing cuts max drawdown from 22.06% to 0.60%, prospectively replicated on affine, quadratic, and two-layer targets (partial protection at two layers, 65.23% residual drawdown). Boundaries are mapped honestly: no confirmed benefit on non-abelian S5, and donor embeddings become essential in two-layer models (U+B deficit -4475.6 units). Limitations: toy-scale 1-2 layer transformers (227k-425k params), mechanistic localization resting on one probed trajectory plus four test-conditioned relapse events, relapse studied under a single hyperparameter regime, and no public code URL stated.
Review complete
Assessment at a glance
5
risk flags identified
● high 0
● medium 3
● low 2
5 strengths · 5 weaknesses · 5 author questions
Levels reflect verification complexity — not a quality score. Evolve never rates your paper.
Key claims under review
- Transferring internal attention/MLP weight matrices (B) in addition to E+U provides statistically significant incremental acquisition gains (+5.46 pp early accuracy, Holm p=0.0039; 558-step faster confirmation, Holm p=0.0088), replicating in 2-layer models (+704.67 integral units, 12/12 seeds).
- Transfer acceleration and retention stability are governed by distinct components: freezing the E+U carriers nearly eliminates post-grokking relapse (19.40% to 0.07%), while online validation-triggered freezing cuts True Max Drawdown from 22.06% to 0.60%, prospectively confirmed on affine, quadratic, and (partially) two-layer targets.
- U+B transfer is equivalent to Full within a pre-specified +/-500-step mean-latency margin in one-layer models, though Full is faster in 11/12 seeds, and donor embeddings become essential at two layers (U+B deficit -4475.6 units, far outside the +/-250 margin).
- On a probed trajectory, post-reach relapse is functionally localized to internal blocks (state swap restores 99.85% accuracy) with substantial linear decodability preserved at troughs (readout refitting recovers 8.37-11.03 pp), though this diagnostic scope is explicitly bounded.
- Carrier-shielding protection has empirical boundaries: no confirmed benefit on non-abelian S5 conjugation (+0.15+/-1.14 pp, p=0.895) and only partial attenuation at two layers (residual drawdown 65.23%).
Strengths
Unusually rigorous, pre-specified statistical protocol
Exact two-sided sign-flip permutation tests (2^12=4096 enumerations) with Holm-Bonferroni control over pre-specified 8-contrast families, BCa CIs, TOST equivalence with pre-registered margins, and RMST under administrative censoring. The paper explicitly refuses to over-interpret the censored B-alone condition (10/12 censored, N_non-zero=2, Table 1 Panel B). Internal consistency checks out: the reported TOST t-statistics (3.975/-4.249) are exactly reproducible from the listed 12 paired differences.
Prospective replication cohorts for every major claim
Online gating is confirmed on fresh blocks per claim (Table 2: verification 250-261, affine 270-281, quadratic 290-301, 2-layer 310-321; N=12 each, 12/12 or 11/12 concordant wins), and a scale-mismatch control (embed_orig_scale, RMST 18,025 vs 4,220 for scale-matched E) empirically validates the Frobenius-matching protocol (Sec 3.2).
Honest reporting of negative and boundary results
Genuine conceptual advance over prior embedding-transfer work
Deep reproducibility detail
Weaknesses
Toy-scale scope with unproven generality
Mechanistic localization rests on n=1 trajectory and test-conditioned event selection
Relapse prevalence is regime-specific and inconsistent across cohorts
Equivalence claim has a generous margin and an unresolved tension
No verifiable public artifact link
Full Review Report
$1.99 · one-timeUnlock the complete preliminary analysis Evolve produced for this paper.
- All 5 strengths & 5 weaknesses, analyzed in detail
- 5 risk flags with severity, section & reviewer questions
- Methods & statistical evidence assessment
- 5 questions a peer reviewer would ask
- Full six-dimension evaluation report & recommendations
One-time purchase · this paper only
One-time purchase, this paper only. AI-generated preliminary analysis — not peer review.
Frequently asked questions
What is grokking and why study transfer with it?
Grokking is a phenomenon where small neural networks trained on algorithmic tasks (like modular arithmetic) with strong weight decay first memorize the training set and only much later - after thousands of extra steps - suddenly generalize to held-out data. It is a clean, reproducible testbed for asking what computation a trained model has learned and which parameters carry it, which is why this paper uses it to study which model components transfer between related tasks.
Which model components should be transferred to speed up grokking?
According to this paper, transferring token embeddings and the readout head (the representation carriers) captures most of the benefit, but also transferring the internal attention/MLP weights adds a significant further gain: about 5.5 percentage points higher early accuracy and roughly 560 fewer steps to stable generalization. Surprisingly, the readout head alone is a stronger carrier than the embeddings, and in two-layer models the donor embeddings become essential - omitting them collapses performance.
What is post-grokking relapse and how is it prevented?
After a warm-started model reaches high test accuracy, continued training can sharply degrade it again - the paper observes average accuracy dips up to ~22-80% depending on the setting. The main fix is parameter shielding: freezing the transferred token embeddings and readout head, either from the start or dynamically once validation accuracy crosses 95%. This nearly eliminates the collapse in one-layer models, but only partially protects two-layer models and shows no confirmed benefit on the non-abelian group S5. Simple early stopping also avoids relapse but fixes the model at the trigger point, whereas shielding allows continued optimization.
Is the paper's code and data publicly available?
The paper states that run artifacts include deterministic seeds, checkpoints, and reproduction scripts, and its appendix documents the model code (src/model.py) in detail, but no public repository URL is given in the manuscript text. Availability of the artifacts would need to be confirmed with the author.
About this review
This is an AI-generated preliminary review by Evolve.Pub, produced from the openly available text of the paper. It is not peer review and does not represent the authors' views. The canonical record of the paper is at the publisher link above.
- Figures are available only as captions in the extracted text; visual trajectory claims (e.g., Fig 1a step-50 takeoff) could not be independently verified against plotted data.
- The supplementary archive referenced in Sec 5.1 (geometric ablations, offline shielding details) was not part of the extracted text and could not be assessed.
- Review is based on a v1 preprint (arXiv:2609.18078, 16 Sep 2026) that has not undergone peer review; the metadata abstract is truncated mid-sentence but the full-text extraction appears complete through the references and Appendices A-B.
- No public code/data URL appears in the extracted text; if one exists in uncaptured material, the codeAvailable and reproReadiness judgments would improve.
Go deeper on this paper
Wrote a paper? Stress-test it before reviewers do.
Upload your manuscript for the same structured scrutiny — reasoning, statistics, evidence, figures, and reproducibility.