Understanding Gradient Orthogonalization for Deep Learning via Non-Euclidean Trust-Region Optimization¶

会议: NeurIPS 2026
arXiv: 2503.12645
领域: 优化算法 / Optimization
关键词: 梯度正交化, Muon优化器, 非欧几里得置信域, 谱范数, 权重衰减

一句话总结¶

本文首次建立了梯度正交化方法（如Muon）与非欧几里得谱范数置信域优化之间的理论等价性，提出了带动量、权重衰减及外推的统一随机算法框架，并给出了非凸、星凸与二阶光滑设定下的收敛率分析。

研究背景与动机¶

在大语言模型与深度神经网络的训练中，以AdamW为代表的自适应对角预条件优化器长期占据主导地位。近期Jordan等人提出的Muon优化器与Liu等人在大模型训练中的扩展实践表明，针对隐层矩阵参数采用梯度正交化（orthogonalized gradient updates）能够显著超越AdamW。然而，尽管梯度正交化在经验上表现极其出色，其背后的理论优化机制一直缺乏严格坚实的基础。此前部分文献试图将其解释为非欧矩阵谱范数下的普通梯度下降，但在更新步长中遗漏了核范数缩放因子，无法精确复现实际算法迭代。

理论认知的缺失进一步导致两个关键经验现象无法得到合理解释：其一，形式上高度相似的Orthogonal-SGDM（先正交化再滑动平均）在实践中严重落后于Muon（先滑动平均再正交化），甚至打不过微调后的标准SGDM；其二，在小模型训练中权重衰减（weight decay）并不敏感，但在千亿级大语言模型预训练中，缺失权重衰减会导致参数谱半径发散、Muon相对AdamW的收益迅速消退。

针对上述理论真空，本文从一阶置信域优化的几何视角切入，指出梯度正交化的数学本质正是矩阵谱范数约束下的线性近似置信域极小化。核心 idea：将梯度正交化严格等价刻画为基于矩阵谱范数的非欧几里得置信域梯度方法，并在此统一几何框架下构建带动量、权重衰减与外推机制的随机算法族，系统解释Muon优于Orthogonal-SGDM的方差缩减机理及大模型依赖权重衰减的紧致性保障。

方法详解¶

整体框架¶

本文考虑带有凸正则项的复合非凸优化问题 $\min_{x \in \mathcal{X}} [F(x) = f(x) + R(x)]$，其中向量空间 $\mathcal{X}$ 配备任意非欧几里得范数 $\|\cdot\|$ 且其对偶范数记为 $\|\cdot\|_*$。算法不再依赖Majorization-Minimization（MM）框架下的二阶上界二次正则代理，而是在局部置信域球内直接最小化一阶线性逼近模型。

当空间特化为矩阵空间 $\mathbb{R}^{m \times n}$ 且使用谱范数（算子范数）$\|X\| = \|X\|_{\mathrm{op}} = \|\bm{\sigma}(X)\|_\infty$、正则项 $R \equiv 0$ 时，该置信域极小化步精确导出了矩阵奇异值投影算子 $\text{orth}(G) = (GG^\top)^{\dagger/2}G = UV^\top$。在此基础上，论文构建了融合重球动量、中心平移权重衰减和Nesterov风格外推梯度的完整演化体系。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["随机小批量采样<br/>ξ_k ~ D"] --> B["动量滑动更新<br/>m_{k+1} = (1-α)m_k + α g(x_k; ξ_k)"]
    B --> C["谱范数置信域子问题<br/>min ⟨m_{k+1}, x⟩ + R(x) s.t. ‖x - (1-β)x_k‖ ≤ η"]
    C --> D["正交动量参数落盘<br/>X_{k+1} = (1-β)X_k - η orth(M_{k+1})"]
    D --> E["二阶曲率外推补偿 (可选)<br/>x̄_{k+1} = x_k + (1/α)(x_{k+1} - x_k)"]

关键设计¶

1. 谱范数置信域等价映射：精确推导正交化迭代机制 以往研究常误将正交梯度下降等同于谱范数下的MM梯度下降，导致更新式中出现多余的核范数因子 $\|G_k\|_{\mathrm{nuc}}$，偏离实际实现。本文指出，置信域一阶逼近子问题 $\min_{\|X - X_k\|_{\mathrm{op}} \leq \eta} \langle G_k, X - X_k \rangle$ 的精确解析解正是奇异值分解截断后的正交基乘积： $$X_{k+1} = X_k - \eta \, \text{orth}(G_k), \quad \text{orth}(G) = (GG^\top)^{\frac{\dagger}{2}} G = U V^\top$$ 这一几何视角不仅摆脱了欧氏内积对梯度模长的直接依赖，还将更新限制在谱半径球内，直接赋予了隐层变换在对抗扰动和向前传播时最大的谱稳定性。

2. 动量累积顺序与批次效应：揭示Muon超越Orthogonal-SGDM的本质 针对Orthogonal-SGDM与Muon的先后序之争，本文在非欧置信域框架下给出了方差分析。Muon先对无偏随机梯度进行指数移动平均更新 $m_{k+1} = (1-\alpha)m_k + \alpha g(x_k, \xi_k)$，再将平滑后的动量向量传入非欧置信域步；而Orthogonal-SGDM则是对每步含噪随机梯度的正交化产物做动量累加。由于原始随机噪声 $\zeta_k = \nabla f(x_k) - g(x_k, \xi_k)$ 满足零均值独立性，Muon的线性凸组合能够通过批次累积效应（batching effect）显著压制噪声方差： $$\mathbb{E}[\|m_{k+1} - \nabla f(x_k)\|_*] \leq (1-\alpha)^{k+1}\rho\sigma + \sqrt{\alpha}\rho\sigma + \frac{L\eta}{\alpha}$$ 相反，由于正交化算子是非线性的，$O_{k+1} = \text{orth}(g(x_k, \xi_k))$ 具有严重偏差且期望不再对齐真梯度，对其累加动量无法消除非线性截断引入的随机畸变。

3. 置信域中心平移权重衰减：突破星凸优化下迭代无界困境 在非凸及星凸深度学习损失面下，由于缺乏强凸性或紧致域假定，置信域迭代序列可能无限发散。本文提出将置信域中心向原点收缩 $(1-\beta)$，即求解： $$\min_{x \in \mathcal{X}} \left[\langle m_{k+1}, x \rangle + R(x)\right] \quad \text{s.t.} \quad \|x - (1-\beta)x_k\| \leq \eta$$ 只要满足步长与衰减率几何相容条件 $\eta \geq \beta \max\{\|x_0\|, \|x^*\|\}=\beta D$，即可无条件保证全迭代序列有界性 $\|x_k\| \leq D$，进而将星凸目标下的次优度收敛复杂度从非凸的 $\mathcal{O}(1/\varepsilon^4)$ 提升至与凸优化平行的 $\tilde{\mathcal{O}}(1/\varepsilon^3)$。这从数学上严格验证了Liu等人在千亿LLM长程预训练中必须引入解耦权重衰减以控制矩阵谱增长的实证结论。

4. 动量点外推更新：在二阶光滑非凸目标下实现超欧氏加速 当损失函数具备 $H$-Hessian Lipschitz二阶光滑性质时，本文在动量更新中引入超前外推点 $\bar{x}_{k+1} = x_k + \frac{1}{\alpha}(x_{k+1} - x_k)$。在外推点处评估随机梯度能够抵消曲率变动带来的动量滞后误差，将动量跟踪界中的 $L\eta/\alpha$ 项降为二阶项 $\frac{H\eta^2}{2\alpha^2}$。在非凸复合目标下，该外推机制将寻找 $\varepsilon$-稳态点的复杂度直接优化到 $\tilde{\mathcal{O}}(1/\varepsilon^{3.5})$，首次在任意非欧范数及矩阵谱优化中实现了超越一阶随机下界的加速收敛。

损失函数 / 训练策略¶

统一置信域动量优化器的参数配置与理论步长选取严格遵循对偶范数连续常数 $L$、对偶欧氏转换系数 $\rho$（满足 $\|x\|_* \leq \rho \|x\|_2$）及噪声方差 $\sigma$： - 非凸目标（Algorithm 1）：动量系数 $\alpha = \mathcal{O}\left(\min\{1, \frac{\varepsilon^2}{\rho^2\sigma^2}\}\right)$，步长 $\eta = \mathcal{O}\left(\min\{\frac{\varepsilon}{L}, \frac{\varepsilon^3}{\rho^2\sigma^2 L}\}\right)$； - 星凸目标带权重衰减（Algorithm 2）：选取 $\beta = \mathcal{O}\left(\min\{1, \frac{\varepsilon}{LD^2}, \frac{\varepsilon}{D\rho\sigma}, \frac{\varepsilon^3}{D^3\rho^3\sigma^3}, \frac{\varepsilon^3}{LD^3\rho^2\sigma^2}\}\right)$，并绑定步长 $\eta = \beta D$； - 二阶光滑带外推（Algorithm 3）：外推比例系数设定为 $\gamma = 1/\alpha$，步长自适应引入 $\varepsilon^{1/2}/H^{1/2}$ 缩放。

实验关键数据¶

主实验：算法收敛复杂度全景对比¶

论文对不同假设条件下的置信域算法族、对应特化优化器及收敛复杂度上界进行了严格解析推导与理论对比（下表中所有复杂度均表示达到 $\varepsilon$-精度所需的随机迭代步数 $K$）：

算法 / 实例	适用范数与算子	函数假设	正则项 $R(x)$	评估指标	迭代复杂度 $K$
TR-GD (Theorem 1)	任意非欧范数 / 谱范数	一阶光滑 (L)	闭凸正则项	$\\|\nabla f + \hat{\nabla}R\\|_* \leq \varepsilon$	$\mathcal{O}(L\Delta_0 / \varepsilon^2)$
Algorithm 1 (Muon特化)	任意非欧范数 / 谱范数	一阶光滑 (L)	闭凸正则项	$\mathbb{E}[\\|\nabla f + \hat{\nabla}R\\|_*] \leq \varepsilon$	$\mathcal{O}(L\Delta_0 \rho^2 \sigma^2 / \varepsilon^4)$
Algorithm 2 (+Weight Decay)	任意非欧范数 / 谱范数	星凸 (L, Star-Cvx)	闭凸正则项	$\mathbb{E}[F(x_K) - F(x^*)] \leq \varepsilon$	$\tilde{\mathcal{O}}(LD^3 \rho^2 \sigma^2 / \varepsilon^3)$
Algorithm 1 (+Weight Clip)	任意非欧范数 / 谱范数	星凸 (L, Star-Cvx)	有界域指示函数	$\mathbb{E}[F(x_K) - F(x^*)] \leq \varepsilon$	$\tilde{\mathcal{O}}(LD^4 \rho^2 \sigma^2 / \varepsilon^3)$
Algorithm 3 (+Extrapolation)	任意非欧范数 / 谱范数	二阶光滑 (L, H)	闭凸正则项	$\mathbb{E}[\\|\nabla f + \hat{\nabla}R\\|_*] \leq \varepsilon$	$\tilde{\mathcal{O}}(H^{1/2} \Delta_0 \rho^2 \sigma^2 / \varepsilon^{3.5})$
Algorithm 3 (+Star-Cvx+Extr)	任意非欧范数 / 谱范数	星凸 + 二阶光滑	闭凸正则项	$\mathbb{E}[F(x_K) - F(x^*)] \leq \varepsilon$	$\tilde{\mathcal{O}}(H^{1/2} D^{7/2} \rho^2 \sigma^2 / \varepsilon^{2.5})$

消融实验：不同机制对动量追踪误差与下降项的影响¶

通过理论引理剖析（Lemmas 4, 5, 8），各核心机制对动量误差界与能量单步下降的精确解析贡献如下：

配置 / 理论引理	动量逼近误差 $\mathbb{E}[\\|m_{k+1} - \nabla f(x_k)\\|_*]$	下降引理主要耗散项	核心机制作用解析
标准动量 (Lemma 5)	$(1-\alpha)^{k+1}\rho\sigma + \sqrt{\alpha}\rho\sigma + \frac{L\eta}{\alpha}$	$-\eta \\|\nabla f + \hat{\nabla}R\\|_* + \frac{3}{2}L\eta^2$	滑动平均降低随机方差，步长过大导致漂移
带中心平移 (Lemma 7)	$(1-\alpha)^{k+1}\rho\sigma + \sqrt{\alpha}\rho\sigma + \frac{2L\eta}{\alpha}$	$(1-\beta)(F_k - F^*) + 4L\eta^2$	引入收缩因子 $\beta$，保证序列半径 $\leq D$
二阶外推 (Lemma 8)	$(1-\alpha)^{k+1}\rho\sigma + \sqrt{\alpha}\rho\sigma + \frac{H\eta^2}{2\alpha^2}$	$-\eta \\|\nabla f + \hat{\nabla}R\\|_* + \frac{7}{2}L\eta^2$	外推点抵消一阶变动，线性漂移项化为二阶曲率项

关键发现¶

Muon天然适配神经网络隐藏层结构：在矩阵全连接层 $F(X) = \frac{1}{N}\sum_i \mathcal{L}_i(Xa_i)$ 中，谱范数算子范数自然具备紧致缩放 $\|Xa_i\|_2 \leq \|X\|_{\mathrm{op}}\|a_i\|_2$。由此导出的Lipschitz常数为 $L = \lambda \frac{1}{N}\sum_i \|a_i\|_2^2$，相比Frobenius欧氏范数极其紧致，从几何尺度上证明了谱范数置信域的理论优越性。
权重衰减与权重剪裁（Clipping）的理论等价性：在星凸假设下，带权重衰减的Algorithm 2与带紧致凸集指示正则（Clipping）的Algorithm 1均能达到 $\tilde{\mathcal{O}}(1/\varepsilon^3)$ 复杂度。在无穷范数下，Clipping易于逐元素实现（signSGD with clipping）；但在谱范数下，矩阵谱截断计算代价极高，因此中心平移式的权重衰减是Muon在大模型中控制谱发散的最优可行路径。

亮点与洞察¶

纠正了前人对正交梯度下降的理论误解：指出了Bernstein与Newhouse将其视为谱范数下Majorization-Minimization梯度下降所产生的核范数冗余缺陷，首次用一阶置信域优化严格闭合了理论与算法实现。
严密澄清了Muon与Orthogonal-SGDM的优劣根源：从随机分析角度揭示了动量必须作用在线性无偏梯度上以享受独立随机变量方差相消（batching effect）的原理，解释了为什么先正交后动量（OSGDM）会累积非线性畸变。

局限与展望¶

纯理论分析缺乏真实大模型经验验证：论文为纯理论工作，未在大规模LLM上实操评测所提出的带外推机制Muon（Algorithm 3）的实际吞吐与收敛开销。
星凸性假设与真实神经网络损失面的差距：虽然部分深度学习工作发现SGD沿某些路径呈现星凸特性，但极深多模态网络或LLM注意力层的真实非凸几何更加复杂，理论边界仍有拓宽空间。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次建立谱范数置信域与梯度正交化的精确等价性并统一动量与外推体系。
实验充分度: ⭐⭐⭐⭐ 理论推导完备闭环，各引理证明详尽，但缺乏大模型实操训练曲线。
写作质量: ⭐⭐⭐⭐⭐ 概念清晰，对比尖锐，数学逻辑严谨且物理意义直观明确。
价值: ⭐⭐⭐⭐⭐ 为Muon及次世代大模型非欧几何优化器提供了里程碑式的理论奠基。

配置 / 理论引理	动量逼近误差 \(\mathbb{E}[\\|m_{k+1} - \nabla f(x_k)\\|_*]\)	下降引理主要耗散项	核心机制作用解析
标准动量 (Lemma 5)	\((1-\alpha)^{k+1}\rho\sigma + \sqrt{\alpha}\rho\sigma + \frac{L\eta}{\alpha}\)	\(-\eta \\|\nabla f + \hat{\nabla}R\\|_* + \frac{3}{2}L\eta^2\)	滑动平均降低随机方差，步长过大导致漂移
带中心平移 (Lemma 7)	\((1-\alpha)^{k+1}\rho\sigma + \sqrt{\alpha}\rho\sigma + \frac{2L\eta}{\alpha}\)	\((1-\beta)(F_k - F^*) + 4L\eta^2\)	引入收缩因子 \(\beta\)，保证序列半径 \(\leq D\)
二阶外推 (Lemma 8)	\((1-\alpha)^{k+1}\rho\sigma + \sqrt{\alpha}\rho\sigma + \frac{H\eta^2}{2\alpha^2}\)	\(-\eta \\|\nabla f + \hat{\nabla}R\\|_* + \frac{7}{2}L\eta^2\)	外推点抵消一阶变动，线性漂移项化为二阶曲率项