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

📄 arXiv: 2608.25990v1 📥 PDF

作者: Xiaodong Wu, Wenyi Yu, Chao Zhang, Philip Woodland

分类: cs.LG

发布日期: 2026-08-26

备注: 34 pages, 13 figures, 7 tables


💡 一句话要点

提出Spectral-Aware Muon以提升大语言模型训练效率

🎯 匹配领域: 支柱九:具身大模型 (Embodied Foundation Models)

关键词: 优化器 大语言模型 谱分析 训练效率 机器学习 深度学习 Transformer

📋 核心要点

  1. 现有优化方法如Adam在大语言模型训练中效率不足,尤其在处理大规模数据时表现不佳。
  2. 论文提出Spectral-Aware Muon(SAMuon),通过保持头部规模并利用静态谱先验来优化步长分配,从而提升训练效率。
  3. 实验结果表明,SAMuon在不同模型规模和批量配置下均优于调优后的AdamW和Muon,减少了训练token需求。

📝 摘要(中文)

正交优化器如Muon在大语言模型预训练中相较于Adam显著加速,但其机制尚不完全清晰。本文通过对Transformer损失景观的谱探测分析,揭示了Muon优于Adam的原因,并提出了Spectral-Aware Muon(SAMuon)以改善Muon的均匀缩放限制。SAMuon在保持头部规模的同时,利用静态谱先验放大主体,实验结果显示其在多个模型规模和批量配置中均优于调优后的AdamW和Muon基线,减少了13.3%至24.0%的训练token以达到相同的验证损失。

🔬 方法详解

问题定义:本文旨在解决现有优化器在大语言模型预训练中效率不足的问题,尤其是Adam在处理大规模数据时的局限性。现有方法未能充分利用模型的不同部分,导致训练效率低下。

核心思路:论文的核心思路是通过谱分析来理解优化器的表现,提出SAMuon以优化步长分配,保持头部规模并放大主体,从而提高训练效率。

技术框架:整体架构包括对每个动量缓冲区进行奇异值分解,分析损失最优步长的谱特征,并基于此设计SAMuon。SAMuon有两个变体:完整的SAMuon和简化的SAMuon-lite,分别采用低秩随机SVD和级联的秩一幂迭代。

关键创新:最重要的技术创新在于提出了谱分配的统一框架,解释了Muon为何优于Adam,并通过SAMuon解决了Muon的均匀缩放限制,显著提升了训练效率。

关键设计:SAMuon的设计包括保持头部的Muonscale,利用静态谱先验来放大主体,且两种变体在计算复杂度上与Muon相当,确保了收敛速率的保持。

🖼️ 关键图片

fig_0
img_1
img_2

📊 实验亮点

实验结果显示,SAMuon在124M到1B参数的“modded-nanogpt”模型中,均优于调优后的AdamW和Muon(Scion实现)基线。在所有评估的模型规模和批量配置中,SAMuon减少了13.3%至24.0%的训练token以达到相同的验证损失,且SAMuon-lite几乎没有额外的时间开销。

🎯 应用场景

该研究的潜在应用领域包括自然语言处理、机器翻译和对话系统等大语言模型的训练。通过提高训练效率,SAMuon能够降低计算资源消耗,缩短模型训练时间,推动更大规模模型的开发与应用,具有重要的实际价值和未来影响。

📄 摘要(原文)

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.