← 首页|学术|Spectral Allocation: Why Muon Outperforms Adam, and How to Improve Muon
cs.LG · 2608.25990 · 2026-08-26

Spectral Allocation: Why Muon Outperforms Adam, and How to Improve Muon

Xiaodong Wu, Wenyi Yu, Chao Zhang, Philip Woodland
优化器Muon训练效率
💬 对训练轨迹上真实checkpoint做频谱探测,发现动量缓冲的奇异值方向存在稳定的各向异性——"头部"方向在稳定性边缘只能用小步长、"体部"方向容忍度更高,这统一解释了Muon为何优于Adam;据此提出的SAMuon比调优过的Muon节省13-24%训练token。

🎯 背景

正交优化器如Muon能相对Adam大幅加速大语言模型预训练,但其机制尚未被完全理解。作者通过对Transformer损失landscape做样本外的频谱探测分析来研究这一问题。

🔬 方法

在真实训练轨迹的多个checkpoint处,将每个动量缓冲分解为其奇异方向,并在留出数据上估计每个方向上的损失最优步长。得到的频谱profile在不同batch和训练阶段间各向异性但稳定,且在不同优化器和模型规模间保持一致:一个运行在稳定性边缘的"波动头部"只支持小得多的步长,而容忍度更高的"体部"允许大得多的步长。据此,作者提出Spectral-Aware Muon(SAMuon),保持头部在Muon尺度、用静态频谱先验放大体部,提供完整版(低秩随机SVD)和简化版SAMuon-lite(基于rank-one幂迭代的两级近似)。

📊 结果

两个变体都不引入持久优化器状态或明显额外FLOPs;理想化精确白化版本在标准假设下保持Muon的渐近收敛率。在124M到1B参数的"modded-nanogpt"模型上,两个变体在所有评测的模型规模与batch size配置下都超过调优过的AdamW和Muon(Scion实现)基线,SAMuon比Muon少用13.3%-24.0%的训练token达到相同验证损失,SAMuon-lite以近零额外开销保留大部分增益。

原文摘要

▶ Abstract
Orthogonal optimisers such as Muon can substantially accelerate large language model pretraining relative to Adam, yet the mechanism remains incompletely understood. We investigate this through an out-of-sample spectral probing analysis of Transformer loss landscapes. At checkpoints along real training trajectories, we decompose each momentum buffer into its singular directions and estimate the loss-optimal step size along each direction on held-out data. The resulting spectral profile is anisotropic yet stable across batches and training stages, and consistent across the optimisers and model scales: a volatile head operating at the Edge-of-Stability supports a much smaller step size than the tolerant bulk, which permits substantially larger steps. This profile provides a unified spectral allocation account of why Muon outperforms Adam, which outperforms SGD. It also exposes a limitation of Muon's uniform scaling: it still underutilises the bulk. Guided by this finding, we introduce Spectral-Aware Muon (SAMuon), which holds the head at the Muon scale and amplifies the bulk using a static spectral prior. We provide two variants: the complete SAMuon follows the measured profile using a low-rank randomised SVD and the simplified SAMuon-lite uses a two-level approximation via rank-one power iteration. Neither method adds persistent optimiser state or notable extra FLOPs beyond Muon at scale, and the idealised exact-whitening versions of both retain Muon's asymptotic convergence rate under standard assumptions. Across "modded-nanogpt" models from 124M to 1B parameters, both variants outperform tuned AdamW and Muon (Scion implementation) baselines in all evaluated model-scale and batch-size configurations. SAMuon requires 13.3% to 24.0% fewer training tokens to reach the same validation loss as Muon, while SAMuon-lite retains most of this gain with near-zero wall-clock overhead.
来源: arXiv:2608.25990 · 精读由高松灯生成,基于摘要与 arXiv 页面信息