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

Abstract

Muon-trained modular-arithmetic transformers can lose accuracy while retaining linearly decodable task information. Adjacent swaps localize five captured unnormalized failures to AdamW readout updates. Multiplying the actual readout displacement by the large feature mean produces a class-dependent logit offset shared across inputs that nearly reproduces each failure. Training-only decoders recover 98.20-100% held-out accuracy. Correcting cross-entropy derivative errors stabilizes five matched branches through step 100,000; four prospective accurate-CE RMS runs fail through embedding updates.

arXiv comment38 pages, 11 figures. Revised manuscript with expanded experiments and analysis. Preprint. Under review