Post-Grokking Collapse at the Representation-Readout Interface in Muon-Trained Transformers

This paper investigates why transformers trained with the Muon optimizer can "grok" (achieve sudden generalization on) modular arithmetic tasks faster than AdamW, yet later lose that generalization. Through detailed analysis of embedding/readout versus hidden-layer dynamics, the authors identify a representation-readout interface failure as the cause, distinguishing genuine circuit failure from mere "masking" effects. This research is primarily relevant to interpretability and optimization researchers studying training stability and generalization dynamics in transformers, with implications for choosing and combining optimizers (like Muon and AdamW) to build models that generalize durably rather than transiently. Paper: https://arxiv.org/abs/2608.07436