Gefen：优化的随机优化器

# Gefen：优化随机优化器
来源：https://arxiv.org/html/2606.13894
Nadav Benedek Reichman University & Tomer Koren Tel Aviv University, Google Research & Ohad Fried Reichman University

###### 摘要

AdamW 是现代深度学习中的默认优化器，但其一阶和二阶矩状态为训练内存增加了大约两倍参数大小的缓冲区。我们提出 Gefen，一种内存高效的优化器，它自动在参数块之间共享二阶矩估计，并使用学习到的码本量化一阶矩，从而将 AdamW 的内存占用减少约 **8 倍**，同时保持相同的性能，相当于每十亿参数减少 6.5 GiB。该方法的动机源于一个理论结果，该结果表明，大型混合 Hessian 矩阵元素将平方梯度之比约束在 1 附近，这表明与 Hessian 对齐的参数是共享二阶矩统计量的自然候选者。由于在大规模场景下计算 Hessian 矩阵不可行，Gefen 从初始平方梯度推断块结构，除了 AdamW 默认超参数外，无需特定于架构的元数据或超参数。Gefen 学习一个精确的基于直方图的动态规划量化码本，并重用相同的块进行一阶矩缩放。在各种实验中，Gefen 在比较的类似 AdamW 的方法中实现了最低的峰值优化器内存，同时保持了 AdamW 级别的性能。在 FSDP 和 DDP 训练中，减少的内存占用使得能够使用更大的微批次，并显著提高相对于 AdamW 的吞吐量，从而提供一个实用的即插即用替代方案，具有更低的内存使用，可以增加吞吐量，并能够训练更大的模型或使用更大的批次大小。我们提供完整的 Python 实现，包括融合的 CUDA 内核，地址为 https://github.com/ndvbd/Gefen。

## 1 引言

Adam (Kingma & Ba, 2014 (https://arxiv.org/html/2606.13894#bib.bib15)) 和 AdamW (Loshchilov & Hutter, 2017 (https://arxiv.org/html/2606.13894#bib.bib24)) 是深度学习中广泛使用的主力优化器，因为它们在各种架构和数据集上都具有强大且一致的性能。然而，由于它们同时存储一阶和二阶矩的移动平均，仅其优化器状态就增加了大约两倍的参数内存。

减少优化器内存可以支持训练更大的模型或使用更大的批次大小，进而可以提高训练吞吐量和最终模型质量。因此，我们有动力设计一种优化器，它既能保持类似 AdamW 的性能，同时又能减少内存使用。

例如，Adam-mini (Zhang et al., 2025 (https://arxiv.org/html/2606.13894#bib.bib37)) 是 Adam 的一个近期变体，它通过使用模仿参数 Hessian 矩阵块对角结构的手动规则来共享二阶矩值，从而减少优化器内存。然而，为什么与 Hessian 对齐的分组是有益的，仍然缺乏理论依据。在实践中，用户可能需要提供特定于架构的信息（例如，注意力头数），并且分组规则依赖于张量名称，而张量名称可能因实现而异。量化优化器 (Dettmers et al., 2021 (https://arxiv.org/html/2606.13894#bib.bib8); Li et al., 2023 (https://arxiv.org/html/2606.13894#bib.bib18)) 能减少内存，但依赖于手动设计的次优码本，并通过量化块大小引入了一个隐式超参数。

受这些方法的启发，我们的目标是消除手动分组需求，避免 Adam-mini 依赖于张量名称的决策，为具有高 Hessian 相关性的参数分组提供更坚实的理论基础，在 LLM 和非 LLM 网络上进行验证，并在保持通用性和 AdamW 级别性能的同时，进一步降低内存占用，超越 Adam-mini。

在 LABEL:sec:desiderata 中，我们描述了一个理想优化器应具备的特性，用于指导我们的设计和评估。在第 3 节 (https://arxiv.org/html/2606.13894#S3) 中，我们提供了理论分析，表明高 Hessian 亲和性意味着相似的平方梯度，这证明了我们的分组策略是合理的。在第 4 节 (https://arxiv.org/html/2606.13894#S4) 中，我们介绍了我们的算法，该算法自动分组参数并在这些组内共享二阶矩。实验证明，我们的方法在各种模型和数据集上，在比较的方法中实现了最低的内存占用，同时保持可比或更优的质量和吞吐量，如第 5 节 (https://arxiv.org/html/2606.13894#S5) 所示。熟悉 AdamW 的读者可以直接跳到第 3 节 (https://arxiv.org/html/2606.13894#S3)。

## 2 预备知识与相关工作

##### Adam 和 AdamW。

结合 AdaGrad (Duchi et al., 2011 (https://arxiv.org/html/2606.13894#bib.bib9)) 和 RMSProp (Tieleman & Hinton, 2012 (https://arxiv.org/html/2606.13894#bib.bib35)) 的思想，Adam 是一种自适应一阶优化器，它维护梯度和平方梯度的指数移动平均。在步骤 \(t\)，给定梯度 \(g_t = \nabla_{\theta} \mathcal{L}(\theta_t)\)，Adam 计算

\[
\begin{aligned}
m_t &= \beta_1 m_{t-1} + (1-\beta_1) g_t \quad (\text{一阶矩的 EMA；动量项}) \quad (1) \\
v_t &= \beta_2 v_{t-1} + (1-\beta_2) g_t^2 \quad (\text{二阶原始矩的 EMA}) \quad (2)
\end{aligned}
\]

然后应用偏差校正，因为两个 EMA 都初始化为零，因此在早期步骤中偏向于零，

\[
\hat{m}_t = \frac{m_t}{1-\beta_1^t}, \quad \hat{v}_t = \frac{v_t}{1-\beta_2^t} \quad (\text{偏差校正}) \quad (3)
\]

并更新参数为

\[
\theta_{t+1} = \theta_t - \eta \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon} \quad (4)
\]

AdamW 使用相同的自适应矩，但直接在参数空间应用权重衰减，而不是在梯度中注入 \(L_2\) 惩罚项：

\[
\theta_{t+1} = \theta_t - \eta \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon} - \eta \lambda \theta_t \quad (5)
\]

这种解耦避免了自适应梯度缩放和正则化强度之间的相互作用，通常能产生更好的泛化能力和更稳定的超参数调优。在现代深度学习实践中，Adam 和 AdamW 是事实上的标准优化器。它们通常以 FP32 格式保存一阶和二阶矩状态，因此优化器状态内存大约为每个参数 8 字节。例如，对于一个 13B 参数的模型，Adam 风格的优化器状态单独就需要大约 \(13 \times 10^9 \times 8 \approx 97\) GiB；这可能会使在内存受限的 GPU 上进行训练变得更加困难，并迫使使用更小的批次大小，这也会降低训练吞吐量。

##### 量化优化器状态。

最近的工作通过量化 Adam 风格的矩状态来减少优化器内存。Adam8bit (Dettmers et al., 2021 (https://arxiv.org/html/2606.13894#bib.bib8)) 将张量展平为固定大小的块 (\(B=2048\))，量化一阶和二阶矩，并存储每个块的绝对最大值缩放（每一步重新计算）以及每个元素的 8 位值。其动态树量化将归一化值编码在 \([-1, 1]\) 范围内，在接近零时提供更高精度，同时仍然能表示较大的量级。然而，这种设计存在实际局限性：(1) 动态树量化器可能由于重复的码本条目而造成浪费，(2) 在实践中该方法不适用于所有张量类型（例如，嵌入不应量化），因此它不是一个完全即插即用的通用替代方案，以及 (3) 固定的用户选择块大小引入了额外的超参数负担。Adam4bit (Li et al., 2023 (https://arxiv.org/html/2606.13894#bib.bib18)) 将状态量化推进到 4 位，并改进了对异常值的处理（包括对矩统计量的更细粒度和非对称处理），在保持许多设置中具有竞争力的准确性的同时，进一步减少了内存。该方法使用线性量化，并包含一个块大小超参数，在作者的实现中默认为 128。

详细的扩展相关工作的讨论，包括其他优化器和梯度通信压缩方法，请参见 LABEL:app:extended_related。在这项工作中，我们专注于通用的优化器，它们无需手动配置、额外超参数或学习率重新调整，并且可以在跨训练场景中作为 AdamW 的即插即用替代品，同时实现显著更低的内存占用。

## 3 理论分析

AdamW 为每个参数维护平方梯度的指数移动平均。我们分析在什么条件下，两个具有大幅值混合 Hessian 矩阵元素的权重也具有相似的平方梯度幅值，从而使它们成为共享二阶矩估计的自然候选者。具体来说，对于参数条目索引为 \((k,l)\) 和 \((k', l')\)，我们研究何时

\[
\left(\frac{\partial L}{\partial W[k,l]}\right)^2 \quad \text{和} \quad \left(\frac{\partial L}{\partial W[k',l']}\right)^2
\]

是接近的。这样的关系激发了将具有大 Hessian 耦合的参数对分组并共享它们的二阶矩估计的想法。

###### 定理 3.1 (大 Hessian 条目压缩平方梯度比率)。

该定理表明，对于可行的 \(\|H\| \leq A\)，增加混合 Hessian 的幅度会将平方梯度比率的可接受范围收紧到接近 1。对于共享相同输出坐标的第二层权重以及隐藏偏置项，也观察到相同形式的压缩，而输出偏置的 Hessian 条目要么为零，要么为常数。上述理论推导是针对 LABEL:app:theory-proof 中一个简单的两层 MLP 给出的。然而，在来自真实网络架构的张量中，包括注意力张量和卷积层，也经验性地观察到了相同的行为，如附录 A (https://arxiv.org/html/2606.13894#A1) 所示。

## 4 算法

算法 1 Gefen：我们提出的优化随机优化算法。默认设置与 AdamW 相同。所有向量的操作都是逐元素的。该算法拥有融合的 CUDA 实现，以避免不必要的临时内存分配。
1: **要求：**
    \(\alpha, \lambda \in \mathbb{R}\)：步长和权重衰减，
    \(\beta_1, \beta_2 \in [0,1)\)：矩估计的衰减率。
    \(f(\theta)\)：参数为 \(\theta\) 的随机目标函数。
    初始化
    \(t \leftarrow 0\),
    \(\bm{G}_0 = \nabla_{\theta} f_t(\bm{\theta}_0)\)
2: **对每个** 参数张量 \(\theta\) **执行** [自动块划分](https://arxiv.org/html/2606.13894#alg2) 使用 \(\bm{G}_0\)
3: **执行** [精确 DP 量化码本学习](https://arxiv.org/html/2606.13894#alg3)。
4: **对每个** 参数张量 \(\theta\) **执行**
5:    初始化量化一阶动量码本 \(\bar{\bm{m}}_0 \leftarrow \bm{0}\) 每个权重，和 \(\|\bm{m}_0\|_{\infty}\) 每个块
6:    初始化一个二阶动量标量 \(\bar{\bm{v}}_0 \leftarrow 0\) 每个块
7: **重复**
8:    \(t \leftarrow t + 1\)
9:    **对每个** 参数张量 \(\theta\) **执行**
10:       \(\bm{g}_t \leftarrow \nabla_{\theta} f_t(\bm{\theta}_{t-1})\)  ▷ 获取当前参数的梯度
11:       \(\bm{m}_t \leftarrow \beta_1 \cdot \textsc{dequantize}(\bar{\bm{m}}_{t-1}, \|\bm{m}_{t-1}\|_{\infty}) + (1-\beta_1) \cdot \bm{g}_t\)  ▷ 更新有偏动量估计
12:       \(\bar{\bm{m}}_t, \|\bm{m}_t\|_{\infty} \leftarrow \textsc{Quantize}(\bm{m}_t)\)  ▷ 量化并重新计算每个块的绝对值最大值
13:       \(\hat{\bm{m}}_t \leftarrow \bm{m}_t / (1 - \beta_1^t)\)  ▷ 动量的偏差校正
14:       \(\bm{v}_t \leftarrow \beta_2 \cdot \bm{v}_{t-1} + (1-\beta_2) \cdot \text{mean}(\bm{g}^2_t)\)  ▷ 使用每个块的均值更新有偏二阶矩
15:       \(\hat{\bm{\{v\}}}_t \leftarrow \bm{v}_t / (1 - \beta_2^t)\)  ▷ 紧凑二阶矩估计的偏差校正
16:       \(\bm{\theta}_t \leftarrow \bm{\theta}_{t-1} - \alpha \cdot \left( \hat{\bm{m}}_t / ( \sqrt{ \hat{\bm{v}}_t } + \epsilon) + \lambda \bm{\theta}_{t-1} \right)\)  ▷ 紧凑的 \(\hat{\bm{\{v\}}}_t\) 广播到完整形状
17: **直到** 满足停止准则
18: **返回** 优化后的参数 \(\bm{\theta}_t\)

算法 2 自动块划分
1: **要求：** 参数张量的一步梯度 \(\bm{g}_0\)，\(n = |\bm{g}_0|\)
2: 展平 \(\bm{g}_0\)，令 \(\mathcal{P}\) 为 \(n\) 的所有真因子，按递增排序
3: \(E_{\mathrm{prev}} \leftarrow \emptyset\)，\(\Delta^{\star} \leftarrow 10^{-12}\)，\(p^{\star} \leftarrow 1\)
4: **对每个** 候选周期 \(p \in \mathcal{P}\) **执行**
5:    将 \(\bm{g}_0^2\) 重塑为 \(\bm{B} \in \mathbb{R}^{(n/p) \times p}\)
6:    \(E(p) \leftarrow \sqrt{ \operatorname{mean}_i \left( \mathrm{Var}(\bm{B}_{i,:}) \right) }\)  ▷ 公式 (6) (https://arxiv.org/html/2606.13894#S4.E6)
7:    **如果** \(E_{\mathrm{prev}} \neq \emptyset\) **则**
8:        \(\Delta \leftarrow E(p) - E_{\mathrm{prev}}\)
9:        **如果** \(\Delta < \Delta^{\star}\) **则**
10:           \(\Delta^{\star} \leftarrow \Delta\)，\(p^{\star} \leftarrow p\)
11:   \(E_{\mathrm{prev}} \leftarrow E(p)\)
12: **返回** \(p^{\star}\)，如果 \(p^{\star} \geq 8\)，否则返回 \(1\)

算法 3 精确 DP 量化码本学习
1: **要求：** 梯度 \(\{\bm{g}^{(\theta)}_0\}\)，自动周期 \(\{p_{\theta}\}\)，码本条目数 \(k\)
2: 在 \([-1, 1]\) 上初始化直方图计数 \(\bm{c} \in \mathbb{R}^{16k}\)
3: **对每个** 参数张量 \(\theta\) **执行**
4:    将 \(\bm{g}^{(\theta)}_0\) 重塑为块 \(\bm{B} \in \mathbb{R}^{(n_{\theta}/p_{\theta}) \times p_{\theta}}\)
5:    按每个块的最大绝对值进行归一化：\(\bm{Z}_{i,:} \leftarrow \bm{B}_{i,:} / \|\bm{B}_{i,:}\|_{\infty}\)
6:    将 \(\bm{Z}\) 的所有条目累加到直方图计数 \(\bm{c}\) 中
7: 令 \(\{(m_i, c_i)\}_{i=1}^b\) 为非空直方图桶中心和计数，按 \(m_i\) 排序
8: 计算 \(c_i\)、\(c_i m_i\) 和 \(c_i m_i^2\) 的前缀和
9: 使用中心 \(-1\)（如果 \(k'=1\)）、中心 \(1\)（如果 \(k'=k\)），否则使用桶 \(b_{\ell}, ..., b_r\) 的加权均值作为码本中心，定义区间平方误差代价 \(C(b_{\ell}, b_r, k')\)  ▷ 强制极值码本条目
10: \(D_{b_r, k'} :=\) 使用前 \(k'\) 个码点量化前 \(b_r\) 个桶的代价值；初始化 \(D_{0,0} \leftarrow 0\)，所有其他条目为 \(\infty\)