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

May 11, 2026

Muown: Row-Norm Control for Muon Optimization

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.

Kai Lion, Florian Hübler, Bingcong Li +2

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

Ask PDF

Oct 1, 2026

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

Muon improves large-scale training by applying a spectral-norm steepest-descent update to matrix parameters, but practical models also contain parameter blocks that do not fit dense-matrix geometry. One important case is the tied vocabulary table, which appears in language models and other token generators and can receive multiple structurally different gradient sources, from sparse input lookups to dense output-classifier updates. In the reference recipe these blocks are handed to an auxiliary AdamW optimizer, which restores second-moment state and updates the aliased table as a generic tensor. We propose AF-Muon, an AdamW-free extension of Muon that keeps the Muon matrix update for hidden weight matrices while using a support-aware finite-cap linear minimization oracle for tied vocabulary tables and an RMS-normalized update for one-dimensional auxiliary parameters. AF-Muon therefore trains every parameter class with a single first-moment buffer and no second-moment state, saving around 20% optimizer-state memory relative to Hybrid Muon in our benchmark. Across nine tied-token settings - decoder-only language models from 124M to 1B parameters, a fully shared T5-style encoder-decoder, and ImageGPT-style image-token, protein, and sparse-MoE variants, spanning text, image, and protein-sequence data - AF-Muon improves mean validation loss and perplexity over both Hybrid Muon and a SCION-style Sign endpoint. Long-horizon runs and hyperparameter sensitivity studies confirm the gain is robust, and identical-momentum diagnostics attribute it to the finite cap, which preserves more within-row magnitude than Sign while bounding the coordinate concentration of row-RMS. These results identify tied vocabulary tables as a distinct optimizer geometry and yield a robust AdamW-free Muon variant across models, modalities, and architectures, with about 1% step-time overhead in matched training.

Arash Lagzian, Paniz Halvachi, Junming Zhang +2

National University of Singapore, Singapore · Independent Researcher · LUMIA Lab, School of Artificial Intelligence, Shanghai Jiao Tong University, Shanghai, China

Memory-Efficient Optimization Muon Optimizer LLM Training

Ask PDF

Sep 28, 2026

MeqMuon: Matrix-Equilibrating Muon for LLM Pretraining

The success of large language models (LLMs) has been accompanied by continued growth in model size and pretraining costs. Muon offers high accuracy and training efficiency in LLM pretraining. Recent work introduces row-wise normalization into Muon to balance update magnitudes and improve pretraining performance. However, row-wise normalization alone cannot accommodate different imbalance patterns in update matrices. In this paper, we propose an improved Muon optimizer, called matrix-equilibrating Muon~(MeqMuon), for LLM pretraining. MeqMuon balances both row and column magnitudes through normalization that can be automatically tailored to different imbalance patterns without manual intervention. Moreover, MeqMuon eliminates the need to store AdamW's second-moment estimates, reducing optimizer-state memory usage. Empirical results demonstrate that MeqMuon achieves better convergence performance than AdamW, Muon, and other baselines in LLM pretraining.

Chang-Wei Shi, Xu Wang, Wu-Jun Li

National Key Laboratory for Novel Software Technology, School of Computer Science, Nanjing University, P. R. China

Language Model Pretraining Efficient Language Model Training LLM Training

Ask PDF

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