MuonIO: Principled Norm-Aware Descent for Embedding Tables and Language Model Heads

Linkai Ma; Xinyu Luo; Mengbo Wang; Ananth Grama; Petros Drineas; Brian Bullins

cs.LGOct 2, 2026

MuonIO: Principled Norm-Aware Descent for Embedding Tables and Language Model Heads

Authors: Linkai Ma, Xinyu Luo, Mengbo Wang, Ananth Grama, Petros Drineas, Brian Bullins

Organizations: Department of Computer Science, Purdue University

Abstract

The Muon optimizer derives its update rule for hidden linear layers by solving a local linearization of the loss penalized by the spectral norm, motivated by an RMS-stability argument for dense linear layers. Standard Muon implementations, however, exclude the input (embedding table) and output (language model head) layers from this principled treatment, for which they use AdamW instead. We present MuonIO, a single Muon-style update for both of these layers. For the language model head $\mathbf{L} \in \mathbb{R}^{V \times d}$ , we motivate the use of the $2\to\infty$ operator norm, due to the Lipschitz continuity of the softmax output geometry, while for the embedding table $\mathbf{E} \in \mathbb{R}^{d \times V}$ , we draw on the $1 \to 2$ operator norm, based on the one-hot input geometry identified by Bernstein & Newhouse (2025). The identity $\lVert\mathbf{L}\rVert_{2\to\infty}=\lVert\mathbf{L}^\top\rVert_{1\to2}$ then puts both matrices in the same vocabulary-oriented geometry: MuonIO applies a single normalized-vector rule, which appears as column normalization for $\mathbf{E}$ and row normalization for $\mathbf{L}$ . Empirical evaluations demonstrate the effectiveness of our approach, with MuonIO reducing I/O optimizer state memory by 50% and I/O update FLOPs by $\sim$ 46% compared to Muon for 1B LLaMA pretraining on C4, while also improving validation perplexity.

Figures & tables

Configuration	Embedding	Hidden 2D	LM Head	Gains & Biases
Plain Muon (baseline 1)	AdamW $\eta_{a}$	Muon $\eta_{m}$	AdamW $\eta_{a}$	AdamW $\eta_{a}$
Muon Tuned (baseline 2)	AdamW $\eta_{i}$	Muon $\eta_{m}$	AdamW $\eta_{o}$	AdamW $\eta_{a}$
Ember	Ember $\eta_{i}$	Muon $\eta_{m}$	Ember $\eta_{o}$	AdamW $\eta_{a}$
MuonIO (ours)	IO $\eta_{i}$	Muon $\eta_{m}$	IO $\eta_{o}$	AdamW $\eta_{a}$
Polar Express Tuned	AdamW $\eta_{i}$	Polar Express $\eta_{h}$	AdamW $\eta_{o}$	AdamW $\eta_{a}$
Polar Express IO (ours)	IO $\eta_{i}$	Polar Express $\eta_{h}$	IO $\eta_{o}$	AdamW $\eta_{a}$

Table 1: Optimizer configurations used in our comparisons. Each entry specifies the optimizer and learning rate for the corresponding parameter group. Repeated learning rate symbols indicate that the same learning rate is shared across parameter groups. For each configuration, all distinct learning rates shown in the corresponding row are jointly tuned.

Configuration	60M	130M	1B
Plain Muon (baseline 1)	$36.963\pm 0.054$	$28.513\pm 0.026$	$15.116\pm 0.0003$
Muon Tuned (baseline 2)	$28.502\pm 0.015$	$24.238\pm 0.008$	$14.527\pm 0.007$
Ember	$28.835\pm 0.043$	$24.074\pm 0.004$	$14.516\pm 0.032$
MuonIO (ours)	$\mathbf{28.366}\pm 0.025$	$\mathbf{23.617}\pm 0.025$	$\mathbf{14.374}\pm 0.132$
Polar Express Tuned	$\mathbf{28.385}\pm 0.009$	$24.055\pm 0.010$	$14.361\pm 0.002$
Polar Express IO (ours)	$28.430\pm 0.025$	$\mathbf{23.646}\pm 0.012$	$\mathbf{14.160}\pm 0.008$

Table 2: Final validation perplexity on C4 (lower is better). The upper block uses standard Muon hidden updates, while the lower block uses Polar Express. Bold values indicate the lowest mean perplexity within each block at each model scale. For 60M and 130M, we report the mean and sample standard deviation over the three of five seeds with the lowest final validation perplexity. For 1B, we report the mean and sample standard deviation over three seeds.

Figure 1: Validation loss trajectories on C4 at the 130M scale for Plain Muon (baseline 1), Muon Tuned (baseline 2), Ember and MuonIO. All four configurations use standard Muon hidden updates. The inset highlights the 0.7–0.9B training token range.

I/O-layer optimizer	FLOPs	Optimizer State
AdamW	$13Vd$	$2Vd$ entries
IO	$7Vd+3V$	$Vd$ entries

Table 3: FLOPs and persistent optimizer state counts for one

V\times d

I/O matrix and one optimizer update. Derivations are in Appendix C .

Figure 2: Gradient magnitudes during 130M model training. Rows show the input embedding table (top) and language model head (bottom); columns show the Frobenius norm, spectral norm, and maximum absolute gradient entry. Curves correspond to MuonIO (orange, solid), Ember (blue, dashed), and Muon Tuned (green, dash-dotted). The horizontal axis counts microbatch steps, with eight microbatches per optimizer update. Each curve represents one run using the best set of hyperparamters of the corresponding optimizer.

Appendix figures & tables3 assets

Supplementary material from the paper’s appendix.

Appendix

Stage	Tensor operation	FLOPs
First moment	$\mathbf{M}_{k}\leftarrow\frac{\beta_{1}}{1-\beta_{1}^{t}}\mathbf{M}_{k-1}+\frac{1-\beta_{1}}{1-\beta_{1}^{t}}\mathbf{G}_{k}$	$3N$
Second moment	$\mathbf{V}_{k}\leftarrow\frac{\beta_{2}}{1-\beta_{2}^{t}}\mathbf{V}_{k-1}+\frac{1-\beta_{2}}{1-\beta_{2}^{t}}\mathbf{G}_{k}\odot\mathbf{G}_{k}$	$4N$
Corrected denominator	$\mathbf{D}_{k}\leftarrow\sqrt{\mathbf{V}_{k}}+\varepsilon$	$2N$
Decay and adaptive step	$\mathbf{W}_{k+1}\leftarrow(1-\gamma\lambda)\mathbf{W}_{k}-\gamma(\mathbf{M}_{k}\oslash\mathbf{D}_{k})$	$4N$
Total		$13N=13Vd$

Appendix

Table 4: AdamW tensor-arithmetic count for one matrix and one optimizer update. The elementwise product and division are denoted by

\odot

and

\oslash

, respectively.

Stage	Operation	FLOPs
Momentum	$\mathbf{M}_{k+1}\leftarrow\beta\mathbf{M}_{k}+(1-\beta)\tilde{\mathbf{G}}_{k}$	$3N$
Row norms	$r_{j}\leftarrow\\|(\mathbf{M}_{k+1})_{j,:}\\|_{2}$ , $j\in[V]$	$V[d+(d-1)+1]=2N$
Sum norms	$s_{k+1}\leftarrow\sum_{j=1}^{V}r_{j}$	$V-1$
Form scalar scale	$a_{k+1}\leftarrow(\lambda\sqrt{d})\,s_{k+1}$	$1$
Form row scales	$c_{j}\leftarrow a_{k+1}/\max\{r_{j},\varepsilon\}$ , $j\in[V]$	$2V$
Fused parameter update	$(\mathbf{W}_{k+1})_{j,:}\leftarrow(\mathbf{W}_{k})_{j,:}-c_{j}(\mathbf{M}_{k+1})_{j,:}$ , $j\in[V]$	$2N$

Appendix

Table 5: MuonIO tensor-arithmetic count for one matrix and one iteration of Algorithm 1 . The constant

\lambda\sqrt{d}

is computed once, before training.

Method	Size	$\eta_{h}$	$\eta_{i}$	$\eta_{o}$	$\eta_{a}$	$\mu_{h}$	$\beta_{1}^{A}$	$\beta_{2}^{A}$	$\beta_{\mathrm{I/O}}$
Plain Muon	60M	0.010	$3\times 10^{-4}$	$3\times 10^{-4}$	0.001	0.950	0.900	0.950	–
	130M	0.010	$3\times 10^{-4}$	$3\times 10^{-4}$	0.001	0.950	0.900	0.950	–
	1B	0.019	$3\times 10^{-4}$	$3\times 10^{-4}$	0.023	0.950	0.950	0.950	–
Muon Tuned	60M	0.020	0.040	0.020	0.020	0.900	0.900	0.950	–
	130M	0.020	0.040	0.020	0.020	0.900	0.900	0.950	–
	1B	0.016	0.197	0.004	0.015	0.900	0.900	0.950	–

Appendix

Table 6: Optimizer hyperparameters for the 18 selected LLaMA pre-training runs on C4. Values use three decimal places; nonzero values below

10^{-3}

use scientific notation.

\mu_{h}

is hidden weight momentum;

(\beta_{1}^{A},\beta_{2}^{A})

are Adam/AdamW coefficients;

\beta_{\mathrm{I/O}}

is the separate coefficient for non-Adam I/O updates. PE denotes Polar Express.

Muon Optimizer

Muon optimizer uses momentum and matrix orthogonalization to shape neural-network updates, especially in large language-model training.

See all 146 papers

Muown: Row-Norm Control for Muon Optimization
May 11, 2026Kai Lion, Florian Hübler, Bingcong Li +2Language Model Pretraining Muon Optimizer
Department of Computer Science, ETH Zurich, Switzerland · Department of Mathematics, School of Computation, Information and Technology; Technical University of Munich, Germany · ELLIS Institute Tübingen, MPI-IS, Tübingen AI Center, Germany
AF-Muon: An AdamW-Free Muon Optimizer for Tied-Embedding Models
Oct 1, 2026Arash Lagzian, Paniz Halvachi, Junming Zhang +2Memory-Efficient Optimization Muon Optimizer
National University of Singapore, Singapore · Independent Researcher · LUMIA Lab, School of Artificial Intelligence, Shanghai Jiao Tong University, Shanghai, China
MeqMuon: Matrix-Equilibrating Muon for LLM Pretraining
Sep 28, 2026Chang-Wei Shi, Xu Wang, Wu-Jun LiLanguage Model Pretraining Efficient Language Model Training
National Key Laboratory for Novel Software Technology, School of Computer Science, Nanjing University, P. R. China

May 11, 2026J/K move · Enter open · S save

Muown: Row-Norm Control for Muon Optimization

Kai Lion, Florian Hübler, Bingcong Li, Antonio Orvieto+1

Department of Computer Science, ETH Zurich, Switzerland · Department of Mathematics, School of Computation, Information and Technology; Technical University of Munich, Germany · ELLIS Institute Tübingen, MPI-IS, Tübingen AI Center, Germany

Language Model Pretraining Muon Optimizer

Open paper AskPDF arXiv

Muon has emerged as a strong competitor to AdamW for language model pre-training, yet its behavior at scale is sensitive to weight decay. Recent work has observed that, for Muon without decoupled weight decay, the spectral norm of weight matrices drifts upward over training. Through a decomposition of the spectral norm into a row-magnitude factor and a row-coherence factor, we identify the former as the empirical driver of this drift under Muon, while the latter remains well-behaved along the trajectory. Motivated by this diagnosis, we introduce Muown, a drop-in replacement for Muon that treats the row-magnitude vector as an explicit optimizer variable, updating it under the

\ell_\infty

geometry induced by the decomposition, while applying Muon unchanged to the remaining direction component. We prove that Muown attains the optimal non-convex rates in both deterministic and stochastic regimes under a dual norm aligned with the underlying geometries and with a stochastic noise coefficient that empirically remains below that of Muon throughout training. Across GPT-style pre-training on FineWeb-Edu with model sizes from 124M up to 2.7B parameters, Muown improves perplexity over Muon, SOAP, AdamW, and Lion. It also widens the plateau of near-optimal learning rates across model scales, reduces sensitivity to weight decay, and avoids the spectral norm drift at negligible step-time overhead when appropriately sharded.

Abstract

Muon Optimizer

Explore similar work

Muown: Row-Norm Control for Muon Optimization

AF-Muon: An AdamW-Free Muon Optimizer for Tied-Embedding Models

MeqMuon: Matrix-Equilibrating Muon for LLM Pretraining