Muon-trained transformers grok modular addition and then lose generalization because the hidden representation and the output readout drift apart; freezing the readout and embeddings after grokking removes the collapse.
Diagnosing Catastrophe: Large parts of accuracy loss in continual learning can be accounted for by readout misalignment
1 Pith paper cite this work. Polarity classification is still indexing.
abstract
Unlike primates, training artificial neural networks on changing data distributions leads to a rapid decrease in performance on old tasks. This phenomenon is commonly referred to as catastrophic forgetting. In this paper, we investigate the representational changes that underlie this performance decrease and identify three distinct processes that together account for the phenomenon. The largest component is a misalignment between hidden representations and readout layers. Misalignment occurs due to learning on additional tasks and causes internal representations to shift. Representational geometry is partially conserved under this misalignment and only a small part of the information is irrecoverably lost. All types of representational changes scale with the dimensionality of hidden representations. These insights have implications for deep learning applications that need to be continuously updated, but may also aid aligning ANN models to the rather robust biological vision.
citation-role summary
citation-polarity summary
fields
cs.AI 1years
2026 1verdicts
CONDITIONAL 1roles
background 1polarities
unclear 1representative citing papers
citing papers explorer
-
Post-Grokking Collapse at the Representation-Readout Interface in Muon-Trained Transformers
Muon-trained transformers grok modular addition and then lose generalization because the hidden representation and the output readout drift apart; freezing the readout and embeddings after grokking removes the collapse.