Adafactor：T5 背后的省显存优化器，不存二阶矩状态也能训好模型

我最早训练模型的时候，心里只有一个优化器——Adam。所有教程都在用，所有代码都默认，我甚至没想过还有别的选择。直到我翻开T5 论文，发现他们用的竟然是一个听都没听过的名字：Adafactor。当时的反应是：Google是不是疯了？放着好好的Adam不用？后来我才明白，不是他们疯了，是Adam真的用不起。

T5 的模型规模是 110 亿参数。如果用 Adam 训练，光是优化器状态就比模型本身大 4 倍以上，这还没算梯度。这个显存开销，足以让大多数训练任务在还没开始之前就宣告失败。所以 Google 选择了 Adafactor，一个专门为省显存而生的优化器。

Adam 的隐性成本：每个参数要额外付 8 个字节

Adam 要维护两个状态：一阶矩 m 和二阶矩 v。每个状态都是一个 float32，4 字节。所以每个参数要额外付出 8 字节。听起来不多？算一笔账：一个 10 亿参数模型，光优化器状态就是 8GB。而模型本身如果用 fp16 存储，只有 2GB。也就是说，Adam 状态比模型本身大 4 倍。

对比项	Adam	Adafactor
一阶矩（动量）	每个参数存一个 float	默认不存
二阶矩	每个参数存一个 float	矩阵参数分解为行向量 + 列向量
n×m 矩阵的状态内存	2×n×m×4 字节	(n+m)×4 字节
10 亿参数模型状态量	8GB	视形状可减少 2～3 个数量级

Adafactor 的核心：二阶矩矩阵可以拆成两个向量的外积

Adam 的二阶矩 v 是每个参数的历史梯度平方的指数滑动平均。它的形状和参数一模一样。比如一个全连接层的权重是 n 行 m 列的矩阵，那么 v 也是 n×m 的矩阵。

Adafactor 做了一个大胆的假设：这个 n×m 的矩阵，也许可以用一个 n 维向量和一个 m 维向量的外积来近似。也就是说，你不需要存 n×m 个值，只需要存 n + m 个值。具体做法是：维护行统计量 R（n 维）和列统计量 C（m 维），然后用 R 和 C 的外积除以它们的全局和，来估计实际的二阶矩。这样，内存从 O(nm) 降到了 O(n+m)。对于一个 1024×1024 的矩阵，原来要存 100 万个状态，现在只需要 2048 个，直接少了三个数量级。

说实话，我第一次看到这个分解时心里是怀疑的，因为 v 是逐参数的历史梯度平方，怎么可能用两个向量的外积近似？后来我意识到，神经网络的梯度矩阵本身往往就是低秩的，这个近似比想象中合理。

对于向量参数（比如 bias），没法做矩阵分解，Adafactor 就只存一个标量，或者直接用全局统计。论文里有对应的处理。

没有动量？靠相对步长来凑

Adam 的另一个大开销是动量项 m。Adafactor 默认不存动量。你可能觉得，没有动量，训练会不稳定吧？Adafactor 用了一个替代方案：相对步长。它不是用固定的学习率，而是根据每个参数在当前时刻的梯度尺度，自动调整更新幅度。具体来说，它估计一个全局的梯度尺度，然后让每个参数的更新步长相对于这个尺度变化。这种方式在很多任务上可以代替动量带来的稳定作用。

这个设计有一个好处：你不用像 Adam 那样费心去调学习率，因为学习率是相对自适应的。当然，这也意味着它需要更长的 warmup 来稳定训练。

效果到底行不行？

Adafactor 的论文里，作者在机器翻译、语言建模、图像分类等任务上做了对比。结果显示，在内存减半以上的情况下，Adafactor 的效果和 Adam 基本持平，有些任务上略差，但差距不大。更重要的是，T5 这个 110 亿参数的模型，就是靠 Adafactor 训练出来的。这也证明了它在超大模型上的实用性。

不过要注意，Adafactor 并不总是直接替代 Adam 就能跑出一样的曲线。它对 warmup 步数、相对步长的初始化更敏感。如果你只是把 Adam 换成 Adafactor，其他设置不变，很可能会发现收敛变慢。

用 Adafactor 前，先搞清楚这些设置

在 Hugging Face Transformers 里，使用 Adafactor 非常简单：

from transformers import Adafactor; optimizer = Adafactor(model.parameters(), scale_parameter=True, relative_step=True, warmup_init=True)

这里有三个关键参数：relative_step=True 表示使用相对步长；scale_parameter=True 表示根据参数形状缩放学习率；warmup_init=True 表示在训练开始时用较小的初始步长。如果你把这三个都关掉，那它就是一个普通的 Adam 变体，只是省了二阶矩的内存。

我最早直接用默认参数训练一个中文摘要模型，结果 loss 一直下不去。后来发现是因为 relative_step 开着的，模型还在 warmup 阶段，而我只训了几千步。关掉之后才恢复正常。所以我的经验是：训练 Transformer 模型时，打开这三个参数，效果比较稳；如果你用预训练模型做微调，建议把 relative_step 关掉，设置一个较小的固定学习率。