SOTA-121: Use Muon with decoupled weight decay and AdamW-matched update RMS in place of AdamW

Source

Falcon-LLM Team (2026), LIT-119 — the Falcon-H1-Tiny technical blogpost.

Muon itself is Jordan et al. (2024), LIT-159 — an optimizer for the hidden layers, orthogonalising the momentum update via Newton-Schulz iterations and leaving embeddings and the head to AdamW. It is a blog post rather than a paper, which is why the record holds it under a url: (ADR-009), and it is what every practice in this line modifies.

Muon is one member of the class SOTA-165 recommends: an optimizer that multiplies the gradient by a matrix, here by orthogonalising the momentum, rather than scaling each coordinate on its own. That practice is the claim that membership of the class is what buys the speed-up; this one picks the member the record has production evidence for, and has no case without it.

Muon as modified in ARXIV-2502.16982 (LIT-122): weight decay applied to the orthogonalised update, and the update’s RMS rescaled to match what AdamW would produce, so that the learning rate and weight decay tuned for AdamW carry over. Under that recipe the authors saw stable training at nearly the same optimal learning rate as AdamW and better downstream evaluations, and adopted it for every model in the series, at 90M and 0.6B.

That observation is LIT-119’s, and it is the whole of what the blogpost contributes to this practice: Muon with Liu et al.’s two modifications, tested against AdamW on Falcon-H1-Tiny, trained stably at nearly the same optimal learning rate and evaluated better. It is reported in a sentence, with no table or margin for the Muon-versus-AdamW comparison, which is why the size of the gain below comes from other sources.

Conditions: the comparison here is at tiny scale with a µP-parameterised hybrid Mamba/attention model. The RMS matching is what makes the AdamW hyperparameters transferable; without it the learning rate has to be re-tuned. Muon applies to matrix parameters — embeddings, norms and other vectors keep an Adam-style update.

At a trillion parameters this recipe alone let attention logits run past 1000; SOTA-131 is the addition that bounds them, and the chain is written out there.

How much it buys

Less than the record used to say, and the honest range is wide. LIT-122’s scaling-law runs report roughly 2× the compute efficiency of AdamW. An outside comparison tuning both optimizers separately and judging at the end of training rather than mid-run (LIT-156) gets 1.4× at 0.1B, falling to 1.1× at 1.2B — and finds that ranking two optimizers on intermediate checkpoints can reverse the answer, which is one way the larger figures were reached.

The practice stands at 1.1×, for reasons the number does not carry: the gain is free once the recipe is in place, Muon’s smaller optimizer state and hyperparameter transferability are not what that study measures, and every frontier adopter in the record (LIT-131, LIT-132, LIT-139) trains far above its largest scale. What should not be quoted any more is the 2×.

LIT-153 argues the shrinkage is not intrinsic but an artefact of constant decoupled weight decay fixing the equilibrium weight norm, and recovers 20–30% by pinning the norms instead. If that holds up outside its authors’ group it changes this section again.

The same study is also the comparison with SOAP (SOTA-168), the other matrix preconditioner in the record. Tuning both separately, LIT-156 found Muon best at 1–4× the Chinchilla data-to-model ratio, but overtaken by SOAP (and Kron) at 8× on its 130M and 520M models and at 16× on 130M and 300M; at 1.2B both fall to about 1.1× over AdamW. So the choice of Muon over SOAP here rests on production evidence at scale, not on that comparison, which favours SOAP as training runs further past Chinchilla.

Which parameters the gain is actually paid on

This practice says to use Muon in place of AdamW and does not say where the advantage comes from. Wang et al. (LIT-654) ran the ablation: train with Muon on some blocks and Adam on the rest, everything else matched.

• VO + FFN nearly recovers the full-Muon trajectory, in both gated and ungated FFN settings.
• Query-key contributes little. Applying Muon to W_V alone, or W_O alone, already beats applying it to the whole of QK.
• Within the effective set, W_O matters more than W_V, and W_out more than W_in.

Not a parameter-count effect, and the paper says so: QK and VO are the same size. Their account is that VO and the FFN are the blocks that behave as associative memories, and that orthogonalising an update is what treats each stored outer-product direction alike — THEORY-099.

Nothing here changes the recommendation, for two reasons. The recovery is architecture-sensitive: in the ungated setting VO + W_out alone nearly recovers full Muon, and in the gated setting the same combination falls short. And VO+FFN nearly recovers full Muon rather than beating it, so this is not a reason to run a hybrid — it is an answer to what you are buying.

It does bear on cost. If a deployment cannot afford Muon everywhere, the ablation says where to spend it first; that is a claim from one group at small scale with a 0.7B check, and is not yet a practice.

Known implementations

• Falcon-H1-Tiny (all released checkpoints); Kimi K2, Kimi K3 and DeepSeek-V4 with the additions in SOTA-131