乐观对偶平均化统一了现代优化器

# 乐观双重平均法统一了现代优化器

来源: https://arxiv.org/html/2605.11172
Thomas Pethick 独立研究员 [email protected] & Wanyun Xie EPFL (LIONS) [email protected] Roman Machacek 伯尔尼大学 [email protected] & Volkan Cevher EPFL (LIONS) [email protected]

###### 摘要

我们提出了 SODA，这是乐观双重平均法（Optimistic Dual Averaging）的一种广义形式。它为 Muon、Lion、AdEMAMix 和 NAdam 等最先进的优化器提供了一个统一的视角，表明它们都可以被视为该框架下的乐观实例。基于这一框架，我们提出了一种实用的 SODA 封装器，可用于任何基础优化器，通过理论上依据充分的 $1/k$ 衰减调度来消除对权重衰减（weight decay）的调优需求。在各种规模和训练周期内的实证结果表明，SODA 在无需任何额外超参数调优的情况下，能够 consistently 提升性能。

## 1 引言

深度学习优化主要沿两个 largely 互补的轴发展。第一个轴是*几何与自适应*：设计更新步骤，使其隐含范数、预条件子或约束集反映现代模型的结构，从而潜在改善高维情况下的缩放能力。Adam (Kingma and Ba, 2014) 是一个核心例子，它通过元素归一化将梯度平均与 $\ell_\infty$ 型几何耦合。与此同时，*谱*几何和随机对偶化在随机谱下降（Stochastic Spectral Descent, SSD）(Carlson et al., 2015a; 2016) 中得以发展，并最近在 Muon 和 Scion (Jordan et al., 2024b; Pethick et al., 2025a) 等方法中重新受到关注。

第二个轴是*成分组合*：如何通过动量、平均和调度将梯度反馈和迭代点结合起来，以获得稳定且高性能的训练配方。诸如 Lion (Chen et al., 2023) 的算法以及无调度（schedule-free）封装器 (Defazio et al., 2024) 等实用配方表明，这些成分的组合方式可能与其本身同等重要。过去十年的一个显著模式是，许多看似截然不同的新优化器都可以解释为这个二维设计空间中的不同点。

对于这两个轴，经典的对偶化框架提供了一个有用的视角，其中下一个迭代点是通过最小化由几何诱导项正则化的线性代理获得的：

$$
x^{k+1} \in \operatorname*{arg\,min}_{x\in\mathcal{X}} \gamma_k \langle d^k, x \rangle + h_k(x). \tag{1}
$$

这里，$\gamma_k > 0$ 是步长调度，$d^k$ 是（可能经过平均的）梯度反馈，$h_k$ 是（可能随时间变化的）正则化项/镜像映射，决定了几何结构。

例如，选择 $h_k(x) = \frac{1}{2}\|x-x^k\|_2^2$ 可恢复梯度下降，而非欧几里得的选择 $h_k$ 则产生归一化和感知几何的更新，这些更新对问题维度的依赖性更为有利。同时，*如何*构建 $d^k$（动量、平均、外推）以及*如何*平均迭代点（原始平均/无调度）对所得方法的性质有强烈影响，包括噪声鲁棒性、加速能力和任意时刻保证。

一个主要挑战是如何设置这些优化器的超参数，随着模型规模变得过于昂贵而无法调优，我们转而寻求可预测的缩放规则，这使得这一问题尤为重要。

一个特别具有挑战性的参数是权重衰减。在多周期训练的背景下，权重衰减有精确的表征，因为它限制了迭代点的范数，从而起到正则化作用 (Xie and Li, 2024; Pethick et al., 2025a)。然而，即使在单周期训练中，过拟合不是问题，权重衰减仍然 surprisingly 有益。

最近，开发出一种缩放规则，选择权重衰减为 $1/d$，其中 $d$ 为模型维度 (Xiao, 2024; Qiu et al., 2025)。然而，由于许多大规模实验遵循计算最优的训练范式，其中 $d$ 和训练周期 $n$ 是耦合的，很难分解有效依赖关系是在 $d$ 上还是在时间上。此外，通常需要在每个优化器基础上调整权重衰减，正如在 Lion 中观察到的那样，它比 AdamW 需要大得多的权重衰减 (Chen et al., 2023)。这引出了以下问题：

*是否有可能将近期进展的这两个轴建立在经典方法的基础上，并提供关于超参数的指导？*

表 1: $\bar{\lambda}_k=0$ 时 SODA 的实例（参见 第 3 节）.- 1 重新发现为 Simplified-AdEMAMix (Morwani et al., 2025)。

在本文中，我们认为算法开发的这两个轴并非独立的线索。为此，我们引入了 SODA，这是乐观双重平均法 (Rakhlin and Sridharan, 2013) 的一种广义形式，显式地耦合了*对偶处理*（梯度平均 + 乐观性）和*原始处理*（迭代点平均 + 原始外推）。所得框架提供了一个单一模板，将几种广泛使用的优化器作为特例恢复，并使注入非欧几里得几何（如 $\ell_\infty$ 或谱几何）到现代训练配方中变得 straightforward。

特别是，当其中一个平均参数简化时（我们的“现代化”机制），SODA 产生了一个围绕任何基础优化器的实用封装器，通过理论上依据充分的 $1/k$ 衰减消除权重衰减调优。

#### 贡献。

我们的贡献如下：

1. **通过乐观性和对偶化实现统一：** 我们表明 SODA 为几种最先进的优化器提供了统一的视角。最值得注意的，它在如此广泛的框架中将 Muon、Lion-$\mathcal{K}$ 和 NAdam 捕获为*乐观*实例（参见 表 1）。
2. **理论保证：** SODA 框架使我们能够从理论上推导出在实践中表现出色的超参数选择。我们的公式为权重衰减提供了新视角，不仅将其视为正则化器，而且视为一种原始平均形式。这导致了一种新的更新规则，将每次更新锚定在初始化处，并诱导 $1/k$ 权重衰减调度。
3. **一个简单的封装器：** 我们提出了一种围绕任何基础优化器（例如，Adam, Lion, Muon, Scion）的实用封装器，不引入新的超参数，并消除了调整权重衰减参数的需要。
4. **无需额外调优的实证结果：** 我们证明，当封装 Adam、Muon 和 Scion 时，在各种模型大小和训练周期内均 consistent 地提升性能，甚至优于经过调优权重衰减的基线。我们提出的框架实例化 SODA†，消除了调优权重衰减的需要，同时在大规模上带来更好的性能。

#### 局限性

这项工作专注于单周期设置。对于多周期训练，由于过拟合问题，可能需要有一个控制正则化强度的超参数。

## 2 预备知识

我们要解决以下最小化问题

$$
\min_{x\in\mathcal{X}} f(x) := \mathbb{E}_{\xi}[f(x,\xi)], \tag{2}
$$

这涵盖了许多机器学习问题，其中 $\xi$ 代表数据样本，$x$ 是被优化模型的参数，$f(x,\xi)$ 是损失函数。

#### 双重平均和 Fenchel 共轭

双重平均框架 (Nesterov, 2009) 及其变体通常依赖于 Fenchel 共轭将梯度信息映射回原始空间。函数 $h:\mathcal{X}\to\mathbb{R}\cup\{\infty\}$ 的 Fenchel 共轭定义为：

$$
h^*(d) = \sup_{x\in\mathcal{X}}\left\{\langle d,x\rangle - h(x)\right\},
$$

其中 $d\in\mathcal{X}^*$。次微分 $\partial h^*$ 等价于该共轭操作的 maximizers 集合：

$$
\partial h^*(d) = \operatorname*{argmax}_{x\in\mathcal{X}}\left\{\langle d,x\rangle - h(x)\right\}.
$$

这个恒等式是 Fenchel-Young 不等式 (Bauschke and Lucet, 2012) 的直接结果。然后，一般优化模板 (1) 可以紧凑地写为：

$$
x^{k+1} \in \partial h_k^*(-\gamma_k d^k), \tag{3}
$$

其中 $d^k$ 是过去随机梯度的平均（或动量），$\gamma_k$ 是步长调度，$h_k$ 是一系列定义几何的正则化器（或镜像映射）。

#### 标准正则化器

正则化器 $h(x)$ 的不同选择恢复了著名的优化步骤。我们突出显示三个特别与我们工作相关的实例，特别是当约束集是某些任意范数 $\|\cdot\|$ 的范数球 $\mathcal{D}=\{x\in\mathcal{X} \mid \|x\|\leq 1\}$ 时：

1. **无约束：** 令 $h(x) = \frac{1}{2}\|x\|^2$。那么，$\partial h^*$ 是 sharp 算子 (Nesterov, 2012)：
   $$
   \begin{split}
   \partial h^*(d) &= [d]^\sharp := \operatorname*{arg\,max}_{x}\left\{\braket{d,x} - \tfrac{1}{2}\|x\|^2\right\}.
   \end{split}
   $$
2. **线性最小化预言机（lmo）：** 令 $h(x) = \iota_{\mathcal{D}}(x)$ 为集合 $\mathcal{D}$ 的指示函数（即，如果 $x\in\mathcal{D}$ 则为 0，否则为 $\infty$）。那么 $\partial h^*$ 是 lmo：
   $$
   \begin{split}
   \partial h^*(-d) &= \operatorname{lmo}_{\mathcal{D}}(d) := \operatorname*{arg\,min}_{x\in\mathcal{D}} \braket{d,x}.
   \end{split}
   $$
   这是 Frank-Wolfe（条件梯度）方法的核心操作 (Frank and Wolfe, 1956; Clarkson, 2010; Jaggi, 2013)，这些方法是无投影的，但通常需要 $f$ 是光滑的。
3. **裁剪（Clipping）：** 通过组合上述两者，我们可以处理非光滑的 $f$ 通过“裁剪”。令 $h(x) = \frac{1}{2}\|x\|^2 + \iota_{\mathcal{D}}(x)$。那么，我们有
   $$
   \begin{split}
   \partial h^*(-d) &= \operatorname{clip}_{\mathcal{D}}(d) := \operatorname*{arg\,min}_{x\in\mathcal{D}} \left\{\braket{d,x} + \tfrac{1}{2}\|x\|^2\right\}.
   \end{split}
   $$
   此操作将梯度步裁剪到可行集 $\mathcal{D}$，本质上执行负梯度的投影。这在 Pethick et al. (2025b); Crawshaw et al. (2025) 中用于训练神经网络。

这三种 $\partial h^*$ 的选择都可以通过 lmo 联系起来，因为 $[d]^\sharp = -\|d\|_* \operatorname{lmo}_{\mathcal{D}}(d)$ 且 $\operatorname{clip}_{\mathcal{D}}(d) = \min\{1, \|d\|_*\} \operatorname{lmo}_{\mathcal{D}}(d)$（见 LABEL:lem:equiv）。

#### 范数选择

通过取 $h=\iota_{\mathcal{D}}$，其中 $\mathcal{D}$ 是范数的单位球，出现重要的特例。对于 $\ell_\infty$ 球，$\operatorname{lmo}_{\mathcal{D}}(d) = -\operatorname{sign}(d)$，如 SignSGD 和 Lion 中使用 (Bernstein et al., 2018; Chen et al., 2023)。对于谱范数球，$\operatorname{lmo}_{\mathcal{D}}(G) = -\operatorname{msign}(G)$，其中 $\operatorname{msign}(G) := UV^\top$，$U\Sigma V^\top$ 是 $G$ 的奇异值分解，如 Carlson et al. (2016); Jordan et al. (2024b); Pethick et al. (2025a) 中使用。

## 3 方法

我们提出以下算法，通过引入来自 Tseng (2008); Lan (2012); Defazio et al. (2024) 的原始外推序列 $y^k$，推广了乐观双重平均法（ODA）：

$$
\begin{split}
m^{k+1} &= (1-\alpha_k)m^k + \alpha_k \nabla f(y^k, \xi_k) \\
\bar{m}^{k+1} &= (1-\bar{\alpha}_k)m^{k+1} + \bar{\alpha}_k \nabla f(y^k, \xi_k) \quad \text{(optimism)} \\
z^{k+1} &\in \partial h_k^*(-\gamma_k \bar{m}^{k+1}) = \operatorname*{arg\,min}_{x\in\mathcal{X}} \gamma_k \braket{\bar{m}^{k+1}, x} + h_k(x) \\
x^{k+1} &= (1-\lambda_k)x^k + \lambda_k z^{k+1} \\
y^{k+1} &= (1-\bar{\lambda}_k)x^{k+1} + \bar{\lambda}_k z^{k+1} \qquad \text{(primal extrapolation)}
\end{split} \tag{SODA}
$$

其中 $\alpha_k, \bar{\alpha}_k, \lambda_k, \bar{\lambda}_k \in [0,1]$ 且 $\gamma_k > 0$。我们初始化 $m^0=0$ 和 $x^0=y^0=z^0 \in \partial h_0^*(0)$。观察对偶（梯度）和原始（迭代点）处理方式之间优雅的对称性。我们在 第 4 节 中的分析基于 Schedule-Free 框架...