跳转至

Muon Does Not Converge on Convex Lipschitz Functions

会议: NeurIPS 2026
arXiv: 2605.08980
OpenReview: https://openreview.net/forum?id=78YklG22EU
领域: 优化/理论
关键词: 优化理论、Muon优化器、凸利普希茨函数、误差反馈、非欧几里得次梯度法

一句话总结

本文证明了在深度学习中表现优异的谱下降优化器 Muon 即使搭配衰减学习率在凸 Lipschitz 目标上也无法收敛,揭示了其等价于对角符号动量法,并证明虽可通过误差反馈恢复理论收敛,但这层理论补丁在实际视觉与大语言模型预训练中反而导致性能显著退化。

研究背景与动机

以 AdaGrad、RMSProp、Adam、Shampoo 以及近年兴起的 Schedule-Free 为代表的一系列现代优化器,其算法设计与理论收敛性保证绝大部分建立在凸利普希茨(convex Lipschitz)非光滑函数类之上。凸 Lipschitz 分析长期以来被视为理解深度学习损失动态与设计实用优化算法的最成功基准之一。最近,基于谱下降(spectral descent)与极分解(polar decomposition)的 Muon 优化器在 124M 至 1.6T 规模的多模态与大语言模型(如 nanoGPT、Moonlight、Kimi 与 DeepSeek-V4)预训练中展现出远超传统欧氏优化器的经验优势,但现有的收敛性分析几乎全部严重依赖光滑性(smoothness)假设。

这一现状引发了一个根本性的理论悬念:凸 Lipschitz 理论体系能否同样解释 Muon 的收敛行为?直觉上,作为矩阵极分解方向的约束与正则化次梯度下降,Muon 与矩阵预条件方法 Shampoo 享有一脉相承的几何动机,理应在凸非光滑问题上具备某种形式的收敛保证。然而,凸分析在非欧氏几何算子与动量项交互下的动力学极为微妙,次梯度算子的内在偏差极易阻碍迭代点向最优解收敛。

本文明确指出凸 Lipschitz 模型并不适合作为理解 Muon 成功机制的理论基石。核心 idea:通过证明对角矩阵上 Muon 严格等价于符号动量法(signed momentum),构造特定凸 Lipschitz 反例证明 Muon 即使在任意步长调度下也无法收敛;同时证明正则化非欧更新本质是压缩算子,虽能借由误差反馈(Error Feedback)挽救收敛性,但该理论补丁在深度学习实战中会导致性能劣化,从而断定 Muon 的实战优势来源于非凸与光滑几何结构而非凸 Lipschitz 假说。

方法详解

整体框架

为了系统诊断 Muon 及其非欧几里得泛化族在非光滑凸问题上的行为,本文建立了“对角约简 → 反例构造证明发散 → 压缩算子统一视点 → 误差反馈挽救收敛”的理论分析框架。首先,针对矩阵参数目标函数,将 Muon 更新与正则化版本归约到对角矩阵函数族,严格证明其退化为符号动量法。随后,利用低维凸非光滑目标的不变子空间与次梯度跳跃特性,证明迭代点被永久锁死在远离最优解的界外。为了彻底解决这一发散性,本文从线性最大化预言机(LMO)与极分解导出锐算子(sharp operator)的压缩性质,提出带动量的误差反馈通用框架(EF-M / EF-Muon),并在凸 Lipschitz 假设下建立了随时(anytime)收敛速率。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    In["矩阵参数优化目标<br/>凸 Lipschitz 函数 f(W)"] --> Red["对角函数等价性规约<br/>polar 算子化为 elementwise sign"]
    Red --> CE["反例动力学分析<br/>正交对称基下坐标被永久锁死"]
    CE --> Op["非欧次梯度压缩算子化<br/>LMO 极分解映射为 δ-压缩"]
    Op --> EF["误差反馈记忆与动量校正<br/>EF-M 追踪残差 Pt 恢复收敛"]
    EF --> Out["收敛性对比与经验证伪<br/>理论恢复 O(1/√T) 速率但实战性能劣于原生 Muon"]

关键设计

1. 对角函数的极分解符号约简:将谱下降降维映射为符号动量动力学

谱范数约束下的次梯度下降计算目标梯度的极因子 \(\operatorname{polar}(G_t) = U V^\top\)。对于仅依赖矩阵对角线元素的函数 \(f(W) = f_{\operatorname{diag}}(\operatorname{diag}(W))\),其任意次梯度 \(G_t \in \partial f(W_t)\) 均为对角矩阵。当对角矩阵 \(A = \operatorname{Diag}(a)\) 做降维 SVD 时,奇异值矩阵为绝对值 \(|a_i|\),左、右奇异向量外积恰好恢复符号 \(\operatorname{sign}(a_i)\)。由此引理严格证明 \(\operatorname{polar}(A) = \operatorname{sign}(A)\)。从而,在对角函数上,Muon 优化器的动量更新 \(M_t = \beta M_{t-1} + (1-\beta) G_t\) 以及权重更新 \(W_{t+1} = W_t - \lambda_t \operatorname{polar}(M_t)\) 严格等价于多维独立/协同符号动量算法(signed momentum)。这一极简等价性彻底摆脱了复杂的非交换矩阵极分解近似细节,为构造清晰解析的反例打开了突破口。

2. 鞍形凸折线反例构建:对称次梯度共振导致迭代点永久驻留边界

在降维到二维对角向量空间后,定义非光滑凸函数: $\(f(W) = c |W_{11} + W_{22}| + |W_{11} - W_{22}|, \quad c \in (0,1)\)$ 该函数具有唯一的全局极小点 \(W_{11}^* = W_{22}^* = 0\)。通过引入正交基向量 \(u_1 = [1, 1]^\top\) 与 \(u_2 = [1, -1]^\top\),迭代轨迹在 \(u_1\) 方向和 \(u_2\) 方向呈现截然不同的动态。在离线步长调度下(Counterexample 3.3),当选取常数 \(c = \frac{1-\beta}{2}\) 时,次梯度在 \(u_2\) 方向上按步交替翻转符号,驱动动量项 \(m_t\) 的符号严格由 \(u_2\) 轴的主导分量锁定,导致在 \(u_1\) 投影轴上的更新位移完全被抵消或锁死。迭代点在 \(u_1\) 方向恒定维持初值,从而恒有 \(f(W_t) - \inf f \geq 1-\beta\) 且 \(\|W_t - W^*\|_F \geq \frac{1-\beta}{2}\),哪怕步长 \(\lambda_t \to 0\) 且严格递减,也无法消除常数下界误差。进一步地,在适应性步长和低动量 \(\beta \in [0, 1/2)\) 设定下(Counterexample 3.4),构造零测度集合排除了可达路径,证明对于几乎所有初始点,Muon 均在局部与全局意义下发散。

3. 非欧次梯度算子的压缩算子等价性:统一谱下降与有偏压缩通信

在矩阵范数理论中,正则化非欧次梯度更新依赖线性最大化预言机 \(\operatorname{LMO}_{\|\cdot\|}(M) = \operatorname{argmax}_{\|X\| \leq 1} \langle X, M \rangle\) 与对偶范数 \(\|M\|_*\)。本文证明:对于任意在 Frobenius 范数下满足等价性边界 \(\alpha \|W\|_F \leq \|W\| \leq \beta \|W\|_F\) 的赋范空间,锐算子经过尺度缩放后 \(\mathcal{C}(W) := \alpha^2 \|W\|_* \operatorname{LMO}_{\|\cdot\|}(W)\) 严格满足收缩条件: $\(\|W - \mathcal{C}(W)\|_F^2 \leq (1 - \delta) \|W\|_F^2, \quad \delta = \frac{\alpha^2}{\beta^2}\)$ 特别地,对于谱范数 \(\|\cdot\|_{\operatorname{op}}\),对应对偶范数为核范数 \(\|\cdot\|_{\operatorname{nuc}}\),最小 Frobenius 范数的 LMO 解恰为极分解因子 \(\operatorname{polar}(W)\),收缩常数 \(\delta = 1 / \min\{m, n\}\)。该发现首次证明非欧几里得梯度更新在数学本质上是有偏的几何压缩算子,从底层逻辑解释了 Muon 与有偏梯度压缩算法共有的发散病理。

4. 带动量的误差反馈补偿机制(EF-M):记忆累积与一阶收敛恢复

针对有偏压缩导致的收敛失败,本文将分布式通信中的误差反馈机制扩展到非欧几何与动量体系,设计了 EF-M(Error Feedback with Momentum)算法。引入累积残差记忆张量 \(E_t \in \mathbb{R}^{m \times n}\),在每步将学习率缩放的动量与上一轮未释放的误差合成待压缩信息 \(P_t = E_t + \lambda_t M_t\)。模型参数更新直接采用压缩后的方向 \(W_{t+1} = W_t - \mathcal{C}(P_t)\),而把量化残差保留在 \(E_{t+1} = P_t - \mathcal{C}(P_t)\) 中供后续迭代补偿。针对凸 Lipschitz 随机目标函数,在采用平方反比衰减步长 \(\lambda_t = 1/\sqrt{t+1}\) 条件下,定理 4.2 证明了其时间平均迭代点 \(\bar{W}_T = \frac{1}{T+1} \sum_{t=0}^T W_t\) 满足: $\(\mathbb{E}[f(\bar{W}_T)] - f(W^*) \leq \frac{\|W_0 - W^*\|_F^2}{2\sqrt{T+1}} + \sigma^2 \left( \frac{2\sqrt{1-\delta}}{\delta} + \frac{\beta}{1-\beta} + \frac{1}{2} \right) \frac{1 + \log(T+1)}{\sqrt{T+1}}\)$ 成功将凸非光滑环境下的渐近收敛速率恢复至紧致的 \(\mathcal{O}\left(\frac{\log T}{\sqrt{T}}\right)\)。

损失函数 / 训练策略

在将算法应用于多层深度神经网络的实测中,针对全连接层、注意力投影矩阵以及多维张量组成的乘积参数空间 \(\mathbb{R}^{m_1 \times n_1} \times \dots \times \mathbb{R}^{m_L \times n_L} \times \mathbb{R}^k\),算法采用分块复合极分解。对于矩阵层(权重),采用 MuonMax 规范化的 LMO 映射: $\(\mathcal{C}(P_t)^\ell = \min\left\{s, \frac{1}{L}\right\} \frac{\sum_{j=1}^L \frac{1}{\sqrt{d_j}} \|P_t^j\|_{\operatorname{nuc}}}{\sqrt{d_\ell}} \operatorname{polar}(P_t^\ell)\)$ 对于偏置与一维向量层,则配对符号动量或 NAdamW。超参数包含动量衰减 \(\beta\)、网络基础学习率 \(\lambda\)、层级缩放系数 \(s\) 以及学习率倍率乘数(LR multiplier)。

实验关键数据

主实验

论文在真实深度学习任务中对理论修补的有效性进行了严格的实证检验。评估涵盖代表性计算机视觉图像分类任务(CIFAR-10 数据集上的 36M 参数 WideResNet-28-10)与现代大语言模型自回归预训练任务(FineWeb-Edu 10B 子集上的 206M 参数 nanoGPT 结构)。对比的基线包括 Vanilla Muon、无误差反馈的多层正则化基线 MuonMax 以及由定理 4.2 理论指导设计的收敛算法 EF-MuonMax。在所有设置下均进行了网格超参数调优(每种方法 30 至 60 次独立实验探索最优学习率和倍率)。

任务类型 优化器方法 基础学习率 学习率倍率 验证集交叉熵损失 (Val CE) 相对性能表现
图像分类 (CIFAR-10, 36M WRN) Vanilla Muon \(10^{-4}\) 1.0 0.287 次优基准
图像分类 (CIFAR-10, 36M WRN) MuonMax \(5 \times 10^{-4}\) 100.0 0.265 表现最优
图像分类 (CIFAR-10, 36M WRN) EF-MuonMax \(5 \times 10^{-3}\) 30.0 0.388 性能显著劣化 (+46.4%)
语言建模 (FineWeb-Edu 10B, 206M) Vanilla Muon \(1.36 \times 10^{-4}\) 28.6 3.01 表现最优
语言建模 (FineWeb-Edu 10B, 206M) MuonMax 0.01 30.0 3.27 稍逊于 Vanilla (+0.26)
语言建模 (FineWeb-Edu 10B, 206M) EF-MuonMax 0.01 30.0 3.43 性能垫底 (+0.42)

消融实验

通过对理论与实验设定的对比消融,分析凸非光滑理论性质与真实深度学习非凸光滑曲率之间的脱节关系:

机制 / 理论设定 凸 Lipschitz 函数收敛保证 优化器动力学行为 真实深度模型训练损失表现
Vanilla Muon (specGD + Polyak 动量) ❌ 无法收敛 (反例 3.3, 3.4 下界固定) 极因子仅保留奇异方向,对角线等价符号震荡 图像表现优良 (0.287),语言建模最佳 (3.01)
MuonMax (核范数归一化正则化) ❌ 无法收敛 (受适应性步长反例 3.4 制约) 结合整体谱半径平衡更新,仍属有偏非欧梯度 图像分类表现最佳 (0.265),语言次之 (3.27)
EF-MuonMax (加入误差反馈记忆) ✅ 保证收敛 (\(\mathcal{O}(\log T / \sqrt{T})\) 随时速率) 历史残差回补强制抵消次梯度几何方向偏置 两任务全面恶化 (图像 0.388 / 语言 3.43)

关键发现

  • 理论与经验的完全背离:理论上能保证凸 Lipschitz 全局收敛的误差反馈修正版本(EF-MuonMax),在图像分类任务上验证损失从 0.265 恶化至 0.388(相对上升 46.4%);在语言建模任务中验证损失从 3.01 恶化至 3.43,且整个训练进程中收敛曲线全面被原始 Muon 与 MuonMax 压制。
  • LMO 不应被视作纯粹的通信有偏压缩器:在分布式训练中,压缩算子的误差是通信带宽受限下被迫付出的代价,因此通过误差反馈补偿残差极其有效;但在 Muon 中,极分解正交化并非近似更新,而是利用谱范数对模型隐藏层权重矩阵进行特征维度的各向同性伸缩与曲率均一化。引入误差反馈强制模型“偿还”偏离梯度的残差,反而破坏了正交更新带来的几何结构优势。
  • 最坏情况凸分析对 Muon 失去指导意义:与 Adam 在凸非光滑上发散促成 AMSGrad 理论提出却未在深度学习中超越 Adam 相似,Muon 在最坏情况凸 Lipschitz 目标上的数学缺陷,正是其在真实平滑非凸损失面上高效摆脱病态曲率的副产物。

亮点与洞察

  • 极分解与符号函数的桥梁搭建:巧妙地利用对角矩阵在降维 SVD 下极因子 \(\operatorname{polar}(A) = \operatorname{sign}(A)\) 的代数等价性,将高深抽象的高维矩阵谱下降,严密规约到经典的对角符号动量动力学(Signum / signed momentum),使得解析构造发散反例成为可能。
  • 对理论范式的敏锐质疑与实证纠偏:面对深度学习中盲目套用“凸非光滑利普希茨”经典理论框架的惯性,明确指出并非所有优化器都能套用该范式,通过精细的反例证明与反向实证,指明 Muon 的真正理论根基应当在各向异性光滑性(smoothness)或非凸特征学习动力学中寻找。

局限与展望

  • 反例设定的理论局限:Counterexample 3.3 仅适用于单调递减步长,且依赖在狭窄的初值集合中启动;Counterexample 3.4 虽推广到近乎全空间测度的初始点及包含核范数的自适应步长,但动量参数范围被局限在 \(\beta \in [0, 1/2)\),尚未完全覆盖工业界主流使用的较大动量区间(如 \(\beta = 0.9\) 或 \(0.95\))。
  • 动量调度与加速机制未被纳入:论文主要分析了 Polyak 形式的固定动量系数,未涵盖实际大模型训练中逐步递增到 1 的动量退火策略(\(\beta_t \uparrow 1\)),也未涵盖 Nesterov 加速形式或 Newton-Schulz 矩阵多项式迭代带来的微小数值近似漂移。
  • 未来方向:需要从局部平滑性(generalized smoothness)、非凸谱曲率模型(isotropic curvature model)以及多层神经网络特征学习的动态平衡(feature learning spectral conditions)等全新假说出发,重新构建原生 Muon 的收敛分析基石。

相关工作与启发

  • vs signSGD / Error Feedback (Karimireddy et al., 2019): signSGD 在凸非光滑问题上同样发散,且 EF-signSGD 能够成功修复该问题并在分布式通信压缩中提升鲁棒性。本文借鉴了其反例思路与误差反馈思想,但推向了更高维的矩阵正交谱下降领域,并发现误差反馈在深度学习网络内反而大幅削弱了 Muon 的表现。
  • vs Adam / AMSGrad (Reddi et al., 2018): Reddi 等人揭示了 Adam 在凸非光滑问题上的发散并提出了 AMSGrad 修正;本文对 Muon 进行了一脉相承的“理论证伪与算法修正”剖析,同样论证了这种修补虽具数学完备性但无益于实际工程。
  • vs Shampoo / AdaGrad (Gupta et al., 2018; Duchi et al., 2011): Shampoo 和 AdaGrad 依靠二阶外积矩阵累积预条件变换梯度空间度量,在凸 Lipschitz 条件下具备严密的在线后悔界和收敛性;本文表明 Muon 的正交化动力学与基于度量变换的预条件优化器存在根本性的理论分野,不能套用相同的分析模式。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 首次揭示 Muon 与符号动量在对角矩阵上的严格等价性,并给出非欧谱下降在凸非光滑目标上的发散反例。
  • 实验充分度: ⭐⭐⭐⭐☆ 包含严格的数学证明与反例构造,并在 CIFAR-10 及 206M nanoGPT 上进行了充分的网格搜索验证,有力支撑理论论点。
  • 写作质量: ⭐⭐⭐⭐⭐ 论证严谨,逻辑层层递进,从数学抽象到实证讨论一气呵成,深入浅出。
  • 价值: ⭐⭐⭐⭐⭐ 纠正了社区盲目将凸利普希茨分析套用于 Muon 的研究倾向,为后续谱优化理论指明了依赖光滑性与非凸特征的正确路径。