arXiv：cs.LG（机器学习，全量分类）· Arash Lagzian, Paniz Halvachi, Junming Zhang, Zhouhan Lin, Dianbo Liu·· 15 小时前AI 评分36
AF-Muon：面向绑定嵌入模型的无 AdamW 的 Muon 优化器
AF-Muon: An AdamW-Free Muon Optimizer for Tied-Embedding Models
AI 导读
AF-Muon 是一种无需 AdamW 的 Muon 扩展优化器，对隐藏权重矩阵保留 Muon 矩阵更新，对绑定词表采用支持感知的有限上限线性最小化 oracle，对一维辅助参数使用 RMS 归一化更新，仅需单个一阶矩缓冲、无二阶矩状态，相比 Hybrid Muon 节省约 20% 优化器状态内存。
正文
Abstract: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.
| Comments: | 63 pages, 19 figures. An earlier, shorter version of this work was accepted as a poster at the OPT 2026 workshop (Optimization for Machine Learning) at NeurIPS 2026; this is the complete version |
| Subjects: | Machine Learning (cs.LG) |
| Cite as: | arXiv:2610.01395 [cs.LG] |
| (or arXiv:2610.01395v1 [cs.LG] for this version) | |
| https://doi.org/10.48550/arXiv.2610.01395 arXiv-issued DOI via DataCite (pending registration) |
Submission history
From: Arash Lagzian [view email]
[v1]
Thu, 1 Oct 2026 09:59:36 UTC (7,026 KB)
来源：arXiv：cs.LG（机器学习，全量分类） · arxiv.org