arXiv:2608.25990cs.LG2026-08

通过谱分析揭示穆翁优化器加速原理,并提出更优的谱感知改进版本。

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

论文配图:Spectral Allocation: Why Muon Outperforms Adam, and How to Improve Muon
图 1 · 摘自论文原文
  • 基于损失景观谱分解,发现优化方向存在差异步长需求。
  • 新方法在124M到10亿参数模型上减少13.3%~24.0%训练令牌数。
  • 无需额外状态或算力,轻量版可近乎零延迟提升性能。

正交优化器如穆翁(Muon)相比Adam能显著加速大语言模型预训练,但其机制仍不明确。本文通过在真实训练轨迹上对Transformer损失景观进行样本外谱探测,将动量缓冲区分解为奇异方向,并在保留数据上估计各方向的损失最优步长。结果表明,谱分布具有各向异性但跨批次和训练阶段稳定,且在不同优化器与模型规模间一致:处于稳定边缘的敏感头部需小步长,而容错性强的主体部分允许更大步长。该谱分配模式统一解释了为何穆翁优于Adam,Adam优于SGD。同时暴露了穆翁均匀缩放的局限:仍未充分利用主体部分。据此提出谱感知穆翁(SAMuon),在保持头部按穆翁尺度的同时,利用静态谱先验放大主体部分。提供两种变体:完整版采用低秩随机SVD拟合测量谱,简化版使用单次幂迭代的双层近似。两者均不引入持久优化状态或显著额外浮点运算,在大规模下保持穆翁的渐近收敛率。在124M至10亿参数的modded-nanogpt模型上,所有配置下均优于调优的AdamW与穆翁(Scion实现)基线。SAMuon比穆翁减少13.3%至24.0%训练令牌数达到相同验证损失,而SAMuon-lite仅用近零时钟开销保留大部分收益。

原文摘要 · Abstract (English)

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.

优化器谱分析大模型训练加速

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。