Adam的epsilon如何影响学习率的Scaling Law？

AI大模型入门

于 2024-12-02 10:40:17 发布

阅读量936

点赞数 38

文章标签：学习数据库语言模型自然语言处理人工智能 AI ai大模型

本文链接：https://blog.csdn.net/m0_70486148/article/details/144181993

版权

上一篇文章我们从多个角度讨论了学习率与 Batch Size 之间的缩放规律，其中对于 Adam 优化器我们采用了 SignSGD 近似，这是分析 Adam 优化器常用的手段。那么一个很自然的问题就是：用 SignSGD 来近似 Adam 究竟有多科学呢？

我们知道，Adam 优化器的更新量分母会带有一个，初衷是预防除零错误，所以其值通常很接近于零，以至于我们做理论分析的时候通常选择忽略掉它。然而，当前 LLM 的训练尤其是低精度训练，我们往往会选择偏大的，这导致在训练的中、后期往往已经超过梯度平方大小，所以的存在事实上已经不可忽略。

因此，这篇文章我们试图探索 如何影响 Adam 的学习率与 Batch Size 的 Scaling Law，为相关问题提供一个参考的计算方案。

SoftSign

由于是接着上一篇文章介绍，所以就不再重复相关背景了。为了探究的作用，我们从 SignSGD 换到 SoftSignSGD，即变成，其中：

这个形式无疑更贴近更贴近 Adam。但在此之前，我们需要确认是否真的不可忽略，才能确定是否有进一步研究的价值。

在 Keras 的 Adam 实现中，的默认值是，在 Torch 中则是，这说明的默认值是级别，这时候梯度绝对值小于几率还不算大；但在 LLM 中，的普遍取值是（比如 LLAMA2 [1]），这时候的取值已经来到级别，大概训练进入“正轨”后，梯度绝对值几乎普遍都小于了，所以的影响是显著的。

这个跟 LLM 的参数量也有一定关系。一个能稳定训练的模型，不管参数量多大，它的梯度模长大小大致都在同一数量级，这是反向传播的稳定性决定的（参考训练 1000 层的 Transformer 究竟有什么困难？）。因此，参数量越大的模型，平均下来每个参数的梯度绝度值就相对变小了，从而的作用就更突出了。

值得指出的是，的引入实际上提供了 Adam 与 SGD 之间的一个插值，这是因为当时：

所以，越大，Adam 表现越接近 SGD。

（注：本文 SoftSign 的概念，源于笔者跟 MSR 的刘力源老师、董城昱同学的一个 ongoing collaboration，经我们商量一致后先把这部分结果分享出来，更多后续结论敬请持续关注。）

S型近似

确认了引入必要性后，我们着手开始分析。在分析过程中，我们将会反复遇到 S 型函数，所以还有一个准备工作是探究 S 型函数的简单近似。

S 型函数相比大家已经见怪不怪，上一节引入的函数本身就是之一，上一篇文章分析过程中的函数也是一例，此外还有、等。接下来我们处理的是满足如下特性的 S 型函数：

1. 全局光滑且单调递增；

2. 奇函数，值域是 [-1,1]；

3. 在原点处斜率为。

对于这类函数，我们考虑两种近似。第一种近似跟类似：

它大概是保留如上 3 点性质的最简单函数了；第二种近似是基于函数：

这本质上是一个分段线性函数，放弃了全局光滑的性质，但分段线性会使得积分算起来更容易，我们很快就会看到这一点。

▲ Erf 函数与它的两种近似

均值估计

事不宜迟，沿着上一篇文章的方法，出发点还是：

我们需要做的事情就是估计和。

这一节我们算的是，为此我们需要用函数去近似函数：

然后我们有：

积分形式很复杂，但用 Mathematica 算并不难，结果可以用函数表达出来：

其中。这个函数看起来比较复杂，但它刚好是的 S 型函数，值域为且在处的斜率是，所以利用第一种近似形式

第二个约等号是利用近似来处理分母中的。可以说相当幸运，最终的形式并没有太复杂。接着我们有：

跟上一篇文章一样，最后一个约等号使用了平均场近似，是全体的某种平均，而以及。

方差估计

均值也就是一阶矩解决了，现在轮到二阶矩了：

结果同样可以用函数表示，但更加冗长，这里就不写出来了，还是那句话，对 Mathematica 来说这都不是事。视为的函数时，可以发现结果是一条倒钟形的曲线，关于轴对称，上界是 1，最小值是则在内。

参考的近似式（9），我们选择如下近似：

有一说一，这个近似的精度并不高，主要是为了计算的方便，但它已经保留了倒钟形、轴对称、上界为 1、时结果为 1、结果则为 0 等关键特性。接下来继续应用平均场近似：

所以。其中这一项就代表了的协方差矩阵，它是一个对角阵，这是可以预料的，因为我们的假设之一是各分量之间的独立性，所以协方差矩阵必然是对角阵。

结果初探

由此我们得到：

注意，除了外，剩余的其他符号都不依赖于，所以上式已经给出与的依赖关系。注意为了保证极小值的存在性，我们都会假设矩阵的正定性，而在此假设之下必然有和。

上一篇文章我们说 Adam 最重要的特性是可能会出现 “Surge 现象”，即关于不再是全局的单调递增函数。接下来我们将会证明，的引入会降低现象出现的几率，并且时完全消失。这个证明并不难，很明显 Surge 现象出现的必要条件是：

若否，整个关于便是单调递增的，而关于是单调递增的，所以整个关于单调递增，不存在 Surge 现象。

别忘了是关于的单调递减函数，所以当增大时会更小，从而上述不等式成立的可能性更低，并且时为零，上述不等式不可能再成立，因此 Surge 现象消失。

进一步，我们可以证明时，结果跟 SGD 的一致，这只需要留意到：

我们有极限：

这里是全体的某种平均。于是我们得到当足够大时有近似：

右端就是假设梯度协方差矩阵为 ( 时的 SGD 结果。

文章小结

本文延续了上一篇文章的方法，尝试分析了 Adam 的对学习率与 Batch Size 之间的 Scaling Law 的影响，结果是一个介乎 SGD 与 SignSGD 之间的形式，当越大，结果越接近 SGD，“Surge 现象”出现的概率就越低。总的来说，计算结果没有特别让人意外之处，但可以作为分析作用的一个参考过程。