Post-Grokking Collapse at the Representation-Readout Interface in Muon-Trained Transformers
Muon-trained transformers can grok modular arithmetic and then lose it.
The paper reports that all nine tested configurations on modular addition eventually lost generalization after grokking. The authors pin the failure on the interface between learned representations and the AdamW-trained embeddings/output head, where equivalent internal maps are not fixed by the loss. Freezing either optimizer group from identical states prevents the post-grokking collapse. They also separate true circuit failure from masking, where a still-correct task circuit is outvoted by an adversarial remainder. ArXiv · AI/CL/LG's note
The paper reports that all nine tested configurations on modular addition eventually lost generalization after grokking. The authors pin the failure on the interface between learned representations and the AdamW-trained embeddings/output head, where equivalent internal maps are not fixed by the loss. Freezing either optimizer group from identical states prevents the post-grokking collapse. They also separate true circuit failure from masking, where a still-correct task circuit is outvoted by an adversarial remainder. ArXiv · AI/CL/LG's note
score 4