Improved Convergence Rates of Muon Optimizer for Nonconvex Optimization

本文通过一种不依赖限制性假设的直接简化分析，为 Muon 优化器建立了比现有结果更精确且适用范围更广的收敛性保证，显著提升了其非凸优化场景下的理论收敛速率。

要点

这篇论文主要是在研究一种叫做 Muon 的“超级优化器”（Optimizer）。在训练人工智能（AI）模型时，优化器就像是教练，负责指导模型如何一步步调整参数，从而从“不会做题”变成“满分学霸”。

为了让你更容易理解，我们可以把训练 AI 的过程想象成一群人在迷雾中下山（寻找最低点，即损失函数的最小值）。

1. 背景：为什么需要 Muon？

• 旧教练（如 Adam）： 以前的教练（比如 Adam 优化器）很聪明，他们会根据每个人（每个参数）的脚力不同，给每个人分配不同的步长。但这就像给每个人发了一双特制的鞋，虽然走得稳，但在大团队（大规模模型）里，大家步调不一致，容易互相干扰，导致下山效率不够高。
• 新教练（Muon）： Muon 这个新教练有个绝招，叫**“正交化”。想象一下，如果一群人下山，旧教练可能让大家乱跑，而 Muon 教练会强行让大家排成整齐的方阵，确保大家的行进方向是互相垂直**的（互不干扰）。这样，整个团队就能在保持方向感的同时，更稳定、更高效地冲下山。
• 问题： 虽然大家发现 Muon 在实际训练中效果很好（下山很快），但数学家们一直没能给出一个完美的理论证明，解释它到底为什么这么快，以及在什么条件下最快。之前的理论要么太粗糙，要么要求太苛刻（比如假设山特别平滑，或者必须满足某些特殊条件）。

2. 这篇论文做了什么？

作者（Shuntaro Nagashima 和 Hideaki Iiduka）就像两个精算师，他们重新拿起了 Muon 的“训练手册”，用更简单、更直接的方法，算出了 Muon 下山速度的精确数学公式。

他们发现，只要调整几个关键变量，Muon 的下山速度可以比之前认为的快得多。

关键发现（用比喻解释）：

A. 步长（学习率）与队伍规模（批量大小）的魔法组合
在训练 AI 时，有两个重要参数：

1. 步长（Learning Rate）： 每次迈多大步子。
2. 批量大小（Batch Size）： 每次看多少张地图（数据）来决定下一步怎么走。

• 以前的看法： 大家认为，无论怎么调，Muon 下山的速度上限大概是 $1/\sqrt{T}$（T 是步数）。
• 作者的发现： 只要把“批量大小”调得足够大（甚至随着训练过程指数级增长，比如每走一步，看地图的数量就翻倍），Muon 的速度就能提升到 $1/T$！
  • 比喻： 就像以前大家是“盲人摸象”，每次只摸一点点就决定方向。现在，随着训练深入，我们不仅步长变小（更谨慎），而且每次看的地图越来越多（信息更全）。当信息量足够大时，Muon 就能像开了“上帝视角”一样，直接冲向目标，速度翻倍。

B. 不需要“完美地形”的假设

• 以前的局限： 很多理论证明必须假设“山是完美的碗状”（PL 条件），这在现实中很难满足。
• 现在的突破： 作者证明了，即使山是崎岖不平的（非凸优化，这是 AI 训练的真实情况），Muon 依然能保证收敛，而且速度更快。这意味着 Muon 不仅理论漂亮，而且实战性极强。

C. 不同的“下山策略”
论文详细分析了四种不同的步长策略（恒定的、像余弦波一样先大后小的、多项式衰减的、逐渐变小的）和两种批量大小策略（固定的、指数增长的）。

• 结论： 最完美的组合是：步长逐渐变小 + 批量大小指数级变大。这种组合能让 Muon 在不需要任何特殊假设的情况下，达到理论上的最快收敛速度。

3. 这对我们意味着什么？

• 理论更扎实： 以前用 Muon 是“因为好用所以用”，现在有了论文证明“为什么好用”以及“怎么调最好”。
• 训练更快、更省资源： 对于开发大模型（如 LLM）的公司来说，这意味着可以用更少的计算资源、更短的时间训练出更好的模型。
• 指导实践： 论文告诉工程师们，不要死守固定的批量大小。试着让批量大小随着训练过程“膨胀”起来，配合 Muon 优化器，效果会出奇的好。

总结

这就好比以前大家用 Muon 下山，虽然跑得快，但不知道极限在哪里，也不敢跑太快怕摔跟头。
这篇论文就是给 Muon 发了一张**“极速通行证”，并画出了一张“最佳路线图”**：只要按照“步长渐小、视野渐宽（批量变大）”的策略走，Muon 就能在复杂的地形中，以理论允许的最快速度，稳稳地到达终点。

一句话总结： 作者通过数学证明，发现 Muon 优化器在配合“动态扩大的数据视野”时，能比想象中快得多，且不需要苛刻的前提条件，为大模型训练提供了更高效的理论依据。

技术摘要

这是一份关于论文《Improved Convergence Rates of Muon Optimizer for Nonconvex Optimization》（Muon 优化器在非凸优化中的改进收敛速率）的详细技术总结。

1. 研究背景与问题 (Problem)

背景：
随着大规模深度神经网络（DNN）的训练需求增加，优化算法的稳定性与效率至关重要。Muon 优化器（Momentum orthogonalized by Newton-Schulz）作为一种新兴的优化方法，因其通过牛顿 - 舒尔茨（Newton-Schulz）迭代将更新方向正交化，从而在保持数值稳定性的同时有效控制更新方向的相关性，受到了广泛关注。相比 Adam 等自适应方法，Muon 在高维空间中能更好地保留方向信息。

现有问题：
尽管 Muon 在实践中表现出色，但其理论收敛性分析尚不完善。现有的收敛保证存在以下局限性：

• 假设过强： 部分结果依赖于强假设（如 Polyak-Łojasiewicz (PL) 条件），限制了适用范围。
• 收敛速率较慢： 现有理论给出的收敛速率（如 $O(T^{-1/4})$ 或 $O(T^{-1/2})$）往往低于实际观察到的性能，且不如 SGD 及其变体在特定设置下的理论界限。
• 分析不严谨： 部分分析中残留了与网络维度相关的不可忽略项，未能给出严格的收敛结果。
• 缺乏通用性： 现有分析未能涵盖多种学习率调度（如余弦退火、多项式衰减）和批量大小（Batch Size）策略的组合。

核心目标：
本文旨在通过直接且简化的分析，在不依赖限制性假设（如 PL 条件）的情况下，为 Muon 优化器建立更精确、更广泛的收敛保证，并推导出优于现有文献的收敛速率。

2. 方法论 (Methodology)

核心算法：
论文分析了 Muon 优化器的标准变体（包含 Nesterov 动量加速）。其核心步骤包括：

1. 计算小批量随机梯度 $\nabla f_{B_t}(W_t)$。
2. 更新动量项 $M_t = \beta M_{t-1} + (1-\beta)\nabla f_{B_t}(W_t)$。
3. （可选）若使用 Nesterov，则计算 $C_t = \beta M_t + (1-\beta)\nabla f_{B_t}(W_t)$，否则 $C_t = M_t$。
4. 正交化步骤： 将 $C_t$ 投影到正交矩阵集合上，即 $O_t = \arg\min_{O^\top O=I} \|O - C_t\|_F$。
5. 更新参数：$W_{t+1} = W_t - \eta_t O_t$。

理论分析框架：

• 假设条件： 仅假设损失函数是平滑的（Smoothness，满足 Lipschitz 连续梯度）且随机梯度是无偏的、方差有界的。不假设凸性或 PL 条件。
• 关键引理：
  • 下降引理 (Descent Lemma)： 利用平滑性建立 $f(W_t) - f(W_{t+1})$ 的下界，其中涉及梯度范数 $\|\nabla f(W_t)\|_F$ 和正交化误差 $\|\nabla f(W_t) - C_t\|_F$。
  • 正交性结构利用： 利用 $O_t$ 是 $C_t$ 在正交流形上的投影这一性质，结合对偶范数性质，推导出 $C_t \cdot O_t \ge \|\nabla f(W_t)\|_F - \sqrt{n}\|C_t - \nabla f(W_t)\|_F$ 的关键不等式。
  • 误差累积分析： 通过归纳法分析动量项 $M_t$（或 $C_t$）与真实梯度 $\nabla f(W_t)$ 之间的期望误差，该误差由初始误差、学习率累积项和梯度方差项组成。
• 参数设置： 分析了四种学习率策略（常数、余弦退火、多项式衰减、递减）和两种批量大小策略（常数、指数增长）的组合。

3. 主要贡献 (Key Contributions)

1. 建立了更通用的收敛上界：
   提出了 Muon 优化器生成的全梯度期望上界（Theorem 3.1）。该上界由五项（无 Nesterov）或六项（有 Nesterov）组成，分别依赖于学习率 $\eta_t$、批量大小 $b_t$ 和动量参数 $\beta$。证明过程简化了现有分析，未依赖 PL 条件等强假设。

2. 显著提升了收敛速率：
   在 $O(\cdot)$ 符号意义下，证明了 Muon 在特定参数设置下能达到 $O(1/T)$ 的收敛速率，这优于现有文献中的 $O(1/\sqrt{T})$ 或 $O(T^{-1/4})$。

   • 常数学习率 + 指数增长批量大小： 当 $\eta = O(1/T)$ 且 $b_t = b\delta^t$ 时，收敛速率为 $O(1/T)$。
   • 递减学习率 + 指数增长批量大小： 当 $\eta_t = \eta/\sqrt{t+1}$ 且 $b_t = b\delta^t$ 时，收敛速率为 $O(\log T / \sqrt{T})$。

3. 揭示了批量大小增长的重要性：
   理论分析表明，随着训练步数 $T$ 增加，增大批量大小（特别是指数增长）能显著降低收敛上界中的方差项，从而加速收敛。这一发现与 SGD 领域的最新趋势一致，并首次被严格应用于 Muon 的理论分析中。

4. 广泛的适用性验证：
   涵盖了多种实际常用的超参数配置（如余弦退火学习率、Nesterov 动量），证明了在这些设置下 Muon 均能保持理论上的收敛性。

4. 主要结果 (Results)

论文通过 Corollary 3.1 总结了不同设置下的收敛速率（针对 $\min_{t} \mathbb{E}[\|\nabla f(W_t)\|_F]$）：

设置组合	学习率 ($\eta_t$)	批量大小 ($b_t$)	收敛速率 (无 Nesterov/有 Nesterov)	备注
常数 LR + 常数 BS	$\eta$	$b$	$O(1/T + \eta + 1/\sqrt{b})$	若 $\eta=O(1/T), b=O(T^2)$，可达 $O(1/T)$
常数 LR + 指数增长 BS	$\eta$	$b\delta^t$	$O(1/T + \eta)$	关键发现：若 $\eta=O(1/T)$，可达 $O(1/T)$
余弦退火 LR + 指数增长 BS	$\eta \frac{1}{2}(1+\cos \dots)$	$b\delta^t$	$O(1/T + \eta)$	同样可达 $O(1/T)$
递减 LR + 指数增长 BS	$\eta/\sqrt{t+1}$	$b\delta^t$	$O(\log T / \sqrt{T})$	优于现有 $O(1/\sqrt{T})$ 结果

对比现有工作：

• 优于 Tang et al. (2025) 和 Shen et al. (2025) 的 $O(T^{-1/4})$。
• 优于 Li & Hong (2025) 的 $O(1/\sqrt{T})$（在特定条件下）。
• 无需 Chang et al. (2025) 所需的 PL 条件即可达到类似甚至更好的速率。

5. 意义与影响 (Significance)

• 理论突破： 本文首次在不依赖强假设（如 PL 条件）的情况下，证明了 Muon 优化器在非凸问题中可以达到 $O(1/T)$ 的收敛速率。这填补了 Muon 理论分析与其实用高效性之间的鸿沟。
• 指导实践： 研究结果明确指出了**“指数增长的批量大小”与“适当衰减的学习率”**相结合是发挥 Muon 性能的关键。这为大规模模型训练中的超参数调优提供了坚实的理论依据。
• 通用性启示： 分析中使用的简化技术（利用正交化结构和对偶范数）不仅适用于 Muon，也为其他基于正交化的一阶优化方法提供了新的理论分析范式。
• 替代潜力： 理论上的优越收敛性进一步支持了 Muon 作为 Adam 等传统优化器替代方案的潜力，特别是在对训练稳定性和收敛速度要求极高的大规模预训练场景中。

综上所述，该论文通过严谨的数学推导，不仅提升了 Muon 优化器的理论收敛上限，还给出了具体的超参数配置建议，为深度学习优化领域的理论研究和工程实践提供了重要参考。