Beyond MuP 3: Special Cases, Special Treatment (Embedding, LM Head, RMS Norm)

TL;DR

Part 3 of Jianlin Su’s “Beyond MuP” series on kexue.fm. The post derives stability metrics and steepest-descent directions for three layer types that the standard MuP/Muon machinery handles awkwardly: the input Embedding, the output LM Head, and RMS Norm. Su’s headline observation (from his tweet) is that Embedding and LM Head don’t play well with Muon — Muon’s Newton–Schulz orthogonalization is derived from a steepest-descent argument under the spectral norm, and that derivation breaks down (or yields the wrong update direction) when the layer is fundamentally a lookup table (Embedding) or a wide projection back to the vocabulary (LM Head) rather than a dense hidden-dim-to-hidden-dim matmul. The post works out what the right per-layer treatment is. Note: the original kexue.fm article is in Chinese and was not directly accessible at filing time (403); this page summarizes the tweet abstract and the article title only. An English translation is expected at main-horse.github.io/translations/kexue/ but had not been published as of 2026-05-23.

Key claims

• The post derives stability metrics (forward-pass scale invariance criteria) and steepest-descent directions specifically for the Embedding layer, the LM Head layer, and RMS Norm, rather than treating them as instances of the generic dense-matmul case [tweet abstract].
• Embedding and LM Head are flagged as layers that “don’t play well with Muon” — i.e. Muon’s orthogonalized update is not the right steepest-descent direction under the relevant spectral metric for these layer types [tweet abstract].
• The treatment of RMS Norm parallels the treatment of Embedding and LM Head — i.e. it is also a special case that the generic dense-layer MuP/Muon derivation does not cover [tweet abstract; inferred from the article being grouped under “Special Cases, Special Treatment”].

Method

(Article body not directly accessible at filing time; description below is inferred from the tweet abstract and the prior installments in Su’s “Beyond MuP” series — Muon + Spectral Sphere, the higher-order MuP post, the Stiefel/Orthogonal Muon variants, and the RMS-update derivations.)

The general recipe Su uses across this series is: (1) write down a stability target — typically that the per-token activation RMS at each layer is $\Theta(1)$ under both forward propagation and parameter updates — (2) ask what norm on the parameter update makes that target invariant under reasonable scaling assumptions, and (3) compute the steepest-descent direction under that norm, which defines the optimizer step for that layer. For dense hidden-dim matrices that recipe lands on the spectral norm and yields Muon’s Newton–Schulz orthogonalization. For Embedding (whose left dimension is vocabulary, not hidden width) and for LM Head (the transpose), the right metric is plausibly different — likely a per-row / per-column norm rather than a spectral norm — and so the corresponding optimizer step is something other than Newton–Schulz orthogonalization. For RMS Norm, the parameter is a per-channel scale vector, not a matrix, so spectral-norm steepest descent doesn’t apply at all and the per- element treatment is needed. The post is expected to make all three derivations explicit.

Results

Theoretical/derivational post; no empirical results in the traditional sense are claimed in the tweet abstract.

Why it’s interesting

This sits directly inside the optimizer-side fork of the training-stability debate the wiki has been tracking. The “Muon is half-aligned with µP” argument from Controlled LLM Training on Spectral Sphere is general across dense layers; Su’s piece sharpens the question of whether Muon is even defined for the right object on Embedding and LM Head, which is a different (and more architectural) critique. It also extends Su’s own line on per-layer-type optimizer derivation — same recipe that produced MuonClip / qk-clip and the Quantile Balancing router in Quantile Balancing: A Hyperparameter-Free MoE Load Balancing Method, applied to the embedding/unembedding edges of the network instead of attention logits or MoE routing. The practical question for Luma is whether anyone training large models with Muon today is silently mis-updating the two largest parameter tensors in the model.

See also

• Quantile Balancing: A Hyperparameter-Free MoE Load Balancing Method — companion post from the same author, applying the same “derive the right optimizer for the layer type” methodology to MoE routing
• Controlled LLM Training on Spectral Sphere — SSO; the “Muon is half-aligned with µP” argument that this post sharpens for Embedding / LM Head specifically
• How to Set the Learning Rate for Large-Scale Pre-training? — the per-module-LR debate; per-layer-type optimizer choice is the architectural analogue
• Learning Rate Scaling across LoRA Ranks and Transfer to Full Finetuning — µA, the LoRA-side per-factor analogue (different update rule for $A$ vs $B$)
• Training stability at scale — broader concept this post contributes to
• Hyperparameter scaling laws — Embedding/LM Head treatment affects µP transfer guarantees
• Translation tracker: https://main-horse.github.io/translations/kexue/ (English translations of Su’s kexue.fm posts; this post not yet translated as of 2026-05-23)