When Grokking Runs Backwards: Muon, Modular Addition, and a Gain Knob Nobody Was Watching

Paper: 2608.07436 Authors: Ali Janati, Kaoutar El Maghraoui, Andrei Kanavalau, Anass Belfatmi Categories: cs.AI, cs.LG

The Gap

Grokking research has largely treated the phenomenon as a one-way door. Train a small transformer on modular addition, watch train accuracy hit 100% while validation flatlines, wait, and then validation snaps to 100% as the model discovers the Fourier-multiplication circuit. The standard toolkit that grew around this — Fourier support of the embeddings, power concentrated on a few frequencies, progress measures built from those quantities — implicitly assumes that once the circuit is present, the story is over. What remains is a question of timing: which optimizer, which weight decay, which data fraction gets you through the door sooner.

Muon enters as a timing improvement. Under the now-standard split, Muon takes the hidden matrices while AdamW keeps embeddings and the output head, and it does grok faster. This paper’s move is to keep training past the moment everyone else stops recording. Under that longer horizon, all nine configurations on (a+b) mod 113 grok and then lose generalization. The instability is not a Muon quirk either: the selected AdamW reference falls below threshold on four of five seeds, bottoming out at 27.59%.

The second gap is diagnostic rather than empirical. When accuracy falls, the natural question is “which part broke?” — and the paper argues the question is malformed. The readout composed with the representation is what the loss constrains; the factorization into “a representation” and “a readout” is free up to an invertible map. So the failure can only be located at the interface, jointly. And the standard Fourier progress measures, it turns out, do not move at all across the collapse.

[Problem] grokking is recorded as a one-way door:
          memorize == then generalize == then stop logging
                |
                v
[Assumption under test] after validation saturates the
   solution is stable and optimizer choice only sets
   how fast you arrive
                |
                v
[Method] 1. keep training. 9 configs Muon-hidden + AdamW-embed
         2. measure step-size response to a vanishing gradient
         3. causal test: from bit-identical states freeze one
            group at a time and re-run paired seeds
         4. read the circuit directly via Fourier family filters
                |
                v
[Evidence] 9 of 9 grok then fall below threshold
           post-fit gradient ~ 1e-6
             elasticity: Muon -0.03  vs  AdamW +1.5
             Muon group moves 8.0x faster per parameter
           freeze embed/readout: 0 sub-threshold evals over
             451400 post-grokking steps vs 137-321 unfrozen
           task-aligned family alone = 100 while full model = 45.85
           rescale that family ==> 99.9
                |
                v
[Conclusion] the collapse is a gain failure at the
   representation-readout interface not a lost circuit.
   Standard Fourier support is unchanged. Power cosine 0.9899
   ==> existing progress measures are blind to it

The Increment

One sentence: Before this paper, grokking was a transition you waited for and Muon was a way to wait less; after it, grokking is one resolution of a recurring contest at the representation-readout interface, and the same contest can resolve downward — invisibly to every Fourier progress measure in current use.

Core Mechanism

Start with the state right after the training set is solved. The gradient collapses to order 10^-6, because there is nothing left to fit. Now the two optimizer groups respond to that silence very differently, and the paper quantifies this with what it calls step-size elasticity: how much the update magnitude shrinks when the gradient magnitude shrinks. AdamW’s elasticity is +1.5 — superlinear, so a hundredfold drop in gradient produces roughly a thousandfold drop in step. Muon’s is -0.03, which is to say approximately zero: orthogonalization and normalization make the update size essentially independent of how loud the gradient is. Per parameter, the Muon group keeps moving 8.0x faster than the AdamW group. Nobody asked either group to move; one of them mostly stopped, the other did not.

That drift lands on a product. Logits are readout times representation, and the loss constrains only the product — swap in W M^-1 and M H for any invertible M and nothing observable changes. So there is no fact of the matter about whether “the representation degraded” or “the readout degraded.” What can degrade is their alignment. The causal test is clean: fork from bit-identical states and freeze one group. Freezing embeddings and readout eliminated the failure across five runs and 451,400 post-grokking steps with five paired seeds — unfrozen arms logged 137 to 321 sub-threshold evaluations, frozen arms logged none. Note the asymmetry: the group whose freezing was tested at length is the slower-moving one. Speed is not the same as being the causal handle.

The third component is the readout probe. Filter the logits down to the Fourier family the task uses — (k,k) for addition, and the family swaps to (k,-k) under subtraction, which is a nice confirmation that the task, not the optimizer, selects it. Across 43 checkpoints over five seeds and three regimes, that filtered family alone reaches exactly 100%. This splits failures into two kinds that look identical on an accuracy curve. In circuit failure, the filtered family also stops solving the task: the mechanism is genuinely gone. In masking, the family remains perfect while the full model sits at 45.85% — it assigns a positive margin to the correct answer on every example, including the ones the model gets wrong, and is simply outvoted by a near-equal adversarial remainder. Rescale the family and 99.9% comes back. Grokking, on this account, is the same contest resolving upward.

        a          b
        |          |
   [ embeddings and readout ]  <== AdamW   elasticity +1.5
        |          |                        goes quiet when grad ~ 1e-6
        v          v
   [    hidden matrices     ]  <== Muon    elasticity -0.03
        |                                   8.0x faster per parameter
        v                                   keeps stepping regardless
   H = representation ---.
                          \
                           >  INTERFACE:   logits = W . H
                          /    loss pins only the product
   W = readout -----------'    W M^-1 . M H is unconstrained
        |
        v
   logits decompose into Fourier families
        |
        +-- task-aligned family    (k k) add   /   (k -k) sub
        |      alone = 100  positive margin on EVERY example
        |      including the ones the full model gets wrong
        |
        +-- remainder
        |      near-equal magnitude and opposed ==> outvotes it
        |
        v
   observed accuracy = 45.85     rescale family ==> 99.9
   meanwhile Fourier support unchanged  power cosine 0.9899

   ablation: remove Muon normalization and orthogonalization
     effective conjugate pairs 326 ==> 4
     no recurrent collapse and terminal failure instead
     ==> the thing causing the drift is also what builds the circuit

Think of a recording studio mixing desk. The hidden matrices are the musicians; the task-aligned Fourier family is the one soloist who actually knows the melody, and the remainder is a large ensemble playing something close to noise. The readout and embeddings are the gain knobs — they decide how loud each contribution arrives at the master output, which is the logits. Grokking is the moment the soloist’s fader comes up far enough to carry the mix.

Now the two optimizers are two engineers with different habits. AdamW turns knobs in proportion to how badly the mix sounds; when the mix is already fine, it nearly stops touching anything. Muon turns each knob by a fixed amount every second, whether or not anyone complained, because normalization strips out the “how bad is it” signal from the step size. So after the mix is solved, one engineer sits back and the other keeps twisting. Eventually the ensemble is as loud as the soloist and the master output turns to mush — even though the soloist never played a wrong note. That is masking: solo the correct channel and it is still perfect, push its fader back up and the mix is restored. Circuit failure is the different disaster where the soloist actually forgot the melody, and no amount of fader work helps.

Two extra pieces of the studio map cleanly. Taping the gain knobs in place, i.e. freezing embeddings and readout, prevents the mush for 451,400 steps — the mix stays where it was. And firing Muon’s normalization is not a fix, because that same normalized stepping is what recruited the full ensemble in the first place: without it you go from 326 effective conjugate pairs down to 4, which is a band too small to ever play the piece. You stop having recurrent mush and get permanent silence instead. Finally, the gauge ambiguity: you can rename every musician and re-scale every fader in a compensating way and the master output is bit-identical. So asking “was it the band or the desk?” has no answer. Only the alignment between them is real.

Key Concepts

• Step-size elasticity: Take the log of how big your gradient is and the log of how big a step your optimizer takes, and ask for the slope between them. Plain SGD has a slope of 1: half the gradient, half the step. AdamW here measures +1.5, meaning it decelerates faster than the gradient shrinks, so a near-solved loss basically parks it. Muon measures -0.03, effectively zero, because orthogonalizing the update throws away magnitude information and keeps only direction — a unit-ish step no matter how faint the signal. On its own this is close to a restatement of what normalized optimizers do by design. What makes it consequential is the regime: at gradient 1e-6 the loss has stopped issuing instructions, and a scale-invariant optimizer does not notice, so it keeps performing a random-ish walk through a space of solutions the loss cannot distinguish. Fast arrival and continued wandering are the same property.

• Identified only up to an invertible map: Suppose I hand you the number 12 and tell you it is a product. Was it 3 times 4 or 2 times 6? The question has no answer from the product alone. The interface here is the same situation in matrix form: logits come from readout times representation, and for any invertible M you can put M into the representation and M^-1 into the readout with zero change to any loss, any accuracy, any prediction. So “the representation got worse” is not a well-formed claim; you can always shift blame across the interface at will. What is well-formed is whether the two sides still agree with each other. This is why the paper localizes the failure to the interface rather than to a layer, and why the freezing experiment has to fork from bit-identical states — you need the same 12 to start with before you can say anything about the factors.

• Masking versus circuit failure: Imagine a committee of 100 voting on each question. One member is right every single time. If the other 99 vote randomly and roughly cancel, the committee is right too. Now let the other 99 slowly get louder until they nearly cancel each other but not quite — the residual noise is now as big as the one correct vote, and the committee starts getting things wrong. The expert never became wrong. Measured at 45.85% accuracy, this model still gives the correct answer a positive margin on every example, errors included; it is outvoted, not confused. Rescale the expert’s vote and you are back to 99.9%. The practical sting is that this looks exactly like forgetting from the outside, and the standard Fourier diagnostics agree it is nothing at all — support unchanged, power distribution cosine 0.9899 across an abrupt collapse.

Framework Shift

Before (mainstream)                  After (this paper)
-------------------                  ------------------

val acc                              val acc
100 |      _______________           100 |     /\   /\      /\
    |     /                              |    /  \ /  \    /  \
    |    /                               |   /    v    \  /    \
  0 |___/                               0|__/           \/      \
    +----------------------->             +----------------------->
        time                                    time

one axis of merit:                   two axes of merit:
  is the circuit present  y / n        is the circuit present  y / n
                                       is the circuit loudest  y / n

diagnose with:                       diagnose with:
  Fourier support of embeddings        filter to the task-aligned
  power on top frequencies             family and score it ALONE
  ( both unchanged across a            then compare to full model
    collapse : cosine 0.9899 )         ( 100 vs 45.85 ==> masking )

optimizer question:                  optimizer question:
  how fast do we grok                  what does the optimizer do
                                       AFTER the gradient dies
  [ Muon: faster ]                     [ Muon: keeps stepping.
                                         elasticity -0.03 ]

blame assignment:                    blame assignment:
  layer L degraded                     interface W . H drifted
                                       ( W M^-1 . M H is free
                                         so no layer can be blamed )

From “did the model find the circuit” to “does the circuit still win the vote,” the core shift is that generalization is treated as a contested equilibrium at an interface rather than a discovered object in a layer.

Expert Assessment

Problem choice: Real gap, and a well-chosen one. Grokking papers have a strong incentive to stop logging at the happy ending, and simply extending the horizon is the kind of cheap, obvious experiment that goes unrun for years. The Muon framing is timely given how quickly the hidden-matrices-plus-AdamW-embeddings split became default. My reservation is that the title oversells Muon as the culprit while the abstract itself reports the AdamW reference failing on four of five seeds down to 27.59%. The honest headline is “post-grokking states are not stable equilibria and normalized optimizers make it worse,” which is a weaker but more defensible claim. I also cannot tell from the abstract how carefully this is distinguished from the existing slingshot-effect and ungrokking literature, where cyclic post-convergence instabilities under adaptive optimizers are already documented; that positioning is load-bearing for the novelty claim and its absence from the abstract is a small warning sign.

Method maturity: The diagnostic half is genuinely clever. Refusing to blame a layer because the factorization is gauge-free, then filtering logits to the task-aligned family and scoring it in isolation, is exactly the right instrument, and the (k,k) to (k,-k) swap under subtraction is a satisfying sanity check that the probe tracks the task rather than the optimizer. The “positive margin on every example including the errors” result is the sharpest single sentence in the abstract, because it converts a fuzzy claim about representations into a per-example fact. The mechanism half is thinner. Elasticity -0.03 for an orthogonalized optimizer is close to a definitional consequence of Muon’s update rule, not a discovery; the contribution is noticing that scale invariance is a liability once the gradient dies, which is a good observation dressed as a measurement. And freezing is a fix, not an explanation — it demonstrates that motion in either group is necessary, but says little about why the remainder grows in an adversarial rather than random direction. That “why” is the missing theorem.

Experimental integrity: Better than average on the causal design. Forking from bit-identical states with paired seeds is the right control and I wish more papers did it; 451,400 post-grokking steps with zero sub-threshold evaluations against 137-321 in the unfrozen arms is not a marginal effect. The breadth sweep — two moduli, two widths, two training fractions, subtraction, depth — is appropriate for the claim being made. Three things I would press on. First, “the selected AdamW reference” is doing quiet work; selected how, and against what grid? If the AdamW baseline was chosen for any reason other than a matched-budget sweep, the Muon-versus-AdamW comparison is not load-bearing. Second, five seeds is thin for a claim about intermittent failures, and “137-321 sub-threshold evaluations” depends entirely on evaluation cadence and threshold choice, neither of which the abstract pins down. Third, “326 effective conjugate pairs” exceeds the number of distinct frequency pairs available at modulus 113, so it must be aggregating over neurons or layers — as stated it invites misreading, and the 326 ==> 4 collapse is one of the paper’s headline numbers. I am judging the abstract’s numbers here, not the full tables; several of these concerns could dissolve on a careful read of the methods section.

Writing quality: The abstract is unusually dense and reads like a results table with connective tissue, which I mostly respect — but it front-loads counts and percentages before establishing why the reader should care about the interface at all. The corner cut is theoretical: there is no account of why the remainder ends up near-equal and opposed. A section deriving even a toy version of that — why scale-invariant drift in a loss-flat gauge direction should grow the off-family components at a rate comparable to the on-family ones — would move this from a well-documented phenomenon to an explanation, and would be the single highest-leverage rewrite. Second priority would be an explicit comparison to slingshot and ungrokking phenomenology.

Verdict: weak accept — the causal freezing experiment and the masking-versus-circuit-failure distinction are solid, transferable contributions, and the demonstration that standard Fourier progress measures are blind to a real collapse is worth publishing on its own; the mechanism story is descriptive rather than explanatory and the Muon framing is stronger than the evidence supports.

Takeaways

Specific things worth stealing:

Score the sub-circuit alone, not just the whole model. The most portable technique here is filtering the logits to a hypothesized mechanism and measuring its standalone accuracy alongside the full model’s. When those two numbers diverge — 100% versus 45.85% — you have learned something an ablation cannot tell you: the mechanism is intact and being outvoted. This generalizes anywhere you can decompose a readout into candidate families: attention heads, SAE features, expert routes. The stronger version is the per-example margin check on the failures specifically, since “correct margin on every error case” is a very hard fact to explain away.

Treat scale-invariant optimizers as a liability in the flat regime, not a virtue. Normalized and orthogonalized updates are usually sold as robustness to gradient scale. Reframe it: they do not know when to stop. Once your loss is effectively solved and the gradient is at 1e-6, a scale-invariant optimizer keeps taking full-size steps through directions the loss cannot see. If you use Muon or any normalized update rule with long training tails, measure the log-log slope of update norm against gradient norm for each parameter group — it is cheap — and expect trouble where it is near zero. And note the asymmetry the paper found: the group moving 8x faster per parameter was not the one they had to freeze. Do not assume the fastest-moving group is the causal one.

Look for gauge freedom before you assign blame to a layer. Any time your quantity of interest is a product of two learned factors, you can shift degradation from one factor to the other with an invertible map and change nothing observable. If your interpretability claim would flip under such a reparameterization, it is not a claim about the model. The practical consequence is to define metrics on the composition or on the alignment between factors, and to fork from bit-identical states when you want a causal comparison across the interface.

Your progress measure may be blind by construction. Fourier support unchanged and a power-distribution cosine of 0.9899 across an abrupt accuracy collapse is a clean cautionary result. Support-based and concentration-based metrics discard relative magnitude, which is precisely the variable that failed. Before trusting a progress measure, ask what it is invariant to, then ask whether that invariance covers a failure mode you care about.

And practically: keep training past the good news. If you evaluate at the step where your metric peaks, you cannot distinguish a stable solution from a transient one.