AMUSE: Anytime Muon with Stable Gradient Evaluation¶
会议: NeurIPS 2026
arXiv: 2605.22432
代码: https://github.com/kjeiun/amuse
领域: 优化/理论
关键词: 优化算法, Muon, 无调度优化, 河谷损失地形, 随时训练
一句话总结¶
本文从河谷损失地形几何视角揭示了 Muon 正交化在放大有益低曲率主体更新的同时也会放大高曲率优势方向噪声的机理，并提出结合时变插值系数与无调度迭代平均的 AMUSE 优化器，无需显式学习率衰减即可实现稳定的随时训练与帕累托前沿提升。
研究背景与动机¶
深度学习模型的优化长期依赖 AdamW 搭配预设的学习率衰减调度（如 Cosine 退火与 WSD 策略）。近期优化器研究从两个互补维度对该范式提出了挑战：一方面，Schedule-Free（无调度）优化通过引入迭代点插值与参数平均消除了显式学习率调度的需要，并支持随时终止训练（Anytime Training）；另一方面，Muon 优化器利用矩阵参数的几何结构，对动量矩阵进行正交化处理（通过 Newton-Schulz 迭代近似），在大型语言模型与视觉模型预训练中展现出远超传统一阶优化器的收敛速度，并被快速集成进多款前沿基础模型的训练流程中。
然而，Muon 取得显著经验增益背后的根本机制仍未被完全解释。从深度神经网络的高维非各向同性损失曲面出发，Hessian 谱通常由少数具有极大特征值的低维优势子空间（高曲率山谷峭壁）和大量相对平坦的高维主体子空间（低曲率河床）构成，即所谓的“河谷地形”（River-Valley Landscape）。在这一几何视角下，有益的长期训练进展主要沿平缓的河道发生，而优势方向的分量则导致跨越峭壁两侧的剧烈振荡。实证分析表明，Muon 的矩阵正交化虽然能够大幅提升沿河床主体方向的步长比例，但由于极性分解等正交化操作无法自适应区分噪声与信号，极易将高曲率方向上的微弱残差放大为恒定幅度的跨步，诱发严重的“撞墙反弹”振荡现象。
若直接将 Muon 作为基优化器引入传统固定插值系数的 Schedule-Free 框架，又会面临“早期快速适应”与“后期稳健防振”的鱼与熊掌困境：固定较小的插值系数保留了基序列的快速探索能力，却在训练后期无法抑制山谷振荡；固定较大的插值系数虽能滤除峭壁噪声，但在训练初期会导致平均序列严重滞后，拖慢沿河床的初始切入。核心 idea：将 Muon 的高强度主体方向推进与 Schedule-Free 沿河床求导的稳定机制深度耦合，设计一种独立于总训练步数、时变递增的插值调度函数，使求导点从早期贴近快速探索的 Muon 基序列平滑过渡到后期贴近低曲率的平均序列，在彻底移除学习率衰减的同时实现全局稳定的随时训练。
方法详解¶
整体框架¶
AMUSE 在维护快速基序列与稳定推理序列双轨轨迹的同时，通过动态调节梯度评估点实现训练全周期的稳定性与高收敛速率。输入为模型参数与批量训练数据，优化器内部针对矩阵层与非矩阵层（如偏置、归一化、分类头）分别采用带阻尼正交化和对应的一阶无调度更新规则，最终对外输出纯净平滑的平均参数序列。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["训练当前步 t & 状态参数 (Z_t, X_t, M_t-1)"] --> B["动态时变插值调度<br/>计算评估权重 beta_t"]
B --> C["河床近邻梯度评估<br/>在插值点 Y_t 计算梯度 G_t"]
C --> D["双动量解耦更新<br/>矩阵动量积累 M_t + Newton-Schulz 正交化"]
D --> E["双轨基序列与平均序列演进<br/>基序列 Z_t+1 与无调度遍历平均 X_t+1"]
E --> F["随时推理与评估模型 X_t+1"]关键设计¶
1. 动态时变插值调度：平衡早期探索与后期收敛窗口 为了解决固定插值参数 \(\beta\) 带来的早期迟滞或后期振荡困境，AMUSE 引入了随时间自适应增长的非线性插值系数 \(\beta_t\)。在预热阶段 \(t \le T_0\) 内，系统保持 \(\beta_t = \beta_{\text{init}}\) 以保证对初始参数的快速偏置修正；当 \(t > T_0\) 时，\(\beta_t\) 遵循基于隐式平均窗口保真推导的幂律增长函数： $$ \beta_t = 1 - \left(\frac{c_{t+1}(1 - c_{T_0+1})}{c_{T_0+1}(1 - c_{t+1})}\right)^\rho (1 - \beta_{\text{init}}) = 1 - \left(\frac{T_0}{t}\right)^\rho (1 - \beta_{\text{init}}) $$ 其中 \(c_{t+1} = \frac{\eta_t^2}{\sum_{i=1}^t \eta_i^2}\)（预热结束后退化为 \(\frac{1}{t+1}\)），\(\rho \in [0, 1]\) 是控制向平均序列收敛速率的超参数。该调度的核心洞察在于防止隐式平均有效记忆窗口坍塌：数学推导证明，当 \(\rho=1\) 时严格锁定平均窗口大小保持恒定，当 \(\rho \in (0, 1)\) 时允许有效记忆窗口实现可控的次线性扩展。这种形式完全独立于训练总轮次 \(T\)，随着 \(t\) 渐进趋向 1，在不引入终点约束的前提下自然赋予模型卓越的随时训练特性。
2. 河床近邻梯度评估：在正交化前抑制优势子空间噪声 Muon 振荡的本质在于极性归一化 \(\text{Polar}(\bm{G}) = \bm{G}(\bm{G}^\top \bm{G})^{-1/2}\) 消除了梯度的幅值信息，使极微弱的高曲率残差也能获得固定步长的跨步冲量。AMUSE 通过将梯度评估点定位在插值序列： $$ \bm{Y}_t = (1 - \beta_t)\bm{Z}_t + \beta_t \bm{X}_t $$ 直接从源头上利用平均序列 \(\bm{X}_t\) 贴近“平坦河床”的几何特性，过滤来自陡峭山壁的垂直方向高曲率扰动。随着 \(\beta_t\) 的持续提升，梯度 \(\bm{G}_t = \nabla \mathcal{L}(\bm{Y}_t)\) 中的优势方向投影分量 \(r^{\text{dom}}(\bm{G}_t; \bm{Y}_t)\) 显著单调下降，从而在将动量送入正交化算子之前便消除高频反弹噪声，确保正交化后的更新向量高度对齐低曲率的主体河道方向。
3. 双动量解耦更新：正交化前梯度平滑与正交化后轨迹平均 传统无调度优化器依靠平均序列位移 \(\Delta \bm{X}_t\) 即可诱发隐式动量效应，因而往往省略内部一阶动量缓冲区。然而在 AMUSE 中，矩阵正交化前后的动量扮演着完全非对称的数学角色：正交化前的显式动量 \(\bm{M}_t = \mu \bm{M}_{t-1} + (1-\mu)\bm{G}_t\) 负责平滑随机小批量采样噪声、沉淀真实的物理更新方向；而正交化后的无调度平均序列 \(\bm{X}_{t+1} = (1 - c_{t+1})\bm{X}_t + c_{t+1}\bm{Z}_{t+1}\) 则构成了已正交化位移的历史加权平均 \(\Delta \bm{X}_t = -\frac{\eta}{t(t+1)}\sum_{j=1}^t j \mathcal{T}(\bm{M}_j)\)。实验消融证实，若直接舍弃正交化前的 \(\bm{M}_t\)，正交化算子将极度放大单步 mini-batch 噪声，导致严重性能坍塌；AMUSE 坚决保留该显式缓冲区，构建起“前向降噪-正交映射-后向平滑”的三层保护。
4. 混合参数适配机制与零冗余内存优化 为适配现代深度神经网络的异构层结构，AMUSE 实行差异化分组更新策略：针对二维及以上隐藏层权重矩阵，采用 Newton-Schulz 快速迭代计算正交更新 \(\bm{O}_t = s_{\text{NS}} \cdot \text{Newton-Schulz}(\bm{M}_t)\)，配合解耦权重衰减更新基序列 \(\bm{Z}_{t+1} = (1 - \eta_t \lambda)\bm{Z}_t - \eta_t \bm{O}_t\)；对于一维偏置、嵌入层和归一化参数，则自适应降级为无一阶动量的 Schedule-Free AdamW 或 Schedule-Free SGD 规则。在工程实现上，虽然数学形式依赖 \((\bm{X}_t, \bm{Y}_t, \bm{Z}_t, \bm{M}_t)\) 四组变量，但训练常态下仅需在显存中保留物理参数状态 \(\bm{Y}_t\)、基状态 \(\bm{Z}_t\) 和矩阵动量 \(\bm{M}_t\)。评估或保存检查点时，通过原地线性变换直接从 \(\bm{Y}_t\) 与 \(\bm{Z}_t\) 恢复 \(\bm{X}_t\) 并在评估后复原，使显存占用丝毫不增加于普通 AdamW/SF-AdamW，甚至因非矩阵层省略一阶动量而略低于传统 AdamW。
损失函数 / 训练策略¶
AMUSE 严格遵循无衰减优化范式。学习率仅需经历标准的初始线性预热阶段： $$ \eta_t = \eta \cdot \min\left(1, \frac{t}{T_0}\right) $$ 在预热结束后，全局学习率 \(\eta\) 恒定不变，无需预设总训练步数 \(T\)，也无需维护任何阶梯、余弦或多项式衰减调度器。权重衰减直接以解耦形式作用于基参数序列 \(\bm{Z}_t\) 上。
实验关键数据¶
主实验¶
在多样化的视觉分类、分割以及覆盖 124M 至 1.3B 参数的大语言模型预训练任务中，AMUSE 在整个训练历程中全面超越带余弦调度的 SGD、AdamW 以及原生 Muon，取得更优的性能-迭代步数帕累托前沿。
| 任务 / 数据集 | 评价指标 | AMUSE (本文) | Muon (Cosine 衰减) | Baseline (AdamW / SGD) | 相对迭代步数加速比 |
|---|---|---|---|---|---|
| Llama-124M (FineWeb, 16k iters) | 验证集困惑度 (PPL) \(\downarrow\) | 28.61 | 29.21 | 30.82 (AdamW) | 1.48× |
| Llama-720M (FineWeb, 16k iters) | 验证集困惑度 (PPL) \(\downarrow\) | 19.42 | 19.95 | 20.68 (AdamW) | 1.51× |
| ImageNet-1k (ResNet-50, 100 epochs) | Top-1 准确率 (%) \(\uparrow\) | 77.17 | 76.97 | 76.24 (SGD) | 1.34× |
| CIFAR-10 (Wide-ResNet-16-8, 300 epochs) | 测试准确率 (%) \(\uparrow\) | 96.42 | 95.83 | 95.96 (SGD) | 1.25× |
| CIFAR-100 (DenseNet, 300 epochs) | 测试准确率 (%) \(\uparrow\) | 81.56 | 80.82 | 80.12 (SGD) | 1.22× |
| SVHN (ResNet-3-96, 300 epochs) | 测试准确率 (%) \(\uparrow\) | 97.88 | 97.45 | 97.51 (SGD) | 1.30× |
| ISIC 2018 (U-Net 分割, 50 epochs) | 验证集 Dice 相似度 (%) \(\uparrow\) | 89.51 | 88.63 | 87.94 (SGD) | 1.40× |
| ImageNet MAE 微调 (ViT-B/16, 100 ep) | Top-1 准确率 (%) \(\uparrow\) | 83.62 | 81.25 | 83.21 (AdamW) | 2.57× |
在更广泛的优化器基准横向对比中，AMUSE 亦一致优于 AdEMAMix、ADOPT、Signum、Prodigy、MARS、Lion、Sophia 和 SOAP 等前沿算法。特别是在 MAE 微调实验中，从 AdamW 预训练权重初始化的 ViT 在微调时直接使用原生 Muon 会由于正交动量与既有坐标系失配而发生严重性能崩塌（仅 81.25%），而 AMUSE 凭借河床稳定插值机制完全消除了该失配缺陷，取得 83.62% 的卓越性能。
消融实验¶
针对插值机制、动量配置与衰减调度的系统性消融实验（基于 FineWeb 预训练 Llama-124M 架构），验证了各组件对稳定河道巡航的必要性：
| 实验配置变体 | 验证集困惑度 (PPL) | 端到端单步耗时 (s) | 关键观察与机理解释 |
|---|---|---|---|
| AMUSE 完整方案 (\(\beta_{\text{init}}=0.6, \rho=0.8\)) | 28.61 | 0.719 | 兼顾早期高速收敛与后期峭壁抑振，始终保持大更新范数与高余弦相似度 |
| 固定插值 SF-Muon (\(\beta = 0.95\), 最优常数) | 28.94 | 0.718 | 初期因平均序列滞后导致探索受限，后期受限于固定窗口无法维持深层增益 |
| 固定插值 SF-Muon (\(\beta = 0.60\), 偏小常数) | 29.45 | 0.718 | 早期下降迅猛，但训练中后期山壁高曲率分量复发，最终发散振荡 |
| 截断时变调度 (当 \(\beta_t \ge 0.95\) 后饱和固定) | 28.85 | 0.719 | 证实 \(\beta_t\) 必须向 1 持续渐进逼近，单纯达到某一绝对大值无法维持后期进展 |
| 移除正交前 Muon 动量 (AMUSE w/o \(\bm{M}_t\)) | 34.12 | 0.692 | 严重退化，表明正交化前缺乏动量滤波会导致 Newton-Schulz 极度放大采样噪声 |
| 恒定学习率 Muon + 后处理 EWA (多超参最优) | 29.35 | 0.706 | 外挂参数滑动平均无法弥补训练过程中评估点走样导致的劣质更新 |
| 恒定学习率 Muon + 显式 WSD 终端线性退火 | 29.08 | 0.706 | 虽能引导参数沉降至局部谷底，但仍不及动态插值全周期沿河道更新的增益 |
关键发现¶
- 对后验平均与退火的免疫性：在 Llama-124M 训练轨迹上进行 EWA 和学习率强制线性衰减测试，原生 Muon 经过后置衰减后困惑度显著下降（证明其平时严重受困于峭壁振荡），而 AMUSE 在后置衰减或 EWA 处理下性能几乎无额外变化，且相邻步更新余弦相似度极高，表明 AMUSE 轨迹已天然沿着河床平滑行进。
- 超参数极度鲁棒：在 \(\beta_{\text{init}} \in [0.4, 0.7]\) 和 \(\rho \in [0.6, 0.9]\) 的广泛扫描范围内，模型验证集困惑度均稳定在 \(28.6 \sim 28.8\) 之间，所有组合均全面碾压调优后的 Cosine-Muon（29.21）。
- 极小的时间开销与卓越的大 Batch 扩展性：相比于基线 AdamW，AMUSE 的迭代时间仅增加约 6.7%（主要耗时集中于 Newton-Schulz 正交化本身）；同时解决了传统 SF-AdamW 在 720M 和 1.3B 等超大 Batch 训练下梯度方差崩塌的问题。
亮点与洞察¶
- 从河谷地形揭示 Muon 的“双刃剑”本质：通过量化 Hessian 投影分析，首次严谨证明了 Muon 正交化之所以快速是因为强力放大了低曲率主体（River）分量，但同时也因非选择性的极性分解放大了微弱的高曲率优势噪声（Valley walls）。
- 通过时间依赖的插值系数完美规避窗口坍塌：从无调度遍历平均逆向展开推导，将 \(\beta_t\) 的增长速率上限定在保持有效记忆窗口恒定的临界条件上，利用 \(\rho\) 参数实现超参数无关的自适应步进。
- 消除微调 AdamW 权重时的优化器失配：传统观点认为 AdamW 预训练模型不可直接用 Muon 纯矩阵更新微调，AMUSE 证实通过河床低曲率点采样梯度，能够无缝继承并微调常规优化器的预训练资产。
局限与展望¶
- 状态复制内存开销：相比于原生单轨 Muon，AMUSE 内部需额外保留一份基序列 \(\bm{Z}_t\) 的状态副本，虽然实际显存占用未超出常规 AdamW，但在极端参数规模下仍存在显存精简的优化空间。
- 非矩阵参数的混合异构设计：当前实现中对嵌入层和 LayerNorm/RMSNorm 等一维参数依旧退化使用 SF-AdamW，尚未构建起能够统一一维向量与高维矩阵的纯粹几何正交算子。
- 理论收敛阶数的严谨界定：论文主要依托河谷几何经验分析与非凸实证，未来可在非凸随机优化与边缘稳定性（Edge of Stability）理论框架下推导 AMUSE 的显式收敛速率界。
相关工作与启发¶
- vs Muon / D-Muon: Muon 直接在基序列上用正交化动量单轨行进，依赖强学习率衰减压制末期振荡且不支持随时终止；AMUSE 引入无调度双轨机制与时变插值梯度采样，彻底废除调度器并改善更新方向。
- vs Schedule-Free AdamW / SGD: 原始 Schedule-Free 采用固定常数 \(\beta\)，无法调和早期适应与后期抗振的矛盾，且无法处理大模型矩阵正交几何；AMUSE 推广了时变幂律调度并首度与正交矩阵优化器融合。
- vs WSD (Warmup-Stable-Decay): WSD 需提前指定总计算预算并在最后的 10%-20% 步骤消耗专门的 Decay 阶段才能使参数沉入谷底；AMUSE 全程运行在河床中心，随时提取模型均为最优状态，极大节省预训练与续训算力。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ 首次从河谷地形谱分解透视 Muon 机制，提出基于恒定窗口临界界导出的动态时变无调度插值优化方案。
- 实验充分度: ⭐⭐⭐⭐⭐ 覆盖视觉分类、分割、MAE 微调到 1.3B LLM 预训练，提供严谨的 Hessian 谱投影、EWA/WSD 诊断及全方位消融。
- 写作质量: ⭐⭐⭐⭐⭐ 几何直觉明晰，数学推导与实验论证环环相扣，从问题诊断到算法落地逻辑极其严密。
- 价值: ⭐⭐⭐⭐⭐ 为大模型预训练提供了摆脱 Cosine 调度绑定、支持随时训练且提速 1.5× 的新一代即插即用工业级优化器。