AB
AiBoss
Wiki

什么是 Adafactor(Adafactor)?一种以次线性内存开销实现自适应学习率的优化器

Adafactor 是 Noam Shazeer 与 Mitchell Stern 在 2018 年论文中提出的一种随机优化方法,它通过只保存平方梯度滑动平均的行和与列和,把自适应学习率所需的内存从与参数量成正比降到次线性规模,并配合更新裁剪、逐步提高的衰减率与参数尺度缩放等技巧,在 Transformer 机器翻译训练上取得与 Adam 相当的结果。

Adafactor 是一种用于训练神经网络的随机优化方法(stochastic optimization method),由 Noam Shazeer 与 Mitchell Stern 在 2018 年的论文《Adafactor: Adaptive Learning Rates with Sublinear Memory Cost》中提出。它要解决的问题是:以 RMSProp、Adam、Adadelta 为代表的自适应学习率方法需要为每个参数单独维护平方梯度的指数滑动平均(即二阶矩估计),这部分辅助存储与参数量成正比;Adafactor 通过只维护矩阵行方向与列方向的滑动平均之和来估计逐参数二阶矩,从而把优化器状态的内存开销降到次线性(sublinear)水平。

为什么重要

在 Adafactor 出现之前,训练深度模型时常用的自适应优化器(如 Adam)会为每个可训练参数保存一份二阶矩估计。论文指出,维护这些逐参数的二阶矩估计所需的内存与参数数量相等。对于参数量很大的神经网络权重矩阵,这部分优化器状态会显著抬高训练所需的显存与存储开销,在模型规模增长时成为实际约束。

另一方面,论文还指出自适应方法存在一个稳定性问题:当二阶矩累加器的衰减率过慢时,参数更新可能比期望的更大。也就是说,内存开销与更新幅度控制这两件事,是 Adafactor 试图同时处理的痛点。

工作机制

按论文的叙述,Adafactor 的核心做法可以拆成以下几个要点:

  • 行列分解的二阶矩估计。 对神经网络权重矩阵,不再逐参数保存平方梯度的指数滑动平均,而是只保存这些滑动平均在每一行上的和与每一列上的和,再基于这两组和来估计每个参数对应的二阶矩。这样优化器需要保存的量从「参数量」量级降为「行数加列数」量级,即论文所说的次线性内存代价。
  • 更新裁剪(update clipping)。 针对衰减率过慢导致更新过大的问题,论文提出对最终梯度更新的均方根(root mean square)设置阈值进行裁剪。Hugging Face 文档中对应的参数 clip_threshold 默认值为 1.0,含义正是「最终梯度更新均方根的阈值」。
  • 逐步提高的衰减率。 论文同时提出让二阶矩累加器的衰减率随时间逐渐增大,作为更新过大的补救手段。Hugging Face 文档中 decay_rate 默认值为 -0.8,用于计算平方梯度的滑动平均。
  • 去掉动量。 论文将上述方法与「丢弃动量(dropping momentum)」结合使用。Hugging Face 实现中 beta1 默认为 None,即不启用一阶矩滑动平均;若给定该系数,则用于计算梯度的滑动平均。
  • 按参数自身尺度缩放更新。 论文最后提出根据参数本身的尺度来缩放参数更新。Hugging Face 文档中 scale_parameter 默认为 True,表示学习率按均方根进行缩放。
  • 相对步长与预热初始化。 Hugging Face 文档说明,该优化器会根据 scale_parameter、relative_step 与 warmup_init 在内部调整学习率;relative_step 为 True 时计算与时间相关的学习率,而不是使用外部学习率;warmup_init 则决定时间相关学习率的计算是否采用预热初始化。

需要强调的是,上述机制是论文作者提出的设计与补救方案,论文以实验说明其效果,并不构成对所有模型、所有任务的普遍结论。

典型例子

论文给出的主要实证场景是机器翻译:作者将上述方法组合起来并去掉动量,在 WMT 2014 英德(English-German)机器翻译任务上训练 Transformer 模型,报告称取得了与已发表的 Adam 训练方案相当的结果,同时优化器占用的辅助存储很少。论文还表示,其实验显示行列分解的二阶矩估计方法与基线产生相似的结果。

在工程实现层面,Hugging Face 的 Transformers 文档提供了 Adafactor 类,其构造参数包括 lr、eps(默认 (1e-30, 0.001),分别为平方梯度与参数尺度的正则常数)、clip_threshold、decay_rate、beta1、weight_decay、scale_parameter、relative_step 与 warmup_init。文档称该 PyTorch 实现可作为 Adam 的直接替代(drop-in replacement),并注明其内部会依据上述选项自动调整学习率;若要使用手动(外部)学习率调度,应设置 scale_parameter=False 与 relative_step=False。文档还提到该实现能处理低精度(FP16、bfloat)数值,但表示尚未经过充分测试,并给出了 T5 微调的推荐设置。

边界与常见误解

首先,Adafactor 的内存优势来自对权重矩阵做行列分解的近似估计,它估计的是逐参数二阶矩,而不是精确保存每一个二阶矩。论文的表述是该方法产生与基线相似的结果,这属于作者在其实验设置下的报告,不应被理解为在任何模型结构上都等价于 Adam。

其次,Adafactor 并不等同于「Adam 加一个更省内存的开关」。论文明确去掉了动量,并额外引入更新裁剪与逐步提高的衰减率来处理自适应方法的更新过大问题;Hugging Face 文档也提醒,其内部学习率行为受 scale_parameter、relative_step、warmup_init 控制,若使用者仍按外部学习率调度的习惯配置,需要显式关闭相关选项,否则实际生效的学习率可能与预期不同。

第三,关于低精度支持,Hugging Face 文档的措辞是「能处理」但「尚未充分测试」,这属于该文档对自身实现的说明,不能推广为 Adafactor 在低精度训练下已被普遍验证。

最后,论文的实证集中在 WMT 2014 英德翻译上的 Transformer 训练,其结论的适用范围应以论文所报告的任务与设置为限;把它当作对所有训练场景都更优的通用替代方案,超出了原始素材的支持范围。

参考资料