训练超大型深度学习模型——尤其是拥有数十亿参数的大型语言模型 (LLM)——是一项艰巨的任务。其中最大的瓶颈之一不仅是计算量，还有所需的海量内存。这部分内存的很大一部分并非由模型权重本身消耗，而是由优化器的状态——即为高效更新模型而需要跟踪的额外数据——占用的。

Adam 优化器是训练 LLM 的首选，但它特别“吃内存”。它为模型中的每个参数额外存储两个状态值，与 SGD 等更简单的方法相比，几乎将内存占用翻倍。以 LLaMA-2 7B 模型为例——仅优化器的状态就需要超过 25 GB 内存！

这促使人们竞相研发更节省内存的优化器。我们见过一些聪明的方案，例如对优化器状态进行量化的 8-bit Adam，以及采用低秩梯度投影的 GaLore。虽然这些方法有用，但大多是启发式的，往往缺乏严格的数学收敛保证。有时这种取舍会导致模型精度显著下降。

于是，一个核心问题摆在我们面前:
能否在不牺牲性能的前提下，同时实现大幅的内存优化和强有力的收敛保证？

来自 ISTA 和 KAUST 的研究人员在近期论文 “MICROADAM: Accurate Adaptive Optimization with Low Space Overhead and Provable Convergence” 中给出了响亮的答案: 可以。他们提出了 MicroAdam——一种新的优化器，通过巧妙压缩梯度信息来显著降低内存使用，同时首次为此类方法提供了严谨的理论收敛性证明。

背景: 自适应的成本与压缩的潜力

要理解 MicroAdam，我们需要先搞清楚为什么 Adam 内存开销如此之大——以及从哪里可以精简。

Adam 为每个参数 \( \theta \) 维护两个滑动平均值:

  1. 一阶矩 (\( m_t \)): 梯度的指数加权移动平均，类似动量，用于平滑并加速参数更新。
  2. 二阶矩 (\( v_t \)): 梯度平方的指数加权移动平均，用于为每个参数自适应调整学习率。梯度更新幅度大的参数，其学习率会减小；更新稀少的参数，学习率则会增大。

同时存储所有参数的 \( m_t \) 和 \( v_t \) 是 Adam 高效但昂贵的根源。如果一个模型有 \(d\) 个参数，且以 16 位浮点 (2 字节) 存储，那么 Adam 的状态开销为:

  • float32: \( 8d \) 字节 (m 与 v 各需 4 字节)
  • float16: \( 4d \) 字节

一个自然的设想是: 不存储完整梯度。
梯度压缩——尤其是梯度稀疏化——只保留一部分“重要”的梯度分量。例如，Top-K 方法选取幅度最大的 \(k\) 个分量，其余全部置零。在超大模型中，这种压缩可非常激进: 99% 稀疏度意味着只存储 1% 的梯度值。

然而，简单稀疏化有个大问题: 小梯度值会被系统性忽略。这会引入偏差，可能使部分权重长期不更新，或者使更新方向偏离正确轨迹。

误差反馈登场

误差反馈 (Error Feedback，EF) 是分布式机器学习中提出的一种巧妙方案，可优雅地解决上述问题。压缩时舍弃的梯度分量？将它们存储在一个误差缓冲区中，并在下一次迭代中加回到梯度里。这样一来，信息并未真正丢失——只是延迟使用。

在分布式训练中，EF 效果显著。但在单机场景中，标准 EF 存在一个令人沮丧的悖论: 误差缓冲区的大小与完整梯度一样！这直接抵消了节省的内存空间。

MicroAdam 的突破既简单又深刻:
连误差缓冲区也进行压缩。

MicroAdam 的工作原理

MicroAdam 将 Top-K 稀疏化与误差反馈相结合——但重点是使用低精度量化来压缩误差反馈，从而显著减少内存占用并保持收敛性。

MicroAdam 及其辅助函数的算法伪代码。左侧展示主循环，右侧展示 ADAMSTATS 和量化过程。

1. 先校正，再压缩

在第 \( t \) 步:

  1. 从当前梯度 \( g_t \) 开始。
  2. 使用反量化后的误差反馈 \( e_t \) 对其作校正:
\[ a_t = g_t + \text{dequantize}(e_t) \]
  1. 应用 Top-K 稀疏化 (如取前 1% 最大值) :
\[ I_t, V_t = \text{TopK}(a_t) \]

其中 \(I_t\) 是最大幅值分量的索引，\(V_t\) 是对应的数值。该稀疏向量用于更新 Adam 的矩估计。

2. 压缩误差

移除 Top-K 分量后的残差为:

\[ e_{t+1}^{\text{raw}} = a_t - \text{TopK}(a_t) \]

MicroAdam 不直接存储这个密集向量，而是将其量化为每值 4 位:

\[ e_{t+1} = \text{quantize}(e_{t+1}^{\text{raw}}, \text{bits}=4) \]

与 float32 相比，这使误差缓冲区大小缩小 8 倍——EF 的内存开销几乎可以忽略。

3. 稀疏窗口的动态统计

Adam 通常为所有 \(d\) 个参数密集存储 \( m_t \) 和 \( v_t \)，MicroAdam 则避免这样做:

  • 仅保留一个包含最近 \( m \) 步 (如 10 步)** 稀疏**梯度 \((I_t, V_t)\) 的滑动窗口；
  • 在需要时基于稀疏历史动态重算 \( m_t \) 和 \( v_t \)；
  • 借助极高稀疏度 (99% 零值) 和定制 CUDA 实现，重算过程既快又省内存。

误差反馈实战

Adam、TopK-Adam 和带 EF 的 TopK-Adam 在 Rosenbrock 函数上的优化轨迹。标准 Adam 轨迹平滑直达最优；TopK-Adam 缺乏 EF 时路径曲折低效；带 EF 的 TopK-Adam 恢复了平滑高效的轨迹。

在经典的 Rosenbrock 函数最小化实验中:

  • Adam (左) : 平滑、直接到达最优目标；
  • Top-K Adam (中) : 路径曲折——缺失 EF 会严重损害收敛性；
  • Top-K Adam + EF (右) : 平滑轨迹重现。EF 会累计被丢弃的梯度分量并在之后回补，确保轨迹保持正确。

MicroAdam 在超大规模训练中高效利用了这一原理。

内存占用

以 LLaMA-2 7B (\(d\) 个参数) 为例:

  • AdamW (bf16): \(4d\) 字节 → 25.10 GB
  • AdamW-8bit: \(2d\) 字节 → 12.55 GB
  • MicroAdam (m=10，1% 稠密度): \(0.5d\) (4-bit EF) + \(4mk\) (稀疏窗口) = 0.9d 字节 → 5.65 GB

不到 8-bit Adam 内存的一半，比标准 AdamW 轻 4 倍以上。

收敛性保证

MicroAdam 不只是实用，它在理论上同样稳固。

这里涉及两个关键压缩算子:

  1. 梯度压缩器 \(C\) ——Top-K 稀疏器，满足 q-收缩性质:
\[ \|\mathcal{C}(x) - x\| \le q \|x\| \]

对于 Top-K，\(q\) 与稀疏度直接相关。

  1. 误差压缩器 \(Q\) ——4 位量化器，无偏且有界:
\[ \mathbb{E}[Q(x)] = x, \quad \|Q(x) - x\| \le \omega \|x\| \]

关键条件为:

\[ (1 + \omega)q < 1 \]

这确保压缩不会造成过多累积信息丢失。

在此条件下，作者证明:

  • 定理 1 (非凸情形) : 收敛率 \( \mathcal{O}(1/\sqrt{T}) \)，与 AMSGrad 相同。
    展示 MicroAdam 在非凸函数上收敛率的公式。

  • 定理 2 (PL 条件) : 对满足 Polyak-Łojasiewicz 性质的函数，收敛率 \( \mathcal{O}(\log T / T) \)。
    展示 MicroAdam 在 PL 条件下收敛率的公式。

渐近收敛速率与未压缩 Adam 一致——压缩仅影响常数项。

实验: 真实世界验证

微调 BERT 与 OPT

BERT-Base (110M) 、BERT-Large (335M) 、OPT-1.3B——与 Adam、Adam-8bit、CAME、GaLore 对比。

表 1: GLUE/MNLI 微调结果。MicroAdam 在准确率和损失上与基线相当或更优，内存占用与其他高效优化器相似。

MicroAdam 的准确率达到或超过 Adam-8bit，内存占用相当——显著优于 GaLore 和 CAME。

训练损失曲线:
BERT-Base、BERT-Large、OPT-1.3B 的训练损失曲线，MicroAdam 曲线始终较低。
BERT-Base、BERT-Large、OPT-1.3B 的训练损失曲线，MicroAdam 曲线始终较低。
BERT-Base、BERT-Large、OPT-1.3B 的训练损失曲线，MicroAdam 曲线始终较低。

微调 LLaMA-2 7B 与 13B

GSM-8k 数学推理

表 2: 在 GSM-8k 数据集上微调 LLaMA-2 的结果。MicroAdam (m=10) 在适应 40GB GPU 的同时，准确率与完整 Adam 相当。

MicroAdam 实现了在单块 40 GB GPU 上完成 LLaMA-2 7B 全参数微调:

  • 准确率: 34.72% (MicroAdam m=10) vs 34.50% (Adam)
  • 轻松适配 GPU 内存；Adam 则不行。

Open-Platypus 指令微调

表 3: Open-Platypus 数据集结果。MicroAdam 在 AdamW、Adam-8b、MicroAdam 三者中取得最高平均准确率且内存占用最低。

在各项任务中平均准确率最高，内存占用最低。

ImageNet 预训练

ResNet-18、ResNet-50 从零开始训练。

表 4: ImageNet 上训练 ResNet 的结果。MicroAdam 在两款 ResNet 上均取得最高验证准确率，甚至超越优化良好的 SGD 基线。

MicroAdam 获得最高验证准确率——甚至超越调优后的 SGD，且优化器状态内存占用最小。

训练曲线:
ResNet-18、ResNet-50 在 ImageNet 上的训练与验证曲线。MicroAdam (蓝色) 在验证准确率上始终领先。
ResNet-18、ResNet-50 在 ImageNet 上的训练与验证曲线。MicroAdam (蓝色) 在验证准确率上始终领先。

作者推测，这可能源自一种隐式正则化效应: 稀疏更新 (每步约 10% 的权重) 能够降低过拟合。

结论与启示

MicroAdam 不只是一个新优化器，它是一个原则性且内存高效的设计，并且具备可证明的收敛性——填补了效率与严谨性之间的长期缺口。

核心要点:

  1. 压缩误差反馈 —— EF 缓冲区量化为 4 位，内存占用几乎为零。
  2. 大幅内存节省 —— 不到 Adam-8bit 内存的一半，比标准 AdamW 轻 4 倍以上。
  3. 坚实理论基础 —— 收敛速率与未压缩 Adam 相同，压缩只影响常数项。
  4. 突出实证表现 —— 在 LLM 微调中匹敌或超越 Adam；在 CV 预训练验证准确率上击败 SGD。

MicroAdam 证明了误差反馈即便经过高度压缩仍能有效，这为未来既节省内存又具理论可靠性的优化器奠定了基础。
对从业者而言，这意味着可以在平价硬件上完成最新模型的全参数微调，从而降低成本，让更多人参与并推动前沿 AI 研究的普及。