Post-Grokking Collapse at the Representation-Readou... | AI Research

Post-Grokking Collapse at the Representation-Readout Interface in Muon-Trained Transformers investigates why transformer models trained with the Muon optimizer frequently lose their ability to generalize after initially "grokking" a task. While Muon accelerates the grokking process for modular addition, the researchers find that these models suffer from a recurring, unstable collapse where test accuracy drops significantly, even after the training set has been successfully solved.

The Nature of the Collapse

The study identifies that the failure is not a loss of the learned algorithm itself, but a "hidden-readout misalignment." Because the transformer’s residual stream lacks a fixed basis, the internal representation and the output head (readout) are only identified jointly. As training progresses with vanishing gradients, the Muon-optimized hidden layers and the AdamW-optimized readout layers begin to diverge.
The researchers observed that the two groups of parameters move at different rates—with the Muon group moving 8.0 times faster per parameter—causing the readout to lose its ability to interpret the internal representation. This is evidenced by the fact that freezing either the embeddings/readout or the hidden layers prevents the collapse entirely.

Experimental Findings

The researchers tested this phenomenon across various configurations, including different moduli, model widths, and depths. Key observations include:

• Optimizer Instability: While Muon is faster, it is not uniquely prone to this failure. AdamW also experiences post-grokking collapse, though the severity is tied to the learning rate.

• Fourier Circuit Integrity: By filtering the model's representation to the specific Fourier family that computes modular addition, the researchers found that the circuit remains intact even when the full model fails. In "circuit masking," the Fourier family continues to solve the task perfectly, but its output is outvoted by an adversarial remainder in the model.

• Task Specificity: The model’s internal representation is highly specific to the task. For example, when trained on subtraction, the model automatically shifts to using the $(k, -k)$ Fourier modes instead of the $(k, k)$ modes used for addition.

Why the Failure Occurs

The paper concludes that the collapse arises because the loss function does not constrain the specific basis in which the model computes. During the late stages of training, when gradients are extremely small (order $10^{-6}$), the different update rules of Muon and AdamW cause the hidden representation and the readout to drift apart in the latent space.
The researchers demonstrate that this is not a failure of the model to learn the task, but a failure of the two optimizer groups to maintain a shared coordinate system. When the researchers manually rescaled the task-aligned Fourier family, they were able to restore the model's accuracy to 99.9% without any additional training, confirming that the necessary information remains present but becomes inaccessible to the readout.

Limitations and Observations

The study highlights that standard progress measures often fail to detect this collapse because they track the presence of the circuit rather than the alignment between the representation and the readout. While ablating Muon’s normalization and orthogonalization prevents the specific "hidden-readout" drift, it leads to a terminal failure where the model's representation collapses entirely. The findings suggest that the speedup provided by Muon acts as an instrument, allowing researchers to observe and intervene in the grokking process at a scale that would be otherwise prohibitive.