基于Amos优化器思想推导出来的一些“炼丹策略”

 

©PaperWeekly 原创 · 作者 | 苏剑林

单位 | 追一科技

研究方向 | NLP、神经网络

如果将训练模型比喻为“炼丹”，那么“炼丹炉”显然就是优化器了。据传 AdamW 优化器是当前训练神经网络最快的方案，这一点笔者也没有一一对比过，具体情况如何不得而知，不过目前做预训练时多数都用 AdamW 或其变种 LAMB 倒是真的。然而，正如有了炼丹炉也未必能炼出好丹，即便我们确定了选择 AdamW 优化器，依然有很多问题还没有确定的答案，比如：

1. 学习率如何适应不同初始化和参数化？

2. 权重衰减率该怎么调？

3. 学习率应该用什么变化策略？

4. 能不能降低优化器的显存占用？

尽管在实际应用时，我们大多数情况下都可以直接套用前人已经调好的参数和策略，但缺乏比较系统的调参指引，始终会让我们在“炼丹”之时感觉没有底气。在这篇文章中，我们基于 Google 最近提出的 Amos 优化器的思路，给出一些参考结果。

基础回顾

Amos 优化器出自最近的论文《Amos: An Adam-style Optimizer with Adaptive Weight Decay towards Model-Oriented Scale》[1]，它对上述几个问题都推导了比较完整的推导，并通过实验证实了它的有效性。然而，原论文的推导实在是不好读，各种记号和估计都过于随意，给人很“凌乱”感觉。不过好在 Amos 的思想还不算复杂，我们可以借用一下。

在开始推导之前，我们不妨先回顾一下对于上述几个问题，现有的解决方案是怎样的。

首先，第一个问题，大家可能不大理解“初始化”和“参数化”分别是什么含义，其实这就是模型权重的两种设置方式，常见的就是一个 的矩阵，一般用“均值为 0、方差为1/n”的方式初始化，详细介绍可以参考笔者之前《从几何视角来理解模型参数的初始化策略》[2]、《浅谈Transformer的初始化、参数化与标准化》[3]。

从“方差为 1/n”我们就可以看到，不同参数有着不同的尺度（或者说数量级），如果我们用同一个学习率更新所有参数，那么就会导致每个参数的更新幅度不一样。这个问题笔者觉得比较优雅的解决方案就是 LAMB 优化器，它每次更新的模长直接取决于参数本身的模长，学习率只是用来描述相对更新量的大小。

至于权重衰减率问题，至少在预训练领域，笔者观察到的是都是沿用最早的选择 0.01，没有发现去调整该参数的工作。而对于学习率变化策略，大家都知道应该要将学习率慢慢降到零，但具体应该选用什么什么下降策略，暂时也没有太多的理论指导，多数结果也只是实验总结出来的。

最后，关于节省显存问题，比较经典的工作就是 AdaFactor 优化器，笔者之前在《AdaFactor优化器浅析（附开源实现）[4]》也有过介绍。降低优化器显存占用的主要就两个思路，一是去掉动量，二是对二阶矩做低秩分解，Amos 本质上也是沿用了这两个思路。

问题设置

本文主要关心开头的前三个问题，希望能够推导出一些“即插即用”的结果。首先，我们将优化器的更新规则简写成：

其实 分别代表 时刻的参数值， 代表 时刻的更新向量（依赖于任务和数据），而标量 （向量的每个元素都大于 0）代表 时刻的学习率。