新优化器 Adan 只需常用优化器的 1/2 到 2/3 的训练成本 (epochs) 达到模型的 SoTA 性能。
最近，新加坡 Sea AI Lab 联合北大 Zero LAB 提出新的深度模型优化器 Adan。该优化器可以仅用一半的计算量就能追平甚至超越现有的 SoTA 优化器。在计算量一样的情况下， Adan 的结果在多个场景（涉及 CV、NLP、RL），多种训练方式（有监督与自监督），多种网络结构（ViT、CNN、LSTM、Transformer 等）上，均已超过其默认以及对应的 SoTA 优化器。具体来说，在不改变结构和训练方式的情况下，Adan 在模型 Swin、ViT、ResNet、ConvNext、MAE、LSTM、BERT、Transformer-XL 以及 RL 中常用的 PPO 算法上，均获得了非平凡的性能提升。
- 代码、配置文件和训练 log 均已开源：https://github.com/sail-sg/Adan
- 原文链接：https://arxiv.org/pdf/2208.06677.pdf
自从 Vision Transformer (ViT) 提出，视觉模型的训练变得越来越复杂，训练时间也变的越来越长。除了涌现出一大批的数据增强技巧，常用的如 MixUp、CutMix、AutoRand，模型训练的 epoch 也越来越大。在 Imagenet 1k 上，训练 epoch 从 ResNet 刚提出的 90 epoch 已经增长到了常用的 300，甚至针对一些自监督学习的模型，例如 MAE，预训练的 epoch 已经达到了 1.6k epoch。训练 epoch 增加意味着训练时间极大的延长，这样极大地增加了学术研究或工业落地的成本。目前一个普遍的做法是增大训练的 batch size 并辅助并行训练以减少训练时间。但是随之而来的一个问题就是，大的 batch size 往往意味着 performance 的下降，并且 batch size 越大，情况越明显。这主要是因为模型参数更新次数的急剧减少，以及较小的梯度噪声不能很好的帮助优化器快速逃离鞍点和尖锐的局部极小点。除此之外，大部分优化器并不能忍受较大的学习率，这极大的限制了优化器的收敛速度，从而增加训练成本。同时，还有一个困扰研究者的问题：在小数据集小 epoch 上有效的超参 / 模型，往往在大数据集大 epoch 会失效。是否存在一个能直接在大数据上快速验证超参 / 模型的优化方法？以上的这些问题，均指向一条解决路径：能否给出一个新的优化器，能以较少的训练成本达到模型的SoTA性能，即提升深度模型优化器收敛速度。由于 Adam 的流行，最近几年提出的深度模型优化器均沿用着 Adam 的冲量范式：重球法，
其中 g_k 是随机噪声，m_k 是 moment，eta 是学习率。
Adam 只是将 m_k 的更新由累积形式换成了平均形式，并沿用 RMSProp 提出的学习率放缩技巧。然而随着 Adam 训练原始 ViT 失败，AdamW 渐渐的变成了训练 ViT 甚至 ConvNext 首选。但是 AdamW 并没有改变 Adam 中的冲量范式，因此在当 batch size 超过 4,096 的时候，AdamW 训练出的 ViT 的性能会急剧下降。Adam 类型优化器中的冲量范式似乎并不能在大 batch size 和较强的数据增强下高效的逃离鞍点和坏的局部极小点，尤其是对于 ViT 这种尖锐局部极小点较多的模型。基于上述原因，我们应当为深度模型优化器考虑一种新的冲量范式。在传统凸优化领域，有一个与重球法其名的冲量技巧 ---Nesterov 冲量算法：
Nesterov 冲量算法在强凸问题上，拥有比重球法更快的理论收敛速度，并且在许多实际问题上也展现出较大的速度优势。同重球法不同的是，Nesterov 算法不在原点计算梯度，而是利用冲量找到一个外推点，在该点算完梯度以后再进行梯度累积。外推点能帮助 Nesterov 算法提前感知周围的 landscape。这种特性直觉上似乎更适应 ViT 的训练，因为他并不是简单的依靠冲量去走出尖锐局部极小点，而是通过提前观察周围的梯度，调整更新的方向。基于此，我们可以考虑用 Nesterov 冲量算法去替代目前 Adam 类优化器中的冲量范式。但是在这之前需要解决一个问题：Nesterov 算法需要在外推点计算梯度，在原点更新，如何避免多次模型参数重载和有效进行 back-propagation (BP)？为了解决这个问题，我们首先对 Nesterov 进行改写：
可以证明，改写的 Nesterov 冲量算法与原算法等价，两者的迭代点可以相互转化，且最终的收敛点相同。可以看到，通过引入梯度的差分项，我们已经可以避免手动的参数重载和人为的在外推点进行 BP。接下来，我们就可以利用改写的 Nesterov 冲量算法去替换原有的 Adam 类算法的冲量范式了。同其他的 adaptive 类型的优化器类似，我们也将 m_k 的更新由累积形式换成了平均形式，并使用二阶项对学习率进行放缩：
可以发现，m_k 的更新将梯度与梯度的差分耦合在一起，但是在实际场景中，我们往往需要对物理意义不同的两项进行单独处理，因此因此我们引入梯度差分的冲量 v_k：
我们可以对梯度的冲量和其差分的冲量设置不同的冲量 / 平均系数。对于带 L2 权重正则的目标函数，目前较流行的 AdamW 优化器通过对 L2 正则解耦，在最近流行的网络上（例如 ViT 和 ConvNext）获得了较好的性能。但是 AdamW 所用的解耦方法偏向于启发式，目前并不能得到其收敛的理论保证。基于对 L2 正则解耦的思想，严格地从优化的角度，我们也给 Adan 引入解耦的权重衰减策略。目前 Adan 的每次迭代可以看成是在最小化优化目标 F 的某种一阶近似：
其中
是基于 n_k 的距离度量。由于 F 中的 L2 权重正则过于简单且光滑性很好，以至于我们不需要对其进行一阶近似。因此，我们可以只对训练 loss 进行一阶近似而忽略 L2 权重正则，那么 Adan 的最后一步迭代将会变成：
有趣的是，我们可以发现 AdamW 的更新准则
是 Adan 更新准则在学习率 eta 接近 0 时的一阶近似。因此，我们可以从 proximal 算子的角度给 Adan 甚至 AdamW 给出合理的解释而不是原来的启发式改进。将上述两个改进结合进 Adan 的基础版本，我们可以得到如下的 Adan 优化器。
这里我们给出了一个重启优化器的条件，即当某种条件满足的时候，Adan 会重置冲量。这是一个在随机优化领域很常见的技巧，Adan 也给出这样的接口。需要说明的是，为了效率，我们并没有在实验中启用重启条件，尽管重启优化器能带来一定的性能提升（见论文，表 12）。这里我们跳过繁复的数学表达式，只给出结论：在给定或未给定 Hessian-smooth 条件的两种情况下，Adan 优化器的收敛速度在非凸随机优化问题上均能达到已知的理论下界。并且该结论在带有解耦的权重衰减策略时仍然成立。目前为止，带有解耦的权重衰减策略的优化器的收敛性分析仍然是个未解的问题。Adan 的收敛性分析为该问题的解决提供了一个值得尝试的思路。- ViT 模型：针对 ViT 模型，我们分别在传统的 ViT 和杰出的 Swin 结构上，测试了 Adan 的性能。
并且在 ViT-small 上在不同的 epoch 下，详细的测试了 Adan 和其他优化器的性能。
可以看到，例如在 ViT-small，ViT-base，Swin-tiny 以及 Swin-base 上，Adan 仅仅消耗了一半的计算资源就获得了同 SoTA 优化器接近的结果。并且在同样的计算量下，Adan 在两种 ViT 模型上均展现出较大的优势。- CNN 模型：除了较难训练的 ViT 模型，我们也在尖锐局部极小点相对较少的 CNN 模型上也测试了 Adan 的性能 --- 包括经典的 ResNet 与较先进的 ConvNext。结果如下：
可以观察到，不管是 ResNet 还是 ConvNext，Adan 均能在大约 2/3 训练 epoch 以内获得超越 SoTA 的性能。在无监督训练框架下，我们在最新提出的 MAE 上测试了 Adan 的表现。其结果如下：
同有监督学习的结论一致，Adan 仅消耗了一半的计算量就追平甚至超过了原来的 SoTA 优化器，并且当训练 Epoch 越小，Adan 的优势就越明显。在 NLP 的有监督学习任务上，我们分别在经典的 LSTM 以及先进的 Transformer-XL 上观察 Adan 的表现。
Adan 在上述两种网络上，均表现出一致的优越性。并且对于 Transformer-XL，Adan 在一半的训练步数内就追平了默认的 Adam 优化器。为了测试 Adan 在 NLP 场景下无监督任务上的模型训练情况。我们从头开始训练 BERT：在经过 1000k 的预训练迭代后，在 GLUE 数据集的 7 个子任务上测试经过 Adan 训练的模型性能，结果如下：
Adan 在所测试的 7 个词句分类任务上均展现出较大的优势。值得一提的是，经过 Adan 训练的 BERT-base 模型，在一些子任务上（例如 RTE，CoLA 以及 SST-2）的结果甚至超过了 Adam 训练的 BERT-large。我们将 RL 常用的 PPO 算法里的优化器替换为了 Adan，并在 MuJoCo 引擎中的 4 个游戏上测试了 Adan 的性能。在 4 个游戏中，用 Adan 作为网络优化器的 PPO 算法，总能获得较高的 reward。
Adan 在 RL 的网络训练中，也表现出较大的潜力。Adan 优化器为目前的深度模型优化器引入了新的冲量范式。相较于之前的优化器，Adan 给出的范式更适合目前的强数据增强大 batch size 的场景。实验显示，Adan 仅需 1/2-2/3 的计算量就能追平或者超越现有的 SoTA 优化器。
Adan 在多个场景（涉及 CV，NLP，RL），多个训练方式（有监督与自监督），多种网络结构（ViT，CNN，LSTM，Transformer 等）上，均展现出非平凡的性能优势。理论角度，Adan 优化器的收敛速度在非凸随机优化上也已经达到了理论下界。掌握「声纹识别技术」：前20小时交给我，后9980小时……
《声纹识别：从理论到编程实战》中文课上线，由谷歌声纹团队负责人王泉博士主讲。课程视频内容共 12 小时，着重介绍基于深度学习的声纹识别系统，包括大量学术界与产业界的最新研究成果。同时课程配有 32 次课后测验、10 次编程练习、10 次大作业，确保课程结束时可以亲自上手从零搭建一个完整的声纹识别系统。© THE END
转载请联系本公众号获得授权
投稿或寻求报道：content@jiqizhixin.com
关注公众号：拾黑（shiheibook）了解更多
[广告]赞助链接：
四季很好，只要有你，文娱排行榜：https://www.yaopaiming.com/
让资讯触达的更精准有趣：https://www.0xu.cn/
机器之心
关注网络尖刀微信公众号
