arXiv：cs.LG· Linkai Ma, Xinyu Luo, Mengbo Wang, Ananth Grama, Petros Drineas, Brian Bullins·· 3 小时前AI 评分38
MuonIO：面向嵌入表与语言模型输出头的规范化下降优化器
MuonIO: Principled Norm-Aware Descent for Embedding Tables and Language Model Heads
AI 导读
MuonIO 为嵌入表和语言模型输出头提供统一的 Muon 式更新规则，分别基于 1→2 和 2→∞ 算子范数，对嵌入表做列归一化、对输出头做行归一化。在 1B LLaMA 的 C4 预训练中，相比 Muon 将优化器状态内存降低 50%、更新 FLOPs 减少约 46%，同时改善验证困惑度。
正文
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.
| Subjects: | Machine Learning (cs.LG); Artificial Intelligence (cs.AI) |
| Cite as: | arXiv:2610.02705 [cs.LG] |
| (or arXiv:2610.02705v1 [cs.LG] for this version) | |
| https://doi.org/10.48550/arXiv.2610.02705 arXiv-issued DOI via DataCite (pending registration) |
Submission history
From: Linkai Ma [view email]
[v1]
Fri, 2 Oct 2026 02:40:51 UTC (95 KB)
来源：arXiv：cs.LG · arxiv.org