通过隐式梯度传输加速基于 LMO 的优化

# 通过隐式梯度传输加速基于 LMO 的优化

来源: https://arxiv.org/html/2605.05577
Won-Jun Jang Si-Hyeon Lee
韩国科学技术院 (KAIST) 电气工程学院
韩国大田
{wonjun_jang,sihyeon}@kaist.ac.kr

###### 摘要

Lion 和 Muon 等最新优化器通过线性最小化预言机（LMOs）对梯度动量进行归一化，展现出了强大的实证性能。虽然方差缩减已被探索用于加速基于 LMO 的方法，但它通常因额外的梯度评估而产生巨大的计算开销。同时，基于 LMO 的方法的理论理解在无约束和有约束公式中仍然碎片化。受这些局限性的启发，我们提出了 *LMO-IGT*，这是一类利用隐式梯度传输（IGT）的新型随机基于 LMO 的方法。我们进一步引入了一个统一的随机基于 LMO 的优化框架以及一个新的平稳性度量——*正则化支撑函数*（RSF），它在统一框架下桥接了梯度范数和 Frank-Wolfe 间隙概念。通过在传输点上评估随机梯度，LMO-IGT 在保留标准随机 LMO 的单梯度每次迭代结构的同时加速了收敛。我们的分析确立：随机 LMO 达到 $\mathcal{O}(\varepsilon^{-4})$ 的迭代复杂度，方差缩减 LMO 以额外梯度评估为代价达到 $\mathcal{O}(\varepsilon^{-3})$，而 LMO-IGT 仅使用每次迭代一个随机梯度就达到 $\mathcal{O}(\varepsilon^{-3.5})$。在实证上，LMO-IGT 始终优于随机 LMO 对应方法，且开销可忽略不计。在其具体实例中，Muon-IGT 在所有评估设置中实现了最强的整体性能，表明 IGT 为现代基于 LMO 的优化提供了一种有效且实用的加速机制。

## 1 引言

近年来，深度学习取得了显著的成功，增加了对训练大规模神经网络可靠优化方法的需求。训练此类模型涉及高维和非凸优化问题，使得高效且稳定的算法至关重要。因此，高效的基于随机梯度的优化器如 Adam (Kingma and Ba, 2017) 和 AdamW (Loshchilov and Hutter, 2019) 长期以来一直是标准选择。最近，Lion (Chen et al., 2023) 和 Muon (Jordan et al., 2024) 等新优化器因其强大的实证性能而引起了广泛关注。这些方法的一个共同特征是它们在优化过程中*归一化*梯度动量，这一特性与线性最小化预言机（LMO）密切相关。

此后，一系列工作致力于改进基于 LMO 的方法的收敛性，其中诸如方差缩减等技术实现了更快的理论速率 (Jiang and Zhang, 2025; Yuan et al., 2025; Chang et al., 2025; Sfyraki and Wang, 2025)。然而，这些方法通常需要进行额外的梯度评估，这在大规模设置中计算成本高昂。这引出了一个自然的问题：基于 LMO 的方法能否在避免与方差缩减相关的巨大计算开销的同时实现加速收敛？

从理论角度来看，LMO 起源于用于有约束优化的 Frank-Wolfe 方法 (Jiang and Zhang, 2025; Sfyraki and Wang, 2025; Pethick et al., 2025)。然而，LMO 类型更新在现代随机优化中的作用尚未被充分理解。特别是，现有分析通常对无约束和有约束问题依赖不同的平稳性度量，阻碍了跨设置统一处理。这导致了第二个问题：随机基于 LMO 的优化能否在一个统一框架下得到理解？

我们通过以下贡献来回答上述问题。

**基于 LMO 的低开销优化加速。** 为了加速基于 LMO 的方法同时避免方差缩减的巨大计算开销，我们从隐式梯度传输（IGT）中汲取灵感，后者利用在前瞻点上评估的梯度来构建动量 (Arnold et al., 2019; Cutkosky and Mehta, 2020)。虽然 IGT 已在归一化梯度方法中进行过研究，但我们将其扩展到基于 LMO 的优化。我们观察到 LMO 更新本质上执行了一种形式的动量归一化，使其非常适合 IGT 风格的加速。基于这一见解，我们提出了一类新方法，称为 LMO-IGT，它在无需额外梯度评估的情况下实现加速收敛。

**统一框架与分析。** 我们开发了随机基于 LMO 的优化的统一框架，涵盖三类方法：随机 LMO、方差缩减 LMO（LMO-VR）以及提出的 LMO-IGT。为了在有无约束设置之间实现统一分析，我们引入了一种新的平稳性度量——*正则化支撑函数（RSF）*，它桥接了基于梯度的和 Frank-Wolfe 类型的度量。在此框架下，我们建立了统一的收敛结果：随机 LMO 达到 $\mathcal{O}(\epsilon^{-4})$ 的迭代复杂度，LMO-VR 达到 $\mathcal{O}(\epsilon^{-3})$，而 LMO-IGT 达到 $\mathcal{O}(\epsilon^{-3.5})$，同时仅产生适度的计算开销。

**实证验证。** 我们用不同的 LMO 集实例化 LMO-IGT，产生了如 Lion-IGT 和 Muon-IGT 等实用算法。特别是，Muon-IGT 在无需额外梯度评估的情况下实现了最强的实证性能，证明了所提方法在大规模设置中的有效性。表 1 总结了我们在统一框架中的三类方法，连同代表性算法、相关的 LMO 集及其相应的迭代复杂度。

**表 1：** 随机基于 LMO 的优化的统一框架及相应的收敛结果。对于梯度动量矩阵 $g$，我们将其奇异值分解（SVD）写为 $g=U\text{diag}(\sigma)V^{\top}$。参见第 II 节了解 LMO 和集合 $\mathcal{C}$ 的定义。

| 类别 | 算法 | LMO 集 $\mathcal{C}$ | $\text{LMO}_{\mathcal{C}}(g)$ | 迭代复杂度 |
| :--- | :--- | :--- | :--- | :--- |
| 随机 LMO<br>(1 次梯度评估) | 归一化 SGD (Hazan et al., 2015) | $\|\cdot\|_2$-球 | $-g/\|g\|$ | $\mathcal{O}(\epsilon^{-4})$ |
| | signSGD (Bernstein et al., 2018) | $\|\cdot\|_\infty$-球 | $-\text{sign}(g)$ | |
| | Signum (Bernstein et al., 2018) | | | |
| | Lion (Chen et al., 2023) | | | |
| | Muon (Jordan et al., 2024) | $\|\cdot\|_\text{op}$-球* | $-UV^{\top}$ | |
| LMO-VR<br>(2 次梯度评估) | Lion-VR (Jiang and Zhang, 2025) | $\|\cdot\|_\infty$-球 | $-\text{sign}(g)$ | $\mathcal{O}(\epsilon^{-3})$ |
| | MARS-Shampoo (Yuan et al., 2025) | $\|\cdot\|_\text{op}$-球* | $-UV^{\top}$ | |
| | Muon-VR (Sfyraki and Wang, 2025) | | | |
| LMO-IGT<br>(1 次梯度评估) | NIGT (Cutkosky and Mehta, 2020) | $\|\cdot\|_2$-球 | $-g/\|g\|$ | $\mathcal{O}(\epsilon^{-3.5})$ |
| | Lion-IGT (本文) | $\|\cdot\|_\infty$-球 | $-\text{sign}(g)$ | |
| | Muon-IGT (本文) | $\|\cdot\|_\text{op}$-球* | $-UV^{\top}$ | |

*矩阵的算子范数球：$\mathcal{C}=\{X \mid \sigma_1(X) \le 1\}$。

## 2 问题设定与预备知识

令 $\xi$ 表示捕获预言机随机性的随机变量，令 $f(\cdot;\xi)$ 表示对应的样本损失函数。我们考虑随机一阶最小化问题

$$
\min_{w \in \mathcal{P}} F(w), \quad F(w) := \mathbb{E}_\xi[f(w;\xi)], \tag{1}
$$

其中 $F:\mathbb{R}^d \to \mathbb{R}$ 和 $f$ 是非凸、可微函数。我们假设可行集 $\mathcal{P} \subseteq \mathbb{R}^d$ 是凸且紧致的。我们做出如下假设：

###### 假设 1（L-平滑性）

函数 $F$ 是 L-平滑的：

$$
\|\nabla F(x) - \nabla F(y)\| \le L \|x - y\|, \quad \forall x, y \in \mathbb{R}^d.
$$

###### 假设 2（有界方差）

随机梯度具有有界方差：

$$
\mathbb{E}_\xi \left[ \|\nabla f(w;\xi) - \nabla F(w)\|^2 \right] \le \sigma^2, \quad \forall w \in \mathbb{R}^d.
$$

###### 假设 3（平均平滑性）

存在 $L > 0$ 使得

$$
\mathbb{E}_\xi \left[ \|\nabla f(x;\xi) - \nabla f(y;\xi)\|^2 \right] \le L^2 \|x - y\|^2, \quad \forall x, y \in \mathbb{R}^d.
$$

###### 假设 4（二阶平滑性）

$F$ 的 Hessian 是 $\rho$-Lipschitz 连续的：

$$
\|\nabla^2 F(x) - \nabla^2 F(y)\|_\text{op} \le \rho \|x - y\|, \quad \forall x, y \in \mathbb{R}^d.
$$

假设 1 和 2 在随机优化中是标准的。假设 3 常用于方差缩减方法，而假设 4 是 IGT 分析中的标准二阶正则性条件 (Cutkosky and Mehta, 2020)。

在本节的其余部分，我们将介绍必要的预备知识。我们首先描述基于 LMO 的优化方法的更新规则，随后是其 Frank-Wolfe 解释及相关的平稳性度量——Frank-Wolfe 间隙。然后我们总结随机 LMO 方法的收敛行为，并讨论作为加速技术的方差缩减。最后，我们提出 IGT 作为一种替代加速机制。相关工作的更详细讨论见附录 A。

### 2.1 基于 LMO 的更新

为了解决 (1)，归一化 SGD 和基于 LMO 的优化方法（例如，Lion 和 Muon）计算，给定梯度或动量估计 $g$，线性最小化预言机（LMO）

$$
\operatorname{LMO}_{\mathcal{C}}(g) \in \arg\min_{v \in \mathcal{C}} \langle g, v \rangle.
$$

这里，$\mathcal{C} \subseteq \mathbb{R}^d$ 是包含原点的紧凸集。等价地，预言机返回 $\mathcal{C}$ 中与下降方向 $-g$ 最对齐的点。$\operatorname{LMO}_{\mathcal{C}}(\cdot)$ 的具体形式取决于 $\mathcal{C}$。例如，如果 $\mathcal{C}$ 是 $\ell_2$ 范数球，则 $\operatorname{LMO}_{\mathcal{C}}(g) = -g/\|g\|$，这恢复了归一化 SGD。对于矩阵值 $g$，选择 $\mathcal{C}$ 为谱范数球，即 $\mathcal{C}=\{X : \|X\|_\text{op} \le 1\}$，产生的更新对应于 Muon。

在每次迭代 $t=0,1,\dots,T-1$，参数 $w_t$ 通过以下方式更新：

$$
w_{t+1} = (1 - \lambda \eta_t) w_t + \eta_t v_t, \quad v_t = \operatorname{LMO}_{\mathcal{C}}(g_t), \tag{2}
$$

其中 $\lambda \ge 0$ 是权重衰减参数，$\eta_t > 0$ 是学习率。当 $\lambda=0$ 时，问题简化为具有 $\mathcal{P}=\mathbb{R}^d$ 的无约束情况。当 $\lambda > 0$ 时，更新允许自然的有约束优化解释，我们将在下面讨论。

### 2.2 Frank-Wolfe 解释

Frank-Wolfe（条件梯度）方法旨在针对凸且紧致的集合 $\mathcal{P}$ 进行有约束优化。在迭代 $t$，给定 $w_t$、随机梯度 $g_t$ 和步长 $\gamma_t \in (0,1)$，更新为

$$
w_{t+1} = (1 - \gamma_t) w_t + \gamma_t \operatorname{LMO}_{\mathcal{P}}(g_t),
$$

通过在 $\mathcal{P}$ 中形成凸组合来保持可行性。当 $\lambda > 0$ 时，(2) 中的基于 LMO 的更新可以解释为 Frank-Wolfe 更新 (Chen et al., 2024; Sfyraki and Wang, 2025)。由于 $\mathcal{C}$ 是紧凸的，缩放后的集合 $\mathcal{P} := \lambda^{-1} \mathcal{C}$ 也是紧凸的。定义 $s_t := \lambda^{-1} v_t$，其中 $v_t = \operatorname{LMO}_{\mathcal{C}}(g_t)$，得到 $s_t = \operatorname{LMO}_{\mathcal{P}}(g_t) \in \mathcal{P}$。因此，更新可以重写为

$$
w_{t+1} = (1 - \lambda \eta_t) w_t + \eta_t v_t = (1 - \gamma_t) w_t + \gamma_t s_t,
$$

其中 $\gamma_t = \lambda \eta_t$，这与 Frank-Wolfe 更新匹配。

针对 $\mathcal{P}$ 上有约束优化的标准平稳性度量是 Frank-Wolfe 间隙 (Jaggi, 2013)：

$$
G_{\mathcal{P}}(w) := \sup_{u \in \mathcal{P}} \langle -\nabla F(w), u - w \rangle.
$$

间隙满足 $G_{\mathcal{P}}(w) = 0$ 当且仅当 $w$ 是 $\mathcal{P}$ 上的一阶平稳点 (Lacoste-Julien, 2016)。此外，其收敛到零等价于满足 KKT 条件 (Xie and Li, 2024; Sfyraki and Wang, 2025)。

### 2.3 收敛分析

从无约束设置中随机非凸优化的角度来看，最近几项工作在标准平滑性假设下分析了 Lion/Muon 类型方法的收敛性 (Don et al., 2024; Jiang and Zhang, 2025; Shen et al., 2025; Chang et al., 2025)。