Technical Deep Dives

The Optimizer Wars: How Muon Is Replacing Adam in Frontier AI Training

Kimi K2 trained 15.5 trillion tokens with zero loss spikes using MuonClip. A year later Kimi K3, DeepSeek-V4 and GLM-5 all orthogonalize their updates — and Muon has its own challenger in Aurora.

Gabriele Masetti · 2026-09-09

The optimizer nobody questioned, until now

For roughly a decade, choosing an optimizer for training a large neural network was not really a choice at all. Adam, published by Diederik Kingma and Jimmy Ba in 2014, and its weight-decay-corrected cousin AdamW, became the default so completely that most papers did not bother justifying the pick. Every GPT model, every Llama release, every open-weight Chinese model through 2024 trained its dense layers with some flavor of AdamW.

That eleven-year run of unchallenged incumbency ended in the summer of 2025, when Moonshot AI shipped Kimi K2 — a trillion-parameter mixture-of-experts model pretrained entirely with a new optimizer called MuonClip. Kimi K2's technical report states plainly that the model "was pre-trained on 15.5 trillion tokens with zero loss spike," a stability claim that AdamW-trained models at comparable scale have never matched.

Fourteen months later, Muon-family optimizers train Zhipu AI's GLM-5, Moonshot's Kimi K3, DeepSeek-V4 and Arcee AI's Trinity models. Adam vs Muon is no longer a research curiosity; it is the live question in frontier pretraining, and the first new AI training optimizer since Adam itself to reach genuine production scale.

Understanding why requires separating two different questions that get conflated in casual comparisons: does Muon train faster per token than AdamW, and does it train more stably at the scale where mixture-of-experts models now live. The evidence below answers both, with numbers traceable to Keller Jordan's original benchmarks, Moonshot's technical reports, and the labs that have since followed.

What is Muon? Momentum orthogonalized by Newton-Schulz

Muon — MomentUm Orthogonalized by Newton-Schulz — was introduced by independent researcher Keller Jordan in a blog post published December 8, 2024, initially as an optimizer for the hidden, two-dimensional weight matrices of a neural network (embeddings and the final classifier layer are left to AdamW). The algorithm runs ordinary SGD-momentum first, then replaces the resulting update matrix with the nearest semi-orthogonal matrix before it touches the parameters.

Jordan's own justification for that step is specific: "updates produced by both SGD-momentum and Adam for the 2D parameters in transformer-based neural networks typically have very high condition number." In plain terms, a handful of directions in the update dominate while many other useful directions are nearly ignored; orthogonalizing the update flattens that spectrum so every direction gets a comparable-magnitude nudge, rather than letting a few dominant singular directions crowd out the rest.

Computing an exact orthogonalization would mean a full singular value decomposition on every matrix, every step — far too slow for GPU training. Muon instead approximates it with a Newton-Schulz iteration, a matrix recurrence that converges toward the orthogonal component of a matrix using only matrix multiplications, which run efficiently in low-precision bfloat16 on GPU tensor cores. Jordan's implementation uses a quintic (fifth-order) Newton-Schulz variant with fixed coefficients (a, b, c) = (3.4445, −4.7750, 2.0315), typically iterated five times per update.

The extra cost is small: Jordan reports the orthogonalization overhead at roughly 0.5–0.7% of total training FLOPs on NanoGPT- and Llama-405B-scale runs. On the NanoGPT speedrunning benchmark, Muon delivered a 1.35x improvement in wall-clock time to a fixed validation loss over the prior AdamW-based record, and it also set the standing CIFAR-10 speed record.

In an internally reported 1.5B-parameter comparison, Muon reached target loss in 10 hours on 8×H100 GPUs against 13.3 hours for AdamW on the same hardware.

Why Adam(W) ruled for a decade — and where it strains at scale

AdamW's staying power came from a genuinely good trade: per-parameter adaptive learning rates via running first- and second-moment estimates, decoupled weight decay, and a body of engineering lore — learning-rate warmup, gradient clipping, WSD schedules — built up over a decade of production runs. None of that made AdamW cheap, however.

Its adaptive moments are notoriously prone to compounding instability as models and token budgets scale past the trillion-parameter, multi-trillion-token range that today's frontier mixture-of-experts models occupy — the exact regime where loss spikes, requiring manual rollbacks and learning-rate surgery, have been reported across multiple public pretraining runs.

Moonshot AI's February 2025 paper, "Muon is Scalable for LLM Training," attacked that scaling gap directly: after adding weight decay to Muon's update rule and adjusting per-parameter update magnitude to keep a consistent root-mean-square scale across matrix and non-matrix parameters, the team's scaling-law experiments found Muon "achieves approximately 2× computational efficiency compared to AdamW" under compute-matched training.

That paper also introduced Moonlight, a 3B/16B-parameter Muon-trained mixture-of-experts model pretrained on 5.7 trillion tokens, as the proof of concept for scaling Muon past small research runs — five months before Kimi K2 took the same idea to a trillion parameters.

What is MuonClip? Kimi K2's fix for exploding attention scores

Scaling Muon to Kimi K2's size exposed a failure mode that smaller runs hadn't surfaced: query-key attention logits growing without bound. Because Muon's orthogonalized updates push weight matrices harder and more uniformly across directions than Adam's adaptive steps do, the query and key projection matrices inside attention heads could drift into a regime where the dot product between query and key vectors — the raw attention logit — exploded before the softmax, producing NaNs or catastrophic loss spikes.

Kimi K2's technical report frames the fix directly: "We propose the MuonClip optimizer, which improves upon Muon with a novel QK-clip technique to address training instability while enjoying the advanced token efficiency of Muon."

QK-Clip works after the fact, not by clamping the attention scores themselves, but by rescaling the weights that produced them. For each attention head h, the algorithm tracks the maximum attention logit reached during the step, S_max^h. If that value exceeds a fixed threshold τ (set to 100 for the full Kimi K2 run), the query and key projection weights for that head are rescaled by a factor γ_h = min(1, τ / S_max^h), applied as √γ_h to each of W_q^h and W_k^h so the product of the two scales down by γ_h overall.

Crucially, this happens as a post-processing step on the weights after the ordinary Muon update — it does not alter the forward or backward pass computation for the step in which the spike was detected; it only constrains the weights going into the next step.

In Kimi K2's run, the clip was actively binding for roughly the first 30% of training steps, capping max logits at the 100 threshold, before attention scores naturally settled into a stable range around 30 and the clip stopped triggering for the remainder of the 15.5-trillion-token run.

Zero loss spikes over 15.5 trillion tokens: why that number matters

"Zero loss spike" sounds like a minor engineering footnote, but at trillion-parameter scale it is closer to the whole ballgame. A loss spike during pretraining typically forces a team to halt the run, roll back to an earlier checkpoint, skip or reorder the offending data batch, and sometimes lower the learning rate — each recovery cycle burns GPU-hours that were already sunk into the spiked steps, and repeated instability can force a full restart of a training phase.

Public post-mortems from several large open pretraining runs have documented exactly this pattern of spike-and-recover cycles. A trillion-token-scale run that completes with no such interventions means every one of the 15.5 trillion tokens Kimi K2 saw contributed to the final model, on the first attempt, with a training schedule set once at the start (constant learning rate of 2e-4 for the first 10 trillion tokens after a 500-step warmup, then cosine decay to 2e-5 over the remaining 5.5 trillion) and never revised mid-run.

That is a reproducibility claim as much as a stability one: a fixed recipe that a team can rerun, port to new hardware, or hand to another lab without expecting a different, spikier outcome each time.

Why is Kimi K2 cheap to train?

Moonshot has not published an audited training budget for Kimi K2, and the company's leadership has pushed back publicly on unofficial cost figures circulating for the later K2 Thinking variant, saying the number "isn't official" because research and experimentation costs are hard to separate from the final pretraining run.

Independent compute estimates nonetheless put the base Kimi K2 model's pretraining at roughly 2.8 million H800-GPU-hours across 14.8 trillion tokens — on the order of $5.6 million at typical cloud GPU rates, a figure comparable to DeepSeek-V3's widely cited training cost and a small fraction of what Western frontier labs have disclosed for models of similar scale.

Three factors compound into that efficiency: Muon's roughly 2x sample efficiency over AdamW per Moonshot's own scaling-law study, meaning fewer tokens are needed to hit a given loss; a sparse mixture-of-experts architecture that activates only 32.6 billion of the model's 1.04 trillion total parameters per token; and, per the stability argument above, a full 15.5-trillion-token run completed without the restart-and-recover cycles that quietly inflate the effective compute bill of less stable training runs.

None of the three would help much without the other two — a sample-efficient optimizer is worth little if it destabilizes at scale, and a sparse architecture is worth little if the optimizer can't train it reliably.

Muon optimizer vs Adam: who else has adopted it

Kimi K2 was not the last word. Zhipu AI's GLM-4.5, released the same month as Kimi K2 in July 2025 at 355 billion total and 32 billion active parameters, trained with standard Muon on its hidden layers. Zhipu's follow-up GLM-5, released February 11, 2026 at 744 billion total and 40 billion active parameters across 80 layers and 256 experts, pushed further with a variant its technical report calls Muon Split: projection matrices are split into smaller, per-head blocks before orthogonalization, which the team found necessary to make Muon compatible with GLM-5's multi-latent attention without a performance regression.

DeepSeek made the same switch at its own frontier. The DeepSeek-V4 report, posted to arXiv on 26 April 2026, covers two models — V4-Pro at 1.6 trillion total and 49 billion activated parameters, V4-Flash at 284 billion and 13 billion — trained on more than 32 trillion tokens using, in the abstract's own words, "the Muon optimizer for faster convergence and greater training stability."

Outside China, Arcee AI's Trinity Large — a roughly 400-billion-parameter, 13-billion-active sparse MoE trained on 2,048 Nvidia B300 GPUs and released in preview in February 2026 — used Muon for its hidden layers and AdamW for embeddings and output projections, the same split Jordan's original recipe recommended, making it one of the first US-trained models to adopt the approach at scale. Arcee shipped a reasoning variant, Trinity-Large-Thinking, under Apache 2.0 on 1 April 2026.

NVIDIA has since built first-class Muon support into Megatron-Core and the NeMo Megatron Bridge; its own GB300 NVL72 benchmarks found Muon training at roughly 1,080 TFLOPs/s/GPU against 1,051 for AdamW on a Kimi K2-scale configuration, a small but positive delta the company attributes to counting the Newton-Schulz matrix multiplications as useful model FLOPs.

Per-head Muon, Kimi K3, and the optimizer chasing Muon

Kimi K3 reached Moonshot's apps and API on 16 July 2026, and the open weights and technical report followed on 27 July, the date the company had given. Between the two the lab shipped three intermediate trillion-parameter models: K2.5 on 27 January 2026, K2.6 in April and K2.7-Code on 12 June.

The published report replaces the estimates that circulated before it. K3 is "a 2.8T parameter Mixture-of-Experts model with 104 billion activated parameters, native vision capabilities, and a 1-million-token context window," built on Kimi Delta Attention and Attention Residuals, with a routing layer called Stable LatentMoE that activates 16 of 896 routed experts per token. Moonshot puts the combined gain at "approximately 2.5x improvement in overall scaling efficiency over Kimi K2."

The optimizer line held through the jump. K3 trains its matrix parameters with Per-Head Muon, orthogonalizing each attention head's projection separately rather than the full matrix at once, paired with the weight-clipping mechanism introduced for Kimi K2. Moonshot's answer to instability is still a constraint applied to the weights after the update, never a rewrite of the update itself. K3 went generally available on Amazon Bedrock on 18 September 2026.

Muon's own turn as incumbent is already being contested. Tilde Research published Aurora on 5 May 2026 after finding that Muon quietly starves part of the network: "by step 500, more than one in four neurons are effectively dead," because weak rows of a tall weight matrix keep receiving weak updates. Aurora constrains each update to be orthogonal and uniform in row norm at once, and Tilde reports it as a drop-in replacement costing about 6% over Muon.

A decade of Adam's default status ended with one trillion-parameter model that didn't spike. Fourteen months on, the optimizer that ended it is the one being patched.