MGUP:一种用于随机优化的动量-梯度对齐更新策略

# MGUP：面向随机优化的动量-梯度对齐更新策略

**来源**: https://arxiv.org/html/2606.17526

Da Chang<sup>134</sup>, Ganzhao Yuan<sup>21</sup>

1. 中国科学院深圳先进技术研究院
2. 深圳理工大学
3. 鹏城实验室
4. 中国科学院大学

###### 摘要

高效优化对于训练大语言模型至关重要。尽管层内选择性更新已有探索，但仍缺乏一种既能实现细粒度控制又能保证收敛性的通用机制。为填补这一空白，我们提出MGUP，一种新颖的选择性更新机制。MGUP 增强标准基于动量的优化器，在每次迭代中对选定的固定比例参数施加较大的步长，而对其余参数施加较小的非零步长。作为一个近乎即插即用的模块，MGUP 可无缝集成到 AdamW、Lion 和 Muon 等优化器中，从而产生强大的变体，如 MGUP-AdamW、MGUP-Lion 和 MGUP-Muon。在标准假设下，我们为随机优化中的 MGUP-AdamW（无权重衰减）提供了理论收敛保证。在包括 MAE 预训练、LLM 预训练和下游微调在内的多种任务上进行的大量实验表明，我们增强后的 MGUP 优化器相比其原始基础优化器，实现了更优或更稳定的性能。我们提供了一种有原则、通用且具有理论基础的层内选择性更新策略，以加速并稳定大规模模型的训练。代码公开发布于：https://github.com/MaeChd/MGUP。

## 1 引言

近期研究表明，大语言模型训练过程中的学习矩阵具有低秩性质，表明学习主要发生在低维空间中（Gur-Ari et al., 2018；Larsen et al., 2022）。这一发现催生了诸如 Galore (Zhao et al., 2024) 和 LDAdam (Robert et al., 2025) 等方法，它们利用梯度低秩分解，在降低内存消耗的同时实现了与全秩更新相当的性能。虽然低秩性质并不直接意味着稀疏性，但优化发生在低维空间这一见解为选择性参数更新提供了重要基础。这一原理在 SIFT (Song et al., 2024) 中得到体现，该方法利用基于梯度的稀疏参数更新，利用大语言模型中固有的低维度和稀疏梯度特性实现高效适配。在此基础上，涌现了几种创新的逐层选择性更新方法，包括 AutoFreeze (Liu et al., 2021)、LOMO (Lv et al., 2023)、LISA (Pan et al., 2024) 和 BAdam (Luo et al., 2024)。这些方法通过策略性地冻结某些层而更新其他层，实现了与全参数更新相当甚至更优的性能。

虽然逐层选择性更新显示出潜力，但更细粒度的参数选择仍未得到充分探索。尽管 SIFT (Song et al., 2024) 研究了稀疏的层内更新，但仍缺乏一种系统的方法来识别每层中最关键的参数。这一研究空白催生了开发新型层内稀疏更新策略的需求。最近，Liang et al. (2024) 提出了谨慎优化器，一种新颖的层内稀疏更新策略。该方法仅选择性地更新动量和梯度对齐的参数（即 \(\mathbf{m}_t \odot \mathbf{g}_t > 0\)），从而在对齐方向进行较大更新，同时跳过未对齐的方向。从概念上讲，它扩展了早期自适应优化器如 AdaBelief (Zhuang et al., 2020)，后者使用 \((\mathbf{m}_t - \mathbf{g}_t)^2\) 调整步长，但引入了基于动量-梯度对齐的参数选择。然而，这两种方法都有明显的局限性。AdaBelief 的更新机制严重依赖 Adam 的二阶矩估计，这限制了其应用于不计算二阶矩的优化器（例如，Lion (Chen et al., 2023) 或 Muon (Jordan et al., 2024)）。此外，谨慎优化器在随机设置下缺乏严格的理论收敛保证。尽管该策略在确定性情况下提供了理论见解，但其在随机条件下的收敛性质尚未得到验证。

这就提出了一个关键问题：在随机优化设置中，基于动量-梯度方向一致性的层内稀疏更新概念是否真正能作为即插即用的机制？如果是，其有效性的边界在哪里？如果不是，背后的原因是什么？我们在第 4 节的理论分析中详细探讨了这个问题。具体而言，我们证明，对于应用掩码的 Adam 变体，简单地将动量和梯度方向未对齐的参数的更新步长设为零，会显著影响随机优化的收敛性质。这促使我们重新思考如何在随机优化设置中更有效地执行选择性参数更新，以保持良好的收敛性质。例如，如果没有引导式的参数选择，可能出现某些极端情况：(i) 只有一小部分参数获得实质性更新（可能导致训练不稳定），或 (ii) 绝大多数参数的更新被过度抑制（可能导致训练缓慢）。因此，我们认为一个有前途的策略不仅要考虑动量与梯度方向的对齐，还要调节接收大量更新和少量更新的参数比例，以在训练效率和稳定性之间取得平衡。

受我们的理论分析和由此产生的设计考虑启发，我们引入一种新的选择性更新方法：MGUP（动量-梯度对齐更新策略）。MGUP 通过对元素级乘积 \(\mathbf{m}_t \odot \mathbf{g}_t\) 的值进行排序，选择性地差异化更新参数。具体来说，根据 \(\mathbf{m}_t \odot \mathbf{g}_t\) 排序的前 \(K\) 个参数接收缩放后的步长 \(\alpha \cdot \eta_t\)（\(\alpha > 1\)），而其余参数接收 \(\gamma \cdot \eta_t\)（\(\gamma < 1\)），其中 \(\eta_t\) 是原始优化器的基步长。MGUP 受谨慎更新策略启发，并依据 AdaBelief 和谨慎优化器的原理进行改进，通过动态调整基于动量-梯度对齐的更新强度。

我们的贡献总结如下：

- • 我们开发了一种新颖的选择性参数更新机制 MGUP，它为一部分参数分配较大的步长，为其余参数分配较小的步长。作为一种即插即用机制，MGUP 可以集成到基于动量的优化器如 AdamW、Lion 和 Muon 中，从而产生我们称之为 MGUP-AdamW、MGUP-Lion 和 MGUP-Muon 的变体。
- • 我们建立了带有 MGUP 机制的 Adam 优化器在随机设置下的收敛性，为其可靠性提供了理论保证。
- • 我们通过关键实验验证了所提出的 MGUP 优化器，包括在 CIFAR-10 上对 ViT-27M 进行 MAE 预训练；在 Wikitext-103 上对 LLaMA2-71M 和 Qwen2.5-150M 进行自回归预训练；以及在 GLUE 和 GSM-8K 上分别对 RoBERTa-base 和 LLaMA2-7B 进行微调。这些结果显示 MGUP 在多种模型和任务上具有鲁棒性和通用性。

## 2 相关工作

在本节中，我们回顾与动量-梯度方法相关的随机优化方法的基本原理。我们考虑最小化如下目标函数：
\[
\min_{\mathbf{x} \in \mathbb{R}^d} \, f(\mathbf{x}), \quad \text{其中} \quad f(\mathbf{x}) = \mathbb{E}_{\xi \sim \mathcal{D}} [f(\mathbf{x}; \xi)].
\tag{1}
\]
这里，\(f: \mathbb{R}^d \rightarrow \mathbb{R}\) 是一个可微且可能非凸的函数，\(\xi\) 表示从未知数据分布 \(\mathcal{D}\) 中采样的随机向量（如训练数据点）。在求解问题 (1) 的背景下，基于动量的方法是大规模机器学习优化的基础，它累积过去的梯度信息以加速收敛并导航复杂的损失景观。标准动量更新，即梯度的指数加权移动平均 (EWMA)，由下式给出：
\[
\mathbf{m}_t = \beta_1 \mathbf{m}_{t-1} + (1 - \beta_1) \mathbf{g}_t,
\]
其中 \(\beta_1\) 是衰减因子，\(\mathbf{m}_t\) 表示动量向量，\(\mathbf{g}_t\) 表示第 \(t\) 次迭代的梯度。该技术平滑了梯度估计，在经验和理论上加速了收敛并增强了训练稳定性 (Sutskever et al., 2013; Chen et al., 2019a; Jelassi and Li, 2022; Fu et al., 2023)。

虽然标准动量是一个稳健的基线，但研究旨在改进它，主要通过：(i) 降低随机梯度估计的方差，以及 (ii) 基于动量和梯度特性自适应学习。方差缩减技术，如 SPIDER (Fang et al., 2018)、STORM (Cutkosky and Orabona, 2019)、SUPER-ADAM (Huang et al., 2021) 和 MARS (Yuan et al., 2024)，通过用方差更低的梯度估计器 \(\mathbf{g}_t'\) 替换原始随机梯度 \(\mathbf{g}_t\) 来工作。这个改进的估计器随后被用于动量更新：\(\mathbf{m}_t = \beta \mathbf{m}_{t-1} + (1 - \beta) \mathbf{g}_t'\)。虽然这些方法在理论上加速了收敛，但它们通常需要额外的计算或存储（例如，存储过去的梯度）。相比之下，MGUP 采用了一种不同的策略，专注于根据动量和当前随机梯度的特性自适应调整更新幅度，而不是直接改变梯度估计的方差。

另一种重要的方法涉及根据动量估计的感知可靠性或特性来调整优化步长。指导这类方法的直觉可以概括为：对于可靠动量增大步长；对于不可靠动量减小步长。这种调整通常通过调制动量向量来实现，可以一般地表示为：
\[
\mathbf{x}_{t+1} = \mathbf{x}_t - \eta_t \mathbf{m}_t \odot \phi_t,
\tag{2}
\]
其中 \(\phi_t\) 是一个缩放因子，通常逐元素应用，由梯度统计量决定。早期的自适应方法，如 Adagrad (Duchi et al., 2011)，通过累积平方梯度引入了每参数学习率。广泛采用的 Adam 优化器 (Kingma and Ba, 2015) 在此基础上，使用梯度的第一矩 \(\mathbf{m}_t\) 和第二矩 \(\mathbf{v}_t\) 的 EWMA：
\[
\mathbf{v}_t = \beta_2 \mathbf{v}_{t-1} + (1 - \beta_2) \mathbf{g}_t^2.
\]
更新步骤随后按元素缩放 \(1 / \sqrt{\hat{\mathbf{v}}_t + \epsilon}\)，其中 \(\hat{\mathbf{v}}_t\) 是偏差校正后的 \(\mathbf{v}_t\)，且 \(\epsilon > 0\) 是一个小常数。这使得 Adam 能够基于历史梯度幅度为每个参数调整学习率。

后续研究探索了各种缩放因子，经常研究当前梯度 \(\mathbf{g}_t\) 与累积动量 \(\mathbf{m}_t\) 之间的相互作用。AdaBelief 优化器 (Zhuang et al., 2020) 修改了 Adam 的第二矩，使用动量与当前梯度之差的平方 \((\mathbf{m}_t - \mathbf{g}_t)^2\)，而不是原始梯度平方 \(\mathbf{g}_t^2\)。第二矩 \(\mathbf{v}_t\) 的更新规则如下，初始条件 \(\mathbf{v}_0 = 0\)：
\[
\mathbf{v}_t = \beta_2 \mathbf{v}_{t-1} + (1 - \beta_2) (\mathbf{m}_t - \mathbf{g}_t)^2 = (1 - \beta_2) \sum_{i=1}^t \beta_2^{t-i} (\mathbf{m}_i - \mathbf{g}_i)^2.
\]
项 \((\mathbf{m}_t - \mathbf{g}_t)^2\) 通过与动量的一致性来衡量对当前梯度的“信任”。显著偏差会增加 \(\mathbf{v}_t\) 中的对应元素，从而降低该参数的有效步长。该机制旨在将 Adam 的快速收敛与 SGD 的泛化能力结合起来。记 \(\mathbf{m}_{t,i}\) 和 \(\mathbf{g}_{t,i}\) 分别为动量向量 \(\mathbf{m}_t\) 和梯度向量 \(\mathbf{g}_t\) 的第 \(i\) 个元素。如果 \(\mathbf{m}_{t,i}\) 和 \(\mathbf{g}_{t,i}\) 符号不同，\((\mathbf{m}_{t,i} - \mathbf{g}_{t,i})^2\) 通常大于 \(\mathbf{g}_{t,i}^2\)（对于相似幅度），从而增加 \(\mathbf{v}_{t,i}\) 并自适应地减小步长。

同时，谨慎优化器 (Liang et al., 2024) 采取了一种更直接的方法来利用动量与梯度之间的符号一致性。它使用一个元素级掩码 \(\varphi_t\) 来选择性地应用动量更新：
\[
\varphi_t = \alpha \cdot \mathbb{I}(\mathbf{m}_t \odot \mathbf{g}_t > 0), \quad \mathbf{x}_{t+1} = \mathbf{x}_t - \eta_t \mathbf{m}_t \odot \varphi_t.
\]
这里，\(\mathbb{I}(\cdot)\) 是指示函数，当其参数为正时等于 1，否则为 0。如果 \(\mathbf{m}_{t,i}\) 和 \(\mathbf{g}_{t,i}\) 的符号对齐，则动量分量 \(\mathbf{m}_{t,i}\) 被缩放 \(\alpha > 1\)；否则，该分量的更新被清零。这种“谨慎更新”策略旨在防止更新受到潜在冲突的梯度信息的影响。

然而，这些先进的自适应方法有明显的局限性。AdaBelief 对二阶矩的依赖限制了其适用范围。