REVIEW 3 major objections 5 minor 18 references
Post-Grokking Collapse at the Representation-Readout Interface in Muon-Trained Transformers
T0 review · 3 major / 5 minor · reviewed 2026-08-10 · deepseek-v4-flash
Pith's one-line read Under the Muon/AdamW split, every grokked modular-addition solution later loses generalization; freezing the readout and embeddings after the circuit forms prevents it.
desk verdict Strong empirical paper on a real Muon/AdamW failure mode; the freeze intervention and Fourier filtering are solid, but the 'unselected basis' mechanism needs a better probe than step-size elasticity. read the letter →
The pith
A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.
The reading
What carries the argument
The central object is the residual stream's basis invariance: because the logits depend on the final residual $h$ and the unembedding $W_U$ only through $W_U h$, any invertible $R$ applied to the stream and absorbed by the surrounding matrices leaves the computed function unchanged, so the loss picks no member of the equivalence class $(R h, W_U R^{-1})$. The paper's measuring instrument is a two-dimensional Fourier decomposition of cached residual states over the operand grid, whose task-aligned addition family $(k,k)$ is exactly sufficient for modular addition at every solved checkpoint and provides the readout that separates circuit failure (the family no longer solves the task) from circuit masking (the family solves it but is outvoted by the rest of the representation). The paper also identifies an optimizer asymmetry: over the quiet window Muon's applied step has elasticity near $-0.03$ while the AdamW groups have elasticities near $+1.5$, so the two sides separate at 8.0 times the per-parameter rate. These objects carry the argument: the basis identity predicts that split optimization can separate the representation from the readout without changing spectral statistics, and the Fourier family turns that silent drift into a functional measurement.
What would settle it
Fit the linear transformation that maps a healthy final residual state onto a collapsed one and apply it to the healthy readout: if the collapse is the unselected change of basis the paper describes, the aligned readout should recover the collapsed representation's accuracy; if the readout has genuinely lost the information, it will not.
Extended reading notes
Core claim
The paper's central claim is that under the standard split, where Muon updates the hidden weight matrices and AdamW updates the embeddings and output head, the speedup in grokking is inseparable from a failure mode: all nine Muon configurations grok modular addition and all nine later fall below 95% test accuracy, with post-grokking minima as low as 0.78%. The collapse is located at the representation–readout interface: logits depend on the final residual representation $h$ and the unembedding $W_U$ only through the product $W_U h$, so for any invertible $R$ the pair $(R h, W_U R^{-1})$ implements the same function, and nothing in the loss selects a particular basis. Once the training set is solved, with loss near $1.5\times10^{-7}$, the gradient provides almost no restoring force on that choice, and the optimizer groups behave differently: over a 691-step window the applied step size has elasticity $-0.026$ to $-0.03$ on gradient magnitude for the Muon group against $+1.47$ to $+1.51$ for the AdamW groups, and the hidden group displaces 8.0 times faster per parameter. Branching from bit-identical states shows that freezing either side suppresses the failure, that within the auxiliary group only the unembedding is individually necessary, and that holding the embeddings and readout fixed after circuit formation removes it across five seeds and 451,400 post-grokking steps. Fourier analysis on the residual stream shows why spectral measures miss it: the task-aligned $(k,k)$ family reaches exactly 100% in isolation at every solved checkpoint, and in circuit masking it still solves the task while the full model reaches 45.85%, outvoted by a near-equal adversarial remainder that rescaling the family alone overcomes.
Load-bearing premise
The argument relies on the premise that once the training loss sits near $1.5\times10^{-7}$, the loss no longer pushes the hidden representation and the readout toward any particular alignment, so the two optimizers can drift apart; this premise is inferred from one 691-step window and one collapse at step 44,700, and would fail if the tiny remaining gradients (hidden norms $2.3\times10^{-7}$ to $8.6\times10^{-7}$, non-hidden $1.95\times10^{-5}$ to $3.52\times10^{-5}$) carry a restoring signal the measurement undercounts, or if the float32 evaluation backend misreports the loss floor.
Editorial extensions
If this is right
- Every tested Muon configuration on modular addition groks and then falls below 95% test accuracy, so the grokking speedup under the split comes with a stability cost that a single-seed strict-stability reading would miss.
- Holding the AdamW-managed embeddings and readout fixed once the circuit has formed removes the collapse in five paired seeds and across 451,400 post-grokking steps, so a two-line intervention converts an unstable configuration into the only strictly stable one in the study.
- At the collapse step the dominant-frequency set is identical (Jaccard index 1.0000) and the power-distribution cosine is 0.9899 while test accuracy falls from 100% to 19.04%, so Fourier support-and-power audits cannot be used to certify that a grokked model still computes the task.
- A model can hold a circuit that solves the task perfectly and still answer incorrectly: in circuit masking the $(k,k)$ family gives 100% in isolation while the full model reaches 45.85%, and rescaling the family alone restores 99.9%, so grokking and post-grokking collapse are the same amplitude competition run in opposite directions.
- The instability appears under both optimizers, with AdamW collapsing at high learning rate and at depth, and it grows with depth, so the failure is a property of the split-optimizer/readout interface rather than of Muon alone.
Reading between the lines
- If the mechanism transfers beyond modular arithmetic, mixed-optimizer training on any residual-stream architecture should show the same vulnerability whenever the loss is flat along the unselected basis direction; a cheap early-warning test would be to track the isolated task-aligned family's margin against the remainder, since the paper shows that ratio, not raw accuracy, is what falls before mas
- The paper's own closing predictions imply a direct architecture test it does not run: because the basis symmetry is exact over $\mathrm{GL}(d)$ and restricted to $O(d)$ under RMS normalization, adding normalization with learned gains should measurably reduce or eliminate the collapse on the same task.
- The freeze result suggests a design rule for mixed optimizers: anchor the input and output coordinates once the circuit forms rather than tuning learning rates or weight decay, since the paper's sweeps show those change the timing and severity of collapse but not its occurrence.
- The Fourier-margin decomposition may provide a more robust definition of circuit health than spectral support, because the standard measures were unchanged at the collapse step; tracking the family's isolated accuracy and its margin against the remainder could detect masking before the full model fails.
Editorial analysis
A structured set of objections, weighed in public.
Referee Report
Summary. The paper studies transformer training on modular arithmetic under the standard split-optimizer routing that assigns Muon to hidden matrices and AdamW to embeddings and the output head. It reports that Muon reaches the grokking threshold faster than AdamW alone, but that the resulting solutions are unstable: all nine configurations swept on (a+b) mod 113 grok and subsequently fall below 95% test accuracy, while the selected AdamW baseline is also unstable on four of five seeds. The paper localizes the failure to the representation–readout interface, arguing that near-saturated loss leaves the residual-stream basis unselected, that Muon and AdamW then respond differently to vanishing gradients, and that freezing the AdamW-managed embeddings and readout after circuit formation eliminates sub-threshold evaluations across five runs and 451,400 post-grokking steps. A Fourier decomposition identifies the task-aligned (k,k) family as exactly sufficient at 43 solved checkpoints, validates the family choice through the subtraction reversal to (k,-k), and distinguishes circuit failure from circuit masking, in which the intact family is outvoted by the remainder. The paper also shows that standard Fourier support and power measures can remain essentially unchanged at the step where the function collapses.
Significance. If the mechanism were fully established, this would be a significant contribution to mechanistic accounts of grokking and to practical guidance on the Muon/AdamW split. The empirical scaffolding is unusually careful: matched branches from bit-identical states, paired five-seed arms verified identical before the freeze, exact family sufficiency and ablation measurements with a subtraction reversal control, a linear margin decomposition, and a consistent operational distinction between circuit failure and circuit masking. These strengths make the collapse phenomenon and the freeze intervention credible. The main gap is explanatory: the unselected-basis mechanism rests on a single 691-step window, on step-norm elasticities that do not test alignment with the loss gradient, and on a loss floor close to float32 precision. Those gaps do not undercut the descriptive results, but they leave the proposed mechanism underdetermined.
major comments (3)
- [§4.1, Table 7] The central mechanism is not established by the reported elasticity. Regressing the applied update norm on the gradient norm measures whether the step size tracks the gradient magnitude; it does not measure whether the step is aligned with the loss gradient. A constant-norm Muon step can be orthogonal to the gradient, supporting the unselected-basis premise, or substantially parallel to it, contradicting that premise, while producing the same -0.026 elasticity. The paper reports no dot product, cosine, or parallel component of the hidden-group update relative to the loss gradient, and it does not report the loss change along the actual hidden trajectory. This is load-bearing because the freeze interventions remove one moving group and would suppress a collapse under either mechanism. A direct alignment measurement, or a demonstration that the loss is flat along the hidden displacement, is needed to support the paper's explanatory claim.
- [§4.1, Table 7 and Figure 2] The training loss values of 1.5–1.8×10^-7 are within one order of magnitude of float32 machine epsilon (~1.19×10^-7), and training runs on Apple MPS in float32. The reported trend that the loss 'falls by a fifth' and the gradient norms of order 10^-7 to 10^-6 could therefore be dominated by round-off. If the tiny gradients are numerical noise, the elasticity regression in the same window is fit to noise, and the conclusion that the loss supplies no restoring force is not supported. Running the same 691-step window, or one further collapse, in float64 on CPU, with gradients recorded in double precision, would settle whether the plateau and the elasticity are real.
- [§4.1 and §11] The mechanistic analysis—the 691-step elasticity window, the matched branch table, the cross-readout substitutions, and the layerwise causal analysis—is carried out at a single seed and, for the main analyzed collapse, at a single event. The paper itself states this limitation in the conclusion, and the five-seed replication covers the freeze outcome and the Fourier family rather than the elasticity or alignment statistics. Since the claimed mechanism is the paper's principal explanatory contribution, at least one additional seed or additional collapse window for the elasticity and for the proposed alignment measurement is needed before the account can be regarded as established. This is a strengthening request; it does not reject the descriptive and interventional results.
minor comments (5)
- [§2.3, §4.1] Define precisely the quantity whose elasticity is reported. It is not clear whether the 'gradient-driven applied update' for Muon is the pre-weight-decay orthogonalized step, the post-decay net displacement, or something else, and the same ambiguity affects the AdamW groups. The comparison of -0.03 versus +1.5 is only interpretable if the same convention is used throughout.
- [Figure 1] State explicitly in the caption that the Jaccard index and the power cosine are computed within the addition family, as Section 8.4 later explains; a reader of the caption alone could otherwise take them as global spectral measures.
- [§5.4 and Table 8] Freezing only the unembedding suppresses collapse over the 2,000-step branch but leaves instability over the 90,000-step run; the text explains this difference, but a sentence in Section 5.4 directly citing the different time scales would prevent the apparent contradiction.
- [§7.3 and Table 12] The replicate count for the equal-power relocation (430) appears only in the text; moving it to Table 12 or reporting it in the methods paragraph would help the reader assess the precision of the 2.64% average.
- [§8.1] In the sentence 'At 95 of the 98 checkpoints the family in isolation is at least as accurate as the model containing it,' clarify in the same sentence what the three exceptions are, rather than leaving that information to the figure caption.
Circularity Check
No significant circularity: the central claims are supported by direct measurements, causal freeze branches, and an externally validated Fourier-family reversal, with no fitted quantity or self-citation doing load-bearing work.
full rationale
The paper's central derivation is not circular. The claim that the collapse occurs at the representation-readout interface is supported by direct observations (test and training accuracy falling together) and by causal matched-branch experiments in which freezing either the hidden group or the auxiliary/readout group suppresses the failure from a bit-identical state; these are interventions, not fitted equivalences. The Fourier-family identification is validated externally by the subtraction reversal, where the "(k,k) for addition, (k,-k) for subtraction" roles exchange exactly, ruling out a family chosen to fit the target result. The collapse is detected by raw test accuracy, and the claim that standard spectral progress measures miss it is confirmed by Jaccard index 1.0000 and cosine 0.9899 computed from the same checkpoints rather than assumed. The only fitted quantity in the paper, the log-log elasticity of applied step size on gradient norm, is a descriptive diagnostic of optimizer behavior and is not used as an input to force the collapse prediction; the freeze results and the spectral-invariance predictions do not depend on its numerical value. No load-bearing self-citation appears: the cited prior work on Muon, grokking, and residual-stream basis invariance is external to the authors, and the paper's own novel measurements stand independently. The skeptic concern that elasticity does not test alignment is a legitimate correctness risk, but it is not a circularity: the mechanism could be wrong while the empirical freeze result remains valid, and the paper does not define the collapse into the elasticity. Overall, the derivation chain is self-contained and externally falsifiable, so the appropriate circularity score is 0.
Assumptions & free parameters
free parameters (2)
- 95% test accuracy threshold =
95%
- Freeze delay after sustained criterion =
2000 steps
assumptions (4)
- standard math An invertible change of basis R of the residual stream leaves the function unchanged when all writers and readers are transformed accordingly; the network contains no normalization and no biases, so this symmetry is exact.
- standard math The two-dimensional Fourier partition over operand pairs is exhaustive, and the task-aligned (k,k) or (k,-k) family is exactly sufficient to compute modular addition or subtraction; an orthonormal DFT preserves power.
- domain assumption At a solved training set with loss around 1e-7 and gradient norms around 1e-6 to 1e-5, the loss does not provide a meaningful restoring force that selects a basis for the hidden/readout pair.
- domain assumption Over the measured windows, Muon's Newton-Schulz update does not track the gradient magnitude while AdamW's update tracks it more than proportionally, so the two parameter groups separate.
Cite this review
Pith. "Pith review of Post-Grokking Collapse at the Representation-Readout Interface in Muon-Trained Transformers." pith.science (2026). https://pith.science/paper/NIBH2VUB
@misc{pith2026260807436,
author = {Pith},
title = {Pith review of: Post-Grokking Collapse at the Representation-Readout Interface in Muon-Trained Transformers},
year = {2026},
howpublished = {\url{https://pith.science/paper/NIBH2VUB}},
note = {Machine review of arXiv:2608.07436}
}
abstract
Under the standard split, Muon gets hidden matrices and AdamW embeddings/output head. Muon groks modular addition faster, but its solutions do not hold. All nine configurations on $(a+b) \bmod 113$ grok and later lose generalization. Across five seeds the selected AdamW reference falls below threshold on four, reaching 27.59%. Instability persists across two moduli, two widths, two training fractions, subtraction, and depth. The failure arises at the representation-readout interface, identified only jointly up to an invertible map unselected by the loss. After solving the training set, the gradient falls to order $10^{-6}$ and the optimizers respond differently: step-size elasticity is -0.03 for Muon versus +1.5 for AdamW, and the Muon group moves 8.0 times faster per parameter. From bit-identical states, freezing either group prevents failure. Freezing embeddings/readout removes it in five runs over 451,400 post-grokking steps and five paired seeds: unfrozen arms record 137-321 sub-threshold evaluations, frozen arms none. Removing Muon's normalization and orthogonalization is no substitute: it collapses representation from 326 effective conjugate pairs to 4, shows no recurrent collapse, and fails terminally. Fourier filtering separates circuit failure from masking. Across 43 checkpoints over five seeds and three regimes, the task-aligned family reaches exactly 100% alone. In circuit failure it no longer solves the task; in masking it remains perfect while the full model reaches 45.85%, giving a positive margin on every example, including errors, but being outvoted by a near-equal adversarial remainder. Rescaling it restores 99.9%; grokking is the same condition resolving upward. The task selects the family, swapping $(k,k)$ for $(k,-k)$ under subtraction. Across an abrupt collapse, standard Fourier support is unchanged and the power-distribution cosine remains 0.9899.
Figures
Figures from the paper (3 more)
Reference graph
Works this paper leans on
-
[1]
Daniel Anthes, Sushrut Thorat, Peter König, and Tim C. Kietzmann. Diagnosing catastrophe: Large partsofaccuracylossincontinuallearningcanbeaccountedforbyreadoutmisalignment. InConference on Cognitive Computational Neuroscience, pages 748–751, 2023. URLhttps://arxiv.org/abs/ 2310.05644
work page Pith review arXiv 2023
-
[2]
Two Speeds of Learning: A Representation-Readout Decomposition of Grokking and Double Descent
Chi-NingChou,OscarUzdelewicz,Neng-ChunChiu,Yao-YuanYang,andSueYeonChung. Twospeeds of learning: A representation-readout decomposition of grokking and double descent.arXiv preprint arXiv:2605.27078, 2026
work page Pith review arXiv 2026
-
[3]
A toy model of universality: Reverse engineering how networks learn group operations
Bilal Chughtai, Lawrence Chan, and Neel Nanda. A toy model of universality: Reverse engineering how networks learn group operations. InInternational Conference on Machine Learning (ICML), 2023. URLhttps://arxiv.org/abs/2302.03025
arXiv 2023
-
[4]
A mathematical framework fortransformercircuits.TransformerCircuitsThread,2021
Nelson Elhage, Neel Nanda, Catherine Olsson, Tom Henighan, Nicholas Joseph, Ben Mann, Amanda Askell,YuntaoBai,AnnaChen,TomConerly,NovaDasSarma,DawnDrain,DeepGanguli,ZacHatfield- Dodds, Danny Hernandez, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. A mathematical...
work page 2021
-
[5]
LauraGomezjuradoGonzalez.Thelongdelaytoarithmeticgeneralization: Whenlearnedrepresentations outrun behavior.arXiv preprint arXiv:2604.13082, 2026
arXiv 2026
-
[6]
Muon: An optimizer for hidden layers in neural networks, 2024.https://github.com/ KellerJordan/Muon
Keller Jordan, Yuchen Jin, Vlado Boza, Jiacheng You, Franz Cesista, Laker Newhouse, and Jeremy Bernstein. Muon: An optimizer for hidden layers in neural networks, 2024.https://github.com/ KellerJordan/Muon
work page 2024
-
[7]
Muon is scalable for LLM training.arXiv preprint arXiv:2502.16982, 2025
Jingyuan Liu et al. Muon is scalable for LLM training.arXiv preprint arXiv:2502.16982, 2025. URL https://arxiv.org/abs/2502.16982
arXiv 2025
-
[8]
Progress measures for grokking via mechanistic interpretability
Neel Nanda, Lawrence Chan, Tom Lieberum, Jess Smith, and Jacob Steinhardt. Progress measures for grokking via mechanistic interpretability. InInternational Conference on Learning Representations (ICLR), 2023. URLhttps://arxiv.org/abs/2301.05217
arXiv 2023
Show all 18 references
-
[9]
Grokking: Generaliza- tion beyond overfitting on small algorithmic datasets.arXiv preprint arXiv:2201.02177, 2022
Alethea Power, Yuri Burda, Harri Edwards, Igor Babuschkin, and Vedant Misra. Grokking: Generaliza- tion beyond overfitting on small algorithmic datasets.arXiv preprint arXiv:2201.02177, 2022. URL https://arxiv.org/abs/2201.02177
2022 arXiv
-
[10]
Prakash and Charles H
Hari K. Prakash and Charles H. Martin. Grokking and generalization collapse: Insights from HTSR theory. InHiLD Workshop at the International Conference on Machine Learning (ICML), 2025. URL https://arxiv.org/abs/2506.04434
2025 arXiv
-
[11]
Lucas Prieto, Melih Barsbey, Pedro A. M. Mediano, and Tolga Birdal. Grokking at the edge of numerical stability. InInternational Conference on Learning Representations (ICLR), 2025. URL https://arxiv.org/abs/2501.04697
2025 arXiv
-
[12]
Muon learns more robust and transferable features than adam.arXiv preprint arXiv:2606.09658, 2026
Tianyu Ruan, Fengzhuo Zhang, Shuche Wang, and Shihua Zhang. Muon learns more robust and transferable features than adam.arXiv preprint arXiv:2606.09658, 2026. URLhttps://arxiv.org/ abs/2606.09658. 33
2026 arXiv
-
[13]
The slingshot mechanism: An empirical study of adaptive optimizers and the grokking phenomenon.arXiv preprint arXiv:2206.04817, 2022
Vimal Thilak, Etai Littwin, Shuangfei Zhai, Omid Saremi, Roni Paiss, and Joshua Susskind. The slingshot mechanism: An empirical study of adaptive optimizers and the grokking phenomenon.arXiv preprint arXiv:2206.04817, 2022. URLhttps://arxiv.org/abs/2206.04817
2022 arXiv
-
[14]
Muonoptimizeracceleratesgrokking.arXivpreprint arXiv:2504.16041, 2025
AmundTveit,BjørnRemseth,andArveSkogvold. Muonoptimizeracceleratesgrokking.arXivpreprint arXiv:2504.16041, 2025. URLhttps://arxiv.org/abs/2504.16041
2025 arXiv
-
[15]
Explaining grokking through circuit efficiency.arXiv preprint arXiv:2309.02390, 2023
Vikrant Varma, Rohin Shah, Zachary Kenton, János Kramár, and Ramana Kumar. Explaining grokking through circuit efficiency.arXiv preprint arXiv:2309.02390, 2023. URLhttps://arxiv.org/abs/ 2309.02390
2023 arXiv
-
[16]
Shuche Wang, Fengzhuo Zhang, Jiaxiang Li, Cunxiao Du, Chao Du, Tianyu Pang, Zhuoran Yang, Mingyi Hong, and Vincent Y. F. Tan. Muon outperforms adam in tail-end associative memory learning. arXiv preprint arXiv:2509.26030, 2025. URLhttps://arxiv.org/abs/2509.26030
2025
-
[17]
The active ingredient in muon’s grokking.arXiv preprint arXiv:2607.20512, 2026
Yufeng Wang. The active ingredient in muon’s grokking.arXiv preprint arXiv:2607.20512, 2026. URL https://arxiv.org/abs/2607.20512
2026 arXiv
-
[18]
The clock and the pizza: Two stories in mechanistic explanation of neural networks
Ziqian Zhong, Ziming Liu, Max Tegmark, and Jacob Andreas. The clock and the pizza: Two stories in mechanistic explanation of neural networks. InAdvances in Neural Information Processing Systems (NeurIPS), 2023. URLhttps://arxiv.org/abs/2306.17844. 34
2023 arXiv
Reviewed August 10, 2026 · model on record in the stance chip above.
Discussion (0). Continue with ORCID to comment.