About Muon Series

• About Muon: Muon At Scale - Current Post
• About Muon: Why Not Adam?

The previous post made the step run. Every team that took that step to a frontier run still had to change it: share a schedule with AdamW, keep a QK product bounded, and orthogonalize a matrix the cluster does not hold.

Not the wrong norm. The wrong scope. The derivation optimizes one matrix. The run has to optimize a model, and the model is heterogeneous.

Why Not Bernstein's Multiplier?

The labs did not keep Bernstein's multiplier. They scaled UV^T so Muon and AdamW can share \eta. The derivation optimized one matrix. The run has to keep two optimizers on one schedule.

The change from the derivation is one assignment. Hidden 2D weights stay on Muon. Embeddings, the LM head, and every 1D parameter stay on AdamW.

Bernstein vs Moonlight

Part 1 left two candidates for the multiplier on UV^T.

• Bernstein [9]: \sqrt{\texttt{fan-out}/\texttt{fan-in}}, which holds the RMS→RMS displacement at \eta. What is fixed is how far the layer can move its own output.

• Moonlight [2]: 0.2\sqrt{\max(m,n)}, which holds the entry RMS of \Delta W near 0.2\eta. What is fixed is the typical size of a single entry.
  • AdamW's update is sign-like, so its entries are O(1) and its step has entry RMS \approx\eta at any shape.
  • Raw UV^T has entry RMS 1/\sqrt{\max(m,n)}, which shrinks as layers widen. A shared \eta would starve the Muon group.
  • The 0.2\sqrt{\max(m,n)} factor cancels that shape term, so both groups move with \eta at a fixed ratio. One learning rate serves both.

What Moonlight, K2, and GLM Ran

This is not settled by argument. Three runs from two labs agree in practice.

Run	Total / active	Tokens	Multiplier convention	Extra
Moonlight [2]	16B / 3B MoE	5.7T	entry RMS 0.2	decoupled weight decay
Kimi K2 [1]	1T / 32B MoE	15.5T	entry RMS, constant not stated	QK-Clip
GLM-4.5 [3]	355B / 32B MoE	23T	entry RMS 0.2, N=5, \mu=0.95	none reported

Two of those rows are Moonshot, so the first two are one team scaling its own recipe from 5.7T tokens to 15.5T. K2's paper describes the optimizer but does not print the constant, so that row inherits Moonlight's. GLM-4.5 is the independent check: a different lab, no shared authors, and the same 0.2.

"we set the Newton-Schulz iteration steps N to 5, momentum \mu to 0.95, and scaled Muon's update RMS to 0.2."

One lab can share a recipe. Two independent teams choosing 0.2 at 15.5T and 23T tokens is a choice of purpose for the multiplier. Bernstein's convention has the cleaner derivation and zero deployments.

The reason is not that one formula is more correct. The multiplier's job is to let two optimizers share a learning rate. Moonlight's scale removes the need to sweep AdamW's learning rate again.

Without the update-scale, a shared \eta starves the widening Muon group.

Moonlight names a second correction as well, and it is required for the same reason the scale is. That correction is weight decay. Without it, nothing bounds \Vert W\Vert. UV^T is unit-spectral even when the gradient is small, so the step size does not shrink with the gradient. Moonlight puts AdamW's decoupled decay on the Muon group for that reason.

u = newton_schulz(g).to(p.dtype)

# Bernstein: hold the RMS->RMS displacement at eta.
# scale = (p.shape[0] / p.shape[1]) ** 0.5

# Moonlight's published step. GLM reports the same RMS target.
scale = 0.2 * max(p.shape) ** 0.5

p.mul_(1 - lr * wd)
p.add_(u, alpha=-lr * scale)

When Does a Full Step Help?

Giving every direction a full-size step is how Muon extracts more from a large batch. It is also how an attention logit leaves its working range. The softmax then collapses onto one token. The gradient dies. The outcome depends on whether the weak directions carry signal.

Past the critical batch

Muon stays data-efficient past the critical batch, where AdamW has already stopped being data-efficient.

The critical batch size is the point past which doubling the batch stops halving the step count. Each extra example carries less new information than the last, so a larger batch no longer reduces the step count. It can still reduce wall-clock time, if the hardware can process the extra examples in parallel.

Essential AI [4] finds the transferable gain:

"We find that Muon is more effective than AdamW in retaining data efficiency at large batch sizes, far beyond the so-called critical batch size, while remaining computationally efficient."

That is the controlled study. GLM-4.5 reports the same thing from production: Muon "can accelerate convergence and tolerate larger batch sizes" [3]. A different lab, a different architecture, 23T tokens.

A large batch produces a gradient whose singular spectrum carries signal further down. An entrywise optimizer cannot use that extra signal: it rescales each coordinate alone. Orthogonalization equalizes across directions, so more of that signal reaches the update. If that is the mechanism, the advantage should grow with batch size, which is what both sources report.

The measured setting is large-batch, multi-node pretraining. A modest batch does not have this advantage.

When Softmax Collapses

The bound that fails is not one matrix. It is the product W_qW_k^T. Kimi K2 [1] is the largest public deployment: 1T parameters, 15.5T tokens, zero loss spikes. That stability is often credited to Muon. The paper points at QK-Clip. The team set out to "leverage the token-efficient Muon optimizer and mitigate its training instabilities through the introduction of QK-Clip" [1]. The its is Muon, not training in general.

"we observe that the maximum attention logits quickly exceed a magnitude of 1000, showing that attention logits explosion is already evident in Muon training to this scale."

A logit of 1000 going into a softmax collapses onto one token, and the gradient dies. That measurement is from a smaller MoE run, before K2 was trained at all. Vanilla Muon had already driven the logits three orders of magnitude past their working range.

The fix watches S_{\max}, the largest attention logit after the step, and rescales the query and key projection weights W_q and W_k when S_{\max} exceeds a threshold \tau (100 in K2's run). Q and K themselves are not touched.

W_q^h \leftarrow \gamma^{\alpha} W_q^h, \qquad W_k^h \leftarrow \gamma^{1-\alpha} W_k^h, \qquad \gamma = \min\!\left(1,\dfrac{\tau}{S_{\max}}\right), \quad \alpha \approx 0.5

With \alpha = 0.5 both weights are multiplied by \sqrt{\gamma}. The next-step logit is a product of Q=XW_q and K=XW_k, so it is multiplied by \gamma and comes back to \tau, and neither projection absorbs the full rescaling. The Muon update is unchanged. This is a check after the step. In K2's run, S_{\max} rose quickly to the cap of 100, stayed there while the clip was binding, then fell below 100 on its own and the clip stopped applying.

The paper already had two standard guards. It rejects both. That is why QK-Clip exists.

• Logit soft-capping "directly clips the attention logits, but the dot products between queries and keys can still grow excessively before capping is applied". Only the logits are clipped. W_q and W_k still grow. The product grows before the cap. Too late.

• Query-Key Normalization is the usual fix. It "is not applicable to multi-head latent attention". K2 uses MLA.

Fireworks [5] gives a longer account of QK-Clip.

Two things most write-ups skip:

• QK-Clip is not part of Muon. It is an attention-logit patch. Any run with the same failure can use it.
  • MuonClip is four things under one name: Muon, the 0.2 RMS scale, decoupled weight decay, and QK-Clip. The first three are optimizer choices. QK-Clip is not.

• The paper already says the instability is Muon's. Orthogonalization sets every singular value to 1. A direction with almost no gradient signal still gets a full-magnitude step.
  • The attention logit is a product of two such matrices. Full-size steps on noise in both factors compound, and the product leaves the working range.

Past the critical batch, weak directions mostly carry signal. Early in a QK pair they mostly do not. QK-Clip does not change the per-tensor norm; it constrains W_qW_k^T. The logit explosion comes from Muon. QK-Clip is the patch Muon needs under MLA.

What Breaks at Scale

Two things break, and they are not the same failure. A sharded run does not have the full matrix. After UV^T, the row norms of the update are still unequal.

The Matrix Does Not Fit

Part 1 justified Newton-Schulz on the grounds that five rounds of matmul are cheap. That is true on one device, and false on a cluster. Dion [6] states the obstruction: Newton-Schulz needs dense multiplication on the full matrix. Sharded training does not have that matrix.

• Under FSDP or tensor parallelism the weight lives in slices. Gathering those slices every step to run Newton-Schulz is inefficient: it brings back the all-gather that sharding existed to remove.

A full-matrix all-gather moves mn entries. Dion's rank- r path moves r(m+n). That is cheaper when r < mn/(m+n). On a 4096 \times 4096 layer the breakeven is r = 2048. r = 256 is eight times less traffic: 2.1M against 16.8M.

The update is a rank- r sketch of B_t=M_{t-1}+G_t, not Newton-Schulz on the full matrix:

B_t=M_{t−1}+G_t, \quad (P_t,R_t)=PowerIter_1(B_t;Q_{t−1}) \\ M_t = B_t - (1-\mu)\,P_t R_t^T, \quad Q_t = \mathrm{ColNorm}(R_t), \quad X_t = X_{t-1} - \eta\, P_t Q_t^T

\mathrm{PowerIter}_1 is one warm-started power step: P \leftarrow BQ, orthogonalize P, R \leftarrow B^T P. Then B \approx PR^T at rank r. Devices exchange P and R, not B.

M_t = B_t - (1-\mu)P_t R_t^T is error feedback. The rank- r piece is subtracted. The leftover stays in M for the next step. Part 1 treated orthogonalization as all-or-nothing. Dion orthogonalizes a subspace and keeps the residual that the rank- r sketch dropped. The weak directions were worth keeping.

Full orthogonalization does not survive sharding. The reported gain is per-step wall-clock against Muon, on models from 160M to 3B. Loss stays in Muon's range. The step gets cheaper.

Certain Neurons Still Dominate

Orthogonalization equalizes singular values. It does not equalize rows. NorMuon [10]: after UV^T, the update still has "highly non-uniform neuron norms, causing certain neurons to dominate".

The fix is after orthogonalization, not instead of it.

• State: one second-moment scalar per row.

• Divide that row of UV^T by the square root of its variance.

• Rescale the whole matrix back to the Frobenius norm it had going in, so the total step size does not change.

Rows share the step. The Frobenius size of the update is restored. That is not a proof that the step is still the spectral maximizer. The paper reports that the update keeps a low condition number. Reported at 1.1B: 21.74% / 11.31% fewer steps to the same loss than Adam / Muon.

The spectral invariant is too coarse inside one tensor, not only across tensors. AdaMuon [11] makes the same point: keep UV^T 's directions, and stop giving every direction the same step.

Which Norm for Which Layer

One norm, chosen once for every 2D weight, is not enough. The runs each patched one part of the model.

Where	What changed	Where it hits
Moonlight, K2, GLM-4.5	operator norm → entry RMS, so AdamW can share \eta	parameter groups
Kimi K2	a constraint on W_qW_k^T	a pair of tensors
Dion	the invariant kept on a rank- r subspace only	directions
NorMuon	rows renormalized after UV^T	rows of one tensor

Bernstein and Newhouse wrote the prediction before Muon existed [7]:

"Different operator norms should be assigned to different tensors based on the role that the tensor plays within the network. For example, while linear and embedding layers may have the same weight space of \mathbb{R}^{m \times n}, these layers play different roles and should be assigned different norms."

W_q and W_k are that case: same shape, and the bound sits on their product.

The table is four patches, each on a different part of the model. Modular Manifolds [8] writes one rule for those patches: the existing norm, UV^T, stays, and each layer also gets a manifold that W must stay on.

Muon has no manifold for W, so after W \leftarrow W - \eta UV^T, the singular values of W can drift. Manifold Muon is the optimizer in Modular Manifolds that puts W on the Stiefel manifold W^T W = I_n, so every singular value of W is 1.

Call the increment A. Adding A to W changes W^T W by A^T W + W^T A + A^T A. A^T A is the square of A, so it is small when the step is small. The condition that remains is A^T W + W^T A = 0. UV^T of the raw gradient G does not meet that condition. Manifold Muon therefore adds a correction 2W(\Lambda+\Lambda^T) to G before it forms UV^T, and tunes \Lambda until A^T W + W^T A = 0:

G + 2W(\Lambda+\Lambda^T) = U\Sigma V^T, \qquad A = -\eta UV^T

Each update of \Lambda recomputes UV^T. One manifold Muon step therefore runs on the order of a hundred UV^T evaluations.

Manifold Muon is one complement of Muon. The limit is compute: keeping W on that manifold is not yet cheap enough to run. Complements of Muon, and analysis of Muon, are still needed and remain under study.

• Polar Express [13] is one such complement: a better polynomial for UV^T. Five Newton-Schulz steps take a singular value of 10^{-4} to 0.05. Ten Polar Express steps take it to 0.994.

• Grokking [12] is one such analysis: Muon speeds the jump from memorization to generalization.

Reference

[1] Kimi Team, Moonshot AI. (2025, July). Kimi K2: Open Agentic Intelligence. arXiv:2507.20534. https://arxiv.org/abs/2507.20534

[2] Liu, J., Su, J., Yao, X., et al. (2025, February). Muon is Scalable for LLM Training. arXiv:2502.16982. https://arxiv.org/abs/2502.16982

[3] GLM-4.5 Team. (2025, August). GLM-4.5: Agentic, Reasoning, and Coding (ARC) Foundation Models. arXiv:2508.06471. https://arxiv.org/abs/2508.06471

[4] Shah, I., et al., with Vaswani, A. (2025, May). Practical Efficiency of Muon for Pretraining. Essential AI. arXiv:2505.02222. https://arxiv.org/abs/2505.02222

[5] Fireworks AI. (n.d.). Deep-dive into MuonClip: Fixing attention score explosions. https://fireworks.ai/blog/muonclip

[6] Ahn, K., Xu, B., Abreu, N., Fan, Y., Magakyan, G., Sharma, P., Zhan, Z., & Langford, J. (2025, April 7). Dion: Distributed Orthonormalized Updates. Microsoft Research. arXiv:2504.05295. https://arxiv.org/abs/2504.05295 · Code: https://github.com/microsoft/dion

[7] Bernstein, J., & Newhouse, L. (2024, September). Old Optimizer, New Norm: An Anthology. arXiv:2409.20325. https://arxiv.org/abs/2409.20325

[8] Bernstein, J. (2025, September 26). Modular Manifolds. Thinking Machines Lab. https://thinkingmachines.ai/blog/modular-manifolds/ · Code: https://github.com/thinking-machines-lab/manifolds

[9] Bernstein, J. (2025, March 7). Deriving Muon. https://jeremybernste.in/writing/deriving-muon

[10] Li, Z., Liu, L., Liang, C., Chen, W., & Zhao, T. (2025, October 7). NorMuon: Making Muon more efficient and scalable. arXiv:2510.05491. https://arxiv.org/abs/2510.05491

[11] Si, C., Zhang, D., & Shen, W. (2025, July 15). AdaMuon: Adaptive Muon Optimizer. arXiv:2507.11005. https://arxiv.org/abs/2507.11005

[12] Tveit, A., Remseth, B., & Skogvold, A. (2025, April 22). Muon Optimizer Accelerates Grokking. arXiv:2504.16041. https://arxiv.org/abs/2504.16041

[13] Amsel, N., Persson, D., Musco, C., & Gower, R. M. (2025, May 22). The Polar Express: Optimal Matrix Sign Methods and Their Application to the Muon Algorithm. arXiv:2505.16932. https://arxiv.org/abs/2505.16932