🤖 machine learning

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

This paper demonstrates that Muon-trained transformers exhibit a "post-grokking collapse" where generalization fails after learning modular arithmetic due to instability at the representation-readout interface, a phenomenon characterized by divergent optimizer dynamics and masking effects that can be resolved by freezing specific model components or rescaling the task-aligned circuit.

Original authors: Ali Janati, Kaoutar El Maghraoui, Andrei Kanavalau, Anass Belfatmi

Published 2026-08-10

📖 7 min read🧠 Deep dive

CC BY 4.0

Original authors: Ali Janati, Kaoutar El Maghraoui, Andrei Kanavalau, Anass Belfatmi

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are teaching a robot to solve a very specific math puzzle: adding two numbers together and finding the remainder when divided by a large prime number. At first, the robot seems to be just memorizing the answers, like a student cramming flashcards. But then, something magical happens called "grokking." Suddenly, after a long period of seemingly no progress, the robot stops memorizing and actually understands the pattern. It can solve the puzzle perfectly, even on numbers it has never seen before. This is a big deal for scientists because it helps them understand how artificial intelligence learns to think rather than just memorize.

However, there's a catch. Sometimes, after the robot figures out the solution, it starts to forget it again, even though it's still being trained on the same puzzle. This is like a student who finally understands algebra, but then starts making silly mistakes on the same problems they just solved. Scientists use special tools to look inside the robot's "brain" to see what's happening. They look at the specific patterns the robot uses to solve the math (like musical notes in a song) to see if the robot is still using the right notes. The big question is: Why does the robot forget? Is it because it stopped learning the notes, or is something else going on?

The Mystery of the Vanishing Genius

In this study, researchers looked at a special type of AI robot called a Transformer, trained to solve modular addition puzzles. They used two different "teachers" (optimizers) to teach the robot. One teacher, named AdamW, is the standard, reliable instructor. The other, named Muon, is a newer, faster teacher that seems to help the robot understand the puzzle much quicker.

The researchers found that Muon is indeed a speed demon. It helps the robot reach that "aha!" moment of grokking in about half the time it takes AdamW. But here is the twist: the solutions Muon finds are incredibly fragile. Once the robot solves the puzzle, it often starts to fail again, dropping from 100% accuracy to near zero, only to recover later. This happens over and over again, even though the robot is still being trained on the same problem. AdamW isn't completely safe either, but it fails much less often and less severely.

The Broken Handshake

So, what is going wrong? The researchers discovered that the problem isn't that the robot forgets the math. The robot still knows the solution perfectly inside its "brain." The problem is a broken handshake between two parts of the robot: the part that thinks (the representation) and the part that answers (the readout).

Imagine the robot's brain is a library full of books (the thoughts), and the robot's mouth is the librarian who reads the books out loud (the answer).

The Library (the hidden layers) is writing new notes on the books every second.
The Librarian (the output head) is trying to read those books to give the answer.

With the Muon teacher, the Library keeps rewriting the books in a slightly different language every time it updates, but the Librarian keeps trying to read them in the old language. At first, they manage to understand each other. But because the Library is changing its language faster than the Librarian can adapt, they eventually stop understanding each other. The Library is still writing the correct math, but the Librarian is reading the wrong words.

The researchers proved this by freezing one side. When they stopped the Library from changing its language, the robot stayed perfect. When they stopped the Librarian from changing its reading style, the robot also stayed perfect. But when both were allowed to move freely, they drifted apart, and the robot started failing.

The Ghost in the Machine

Here is the most confusing part: even when the robot is failing, the "Library" still contains the perfect solution. The researchers used a special filter to look only at the part of the brain that knows the math. When they did this, the solution was still 100% correct!

It turns out the robot is suffering from circuit masking. Imagine the Library is shouting the correct answer, but the rest of the robot is shouting a wrong answer at the exact same volume. The Librarian hears a jumbled mess and gives the wrong answer. The correct answer is still there, but it's being drowned out by noise.

The researchers showed that if you simply turn up the volume on the correct part (the math family) without changing anything else, the robot instantly goes back to being 100% accurate. The math didn't disappear; it just got lost in the noise.

Why the Standard Tools Missed It

Usually, scientists check if a robot is working by looking at the "notes" it uses to solve the problem. They check if the right notes are present and if they are loud enough. In this study, those tools said the robot was fine. The notes were there, and they sounded almost exactly the same as before the failure.

But the robot was failing anyway. This is because the tools were looking at the notes but not at the volume or the language the Librarian was using. The notes were perfect, but the Librarian had stopped listening to them correctly. It's like checking a radio and seeing that the station is broadcasting clearly, but not realizing the radio is tuned to the wrong frequency.

The Deep Dive

The researchers also tested what happens when they make the robot "deeper" (adding more layers of thinking). They found that Muon gets even faster at learning, but the problem of the broken handshake gets worse. In deeper models, the robot's brain writes the math in one section, but the part that reads the answer is in a different section, and they drift apart even faster.

They also tried to fix the Muon teacher by removing its special "normalization" trick (a way of keeping the updates steady). Without this trick, the robot learned the math in a very narrow, concentrated way, but it eventually crashed completely and couldn't recover. This suggests that while the special trick helps the robot learn fast and spread its knowledge out, it also makes the robot more prone to this specific type of drifting failure.

The Takeaway

The main lesson from this paper is that learning fast can be dangerous if the parts of the robot aren't synchronized. The robot didn't forget the math; it just lost the connection between the part that knows the math and the part that speaks it.

The researchers found that if you freeze the "speaking" part (the output and embeddings) after the robot learns the solution, the robot stays stable and never forgets. This suggests that the instability isn't because the robot is bad at math, but because the two halves of the robot are moving at different speeds and losing their shared language.

In short: The robot is a genius who suddenly starts stuttering because its brain and its mouth are speaking different dialects. The math is still perfect inside, but without a way to keep the dialects aligned, the genius can't show off its skills.

Technical Summary: Post-Grokking Collapse at the Representation–Readout Interface in Muon-Trained Transformers

Problem Statement
The paper investigates a specific instability occurring in Transformer models trained on modular arithmetic (specifically modular addition and subtraction) after they have achieved "grokking"—the phenomenon where generalization emerges long after training accuracy has saturated. While prior work established that Muon (an orthogonalizing optimizer) accelerates grokking compared to AdamW, this study reports that Muon-trained models frequently suffer from a recurrent, abrupt collapse in generalization performance post-grokking. Crucially, this failure is not limited to Muon; AdamW models also exhibit instability. The paper demonstrates that neither optimizer is strictly stable across five seeds: the selected AdamW baseline records sub-threshold evaluations on four of five seeds (reaching a minimum of 27.59%), while Muon records them on all five. The core problem is that these models lose the ability to generalize despite the training set remaining solved, and standard mechanistic progress measures (spectral support and power distribution) fail to detect the failure, reporting an intact circuit even as accuracy drops to near-chance levels.

Methodology
The authors employ a split-optimizer training regime where hidden weight matrices are updated by Muon, while token/position embeddings and the output unembedding (readout) are updated by AdamW. The study utilizes modular arithmetic tasks ( $a+b \pmod p$ and $a-b \pmod p$ ) with $p=113$ (and variants with $p=97$ ) on decoder-only Transformers of varying depths (1, 2, and 4).

Key methodological components include:

Matched Branching: Trajectories are branched from bit-identical parameter states to isolate causal factors. One branch continues with standard optimization, while others freeze specific parameter groups (hidden matrices, embeddings, or readout) to test necessity.
Fourier Decomposition: The authors decompose the model's internal residual states into 2D Fourier modes over the operand grid. They define specific "families" of modes (e.g., the addition family $(k, k)$ , subtraction family $(k, -k)$ ) and use filtering interventions to measure the sufficiency and necessity of these families for task performance.
Intervention Suite: The study applies interventions such as phase scrambling, frequency relocation, and cross-readout substitution to determine the exact nature of the learned algorithm and the source of misalignment.
Ablation: The normalized-orthogonalization path of Muon is ablated to test its role in the observed spectral dispersion and instability.

Key Findings and Results

The Nature of the Collapse: The failure is a "hidden–readout misalignment." As training loss approaches zero, the gradient vanishes, but the two optimizers respond differently. Muon's update step size becomes largely independent of the gradient magnitude (elasticity $\approx -0.03$ ), while AdamW's remains proportional (elasticity $\approx +1.5$ ). This causes the hidden representation and the readout to drift apart at a rate of 8.0 times per parameter. The training set remains solved, but the readout can no longer decode the representation it was trained alongside.
Masking vs. Circuit Failure: The paper distinguishes two failure modes. In circuit failure, the Fourier family itself stops computing the task. In circuit masking, the task-aligned family (e.g., the addition family) continues to solve the task perfectly in isolation (100% accuracy), but its contribution is overwhelmed by an adversarial remainder in the full model. The authors show that rescaling the task-aligned component alone restores full model accuracy, proving the circuit was intact but "outvoted."
Optimizer Differences: Muon reaches grokking significantly faster (median 13,011 steps vs. 30,486 for AdamW) but is universally unstable across all tested configurations. AdamW is also unstable across five seeds, though the severity of the collapse (measured by the count of sub-threshold evaluations) is two orders of magnitude lower than Muon's. In a single-seed sweep, AdamW configurations showed conditional instability where severity scaled with the learning rate, but replication across five seeds confirmed that the selected "stable" baseline was not strictly stable across seeds. The instability in Muon is attributed to its update rule, which spreads the representation across a wide spectrum (326 effective conjugate pairs vs. ~5 for AdamW), leaving the task-aligned family with a smaller share of the total power.
Task Selection of Family: The task, not the optimizer, selects the Fourier family. Models trained on addition use the $(k, k)$ modes; those trained on subtraction use $(k, -k)$ . The basis within the family is exact and not shared; two models trained from the same seed diverge to use different frequencies and phases, and their readouts become nearly orthogonal, failing to decode each other's representations.
Depth and Alignment: In deeper models, the Fourier code is written by an MLP in the penultimate block but only becomes "legible" to the unembedding in the final block. The misalignment occurs specifically at the interface between the final block's residual and the unembedding.
Prevention: Freezing either the hidden matrices or the non-hidden parameters (embeddings and readout) after the circuit forms completely prevents the collapse. "Stable Muon," where embeddings and readout are frozen post-grokking, achieves strict stability across 451,400 post-grokking steps and five seeds, whereas standard Muon fails repeatedly.

Significance and Claims
The paper claims to identify a fundamental limitation in split-optimizer training and post-grokking dynamics: the residual stream has no privileged basis, and when the loss surface flattens, the loss provides no pressure to maintain the specific alignment between the representation and the readout. The authors argue that:

Standard measures are insufficient: Spectral measures (support and power distribution) can remain invariant during a collapse because the failure is a change of basis (misalignment) rather than a loss of the algorithm itself.
Grokking is reversible masking: The transition from generalization to collapse is the inverse of the transition from memorization to generalization; both involve the task-aligned family's share of representational power rising or falling relative to the adversarial remainder.
The interface is the locus of failure: The instability is not a property of the hidden representation or the readout alone, but of their joint identification. Fixing one side of the interface stabilizes the system.
Muon's speed comes with a cost: While Muon accelerates the discovery of the circuit, its update rule creates a representation that is more susceptible to misalignment with the readout under standard training conditions, necessitating intervention (freezing) for long-term stability.

The study does not propose new applications but offers a mechanistic account of a specific failure mode in mechanistic interpretability, highlighting that a model can possess a correct circuit while failing to utilize it due to representational drift.