AMUSE: Anytime Muon with Stable Gradient Evaluation

📄 arXiv: 2605.22432v1 📥 PDF

作者: Jueun Kim, Baekrok Shin, Jihun Yun, Beomhan Baek, Minhak Song, Chulhee Yun

分类: cs.LG

发布日期: 2026-05-21

备注: 41 pages, 25 figures


💡 一句话要点

提出AMUSE优化器,结合Muon加速和Schedule-Free稳定,提升深度学习模型训练效率。

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

关键词: 优化器 深度学习 Muon Schedule-Free 自适应优化 学习率策略 梯度下降

📋 核心要点

  1. 现有优化器如AdamW依赖手动调整的学习率策略,且Muon虽加速训练,但易受损失函数景观中高曲率方向噪声影响。
  2. AMUSE通过时变插值系数,结合Muon的快速收敛和Schedule-Free的稳定性,自适应地平衡梯度更新,无需手动调整学习率。
  3. 实验表明,AMUSE在视觉任务和大型语言模型预训练中,相较于AdamW和Muon,在性能-迭代次数帕累托前沿上实现了显著提升。

📝 摘要(中文)

现代深度学习通常依赖于带有预设学习率策略的AdamW优化器,但最近的研究对这两个组成部分都提出了挑战:Schedule-Free优化通过迭代平均消除了显式的学习率策略,而Muon通过正交化矩阵参数的动量来改善更新的几何结构。尽管Muon具有强大的经验性能,但其底层机制仍未完全理解。我们通过河谷损失景观研究Muon,其中有用的训练进展发生在平坦、低曲率的主子空间(河)上,而高曲率的主导方向形成陡峭的谷壁,导致振荡。经验表明,虽然Muon的正交化通过增加主成分来加速河流进展,但它也放大了主导方向的噪声,导致振荡轨迹。在此基础上,我们提出了Anytime MUon with Stable gradient Evaluation (AMUSE),它将Muon的快速主成分进展与Schedule-Free平均的稳定效果相结合。AMUSE使用一个随时间变化的插值系数,该系数最初在快速Muon序列附近评估梯度以进行快速适应,然后逐渐转向稳定的平均序列以抑制谷壁振荡。因此,AMUSE不需要学习率策略,并支持随时训练。在视觉任务和大型语言模型预训练中,AMUSE始终优于(Schedule-Free) AdamW和Muon。

🔬 方法详解

问题定义:现有深度学习优化器,如AdamW,通常需要手动调整学习率策略,这增加了训练的复杂性。Muon虽然通过正交化动量加速了训练,但容易受到损失函数景观中高曲率方向噪声的影响,导致训练过程中的振荡。因此,如何设计一种既能快速收敛,又能稳定训练的优化器是一个关键问题。

核心思路:AMUSE的核心思路是将Muon的快速收敛能力与Schedule-Free优化器的稳定性相结合。通过一个时变的插值系数,AMUSE能够在训练初期利用Muon的快速梯度更新,并在训练后期转向Schedule-Free优化器的稳定平均梯度,从而在加速收敛的同时抑制振荡。

技术框架:AMUSE的整体框架基于Muon优化器,并引入了一个时变插值系数α(t)。该系数用于在Muon的梯度更新和Schedule-Free优化器的平均梯度之间进行插值。具体来说,AMUSE的梯度更新公式为:g_amuse(t) = α(t) * g_muon(t) + (1 - α(t)) * g_avg(t),其中g_muon(t)是Muon的梯度,g_avg(t)是Schedule-Free优化器的平均梯度。

关键创新:AMUSE的关键创新在于引入了时变插值系数α(t),该系数能够自适应地平衡Muon的快速收敛和Schedule-Free优化器的稳定性。通过在训练初期设置较大的α(t)值,AMUSE可以利用Muon的快速梯度更新加速收敛;随着训练的进行,逐渐减小α(t)值,使AMUSE逐渐转向Schedule-Free优化器的稳定平均梯度,从而抑制振荡。

关键设计:时变插值系数α(t)的设计是AMUSE的关键。论文中采用了一种简单的线性衰减策略,即α(t) = 1 - t/T,其中t是当前迭代次数,T是总迭代次数。此外,AMUSE还继承了Muon的正交化动量更新策略,以进一步加速训练。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明,AMUSE在视觉任务(如图像分类)和大型语言模型预训练中,均优于AdamW和Muon。具体而言,在相同的迭代次数下,AMUSE能够达到更高的精度;或者在达到相同精度的情况下,AMUSE所需的迭代次数更少。这表明AMUSE在性能-迭代次数帕累托前沿上实现了显著提升。

🎯 应用场景

AMUSE优化器可广泛应用于各种深度学习任务,尤其适用于需要快速训练和稳定收敛的场景,如计算机视觉中的图像分类、目标检测,自然语言处理中的语言模型预训练等。该方法降低了对学习率调参的需求,具有很高的实际应用价值,并有望推动深度学习模型在资源受限环境下的部署。

📄 摘要(原文)

Modern deep learning commonly relies on AdamW with prescribed learning rate schedules, but recent works challenge both components: Schedule-Free optimization removes explicit schedules via iterate averaging, and Muon improves the update geometry by orthogonalizing momentum for matrix parameters. Despite Muon's strong empirical performance, its underlying mechanism remains partially understood. We study Muon through the river-valley loss landscape, where useful training progress occurs along a flat, low-curvature bulk subspace (the river), while high-curvature dominant directions form steep valley walls that induce oscillations. We empirically show that while Muon's orthogonalization accelerates river progress by increasing the bulk component, it also amplifies dominant-direction noise, causing oscillatory trajectories. Building on this, we propose Anytime MUon with Stable gradient Evaluation (AMUSE), which integrates Muon's rapid bulk progress with the stabilizing effect of Schedule-Free averaging. AMUSE uses a time-varying interpolation coefficient that initially evaluates gradients near the fast Muon sequence for rapid adaptation, then gradually shifts toward the stable averaged sequence to suppress valley-wall oscillations. As a result, AMUSE requires no learning rate schedules and supports anytime training. Across vision tasks and large language model pretraining, AMUSE consistently improves the performance-iteration Pareto frontier over (Schedule-Free) AdamW and Muon.