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

arXiv: https://arxiv.org/abs/2610.02705 · HTML: https://mlsanity.com/paper/2610.02705/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 23 of its 26 references that are in the map.
- First posted on arXiv 2 Oct 2026.

## Builds on (most cited first, 23 of 23)

- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [Kimi K2: Open Agentic Intelligence](https://mlsanity.com/paper/2507.20534/) (2025)
- [GaLore: Memory-Efficient LLM Training by Gradient Low-Rank Projection](https://mlsanity.com/paper/2403.03507/) (2024)
- [Muon is Scalable for LLM Training](https://mlsanity.com/paper/2502.16982/) (2025)
- [Using the Output Embedding to Improve Language Models](https://mlsanity.com/paper/1608.05859/) (2016)
- [Old Optimizer, New Norm: An Anthology](https://mlsanity.com/paper/2409.20325/) (2024)
- [Training Deep Learning Models with Norm-Constrained LMOs](https://mlsanity.com/paper/2502.07529/) (2025)
- [The Polar Express: Optimal Matrix Sign Methods and Their Application to the Muon Algorithm](https://mlsanity.com/paper/2505.16932/) (2025)
- [A Spectral Condition for Feature Learning](https://mlsanity.com/paper/2310.17813/) (2023)
- [Tying Word Vectors and Word Classifiers: A Loss Framework for Language Modeling](https://mlsanity.com/paper/1611.01462/) (2016)
- [NorMuon: Making Muon more efficient and scalable](https://mlsanity.com/paper/2510.05491/) (2025)
- [Modular Duality in Deep Learning](https://mlsanity.com/paper/2410.21265/) (2024)
- [Dion: Distributed Orthonormalized Updates](https://mlsanity.com/paper/2504.05295/) (2025)
- [AdaMuon: Adaptive Muon Optimizer](https://mlsanity.com/paper/2507.11005/) (2025)
- [Softmax is 1/2-Lipschitz: A tight bound across all ℓp norms](https://mlsanity.com/paper/2510.23012/) (2025)
- [Memory-Efficient LLM Pretraining via Minimalist Optimizer Design](https://mlsanity.com/paper/2506.16659/) (2025)
- [Dion2: A Simple Method to Shrink Matrix in Muon](https://mlsanity.com/paper/2512.16928/) (2025)
- [HTMuon: Improving Muon via Heavy-Tailed Spectral Correction](https://mlsanity.com/paper/2603.10067/) (2026)
- [Symmetry-Compatible Principle for Optimizer Design: Embeddings, LM Heads, SwiGLU MLPs, and MoE Routers](https://mlsanity.com/paper/2605.18106/) (2026)
- [Optimal Embedding Learning Rate in LLMs: The Effect of Vocabulary Size](https://mlsanity.com/paper/2506.15025/) (2025)
- [One LR Doesn't Fit All: Heavy-Tail Guided Layerwise Learning Rates for LLMs](https://mlsanity.com/paper/2605.22297/) (2026)
- [Quantifying Hyperparameter Transfer and the Importance of Embedding Layer Learning Rate](https://mlsanity.com/paper/2605.21486/) (2026)
- [Ember: A Kilobyte-State Optimizer for Token Interfaces](https://mlsanity.com/paper/2607.01455/) (2026)

## Abstract (arXiv)

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 L ∈ ℝV × d, we motivate the use of the 2→∞ operator norm, due to the Lipschitz continuity of the softmax output geometry, while for the embedding table E ∈ ℝd × V, we draw on the 1 → 2 operator norm, based on the one-hot input geometry identified by Bernstein & Newhouse (2025). The identity lVertLrVert2→∞=lVertL^toprVert1→2 then puts both matrices in the same vocabulary-oriented geometry: MuonIO applies a single normalized-vector rule, which appears as column normalization for E and row normalization for 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 ~46% compared to Muon for 1B LLaMA pretraining on C4, while also improving validation perplexity.