查看原文
其他

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

苏剑林 PaperWeekly 2023-09-02
 

©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)代表 时刻的学习率。
自 AdamW 起,主流优化器都倾向于把权重衰减(Weight Decay)项从  中独立出来,即

其中 是权重衰减率。本文的主要任务,就是希望能解决 该怎么设置的问题。


权重衰减
我们知道,权重衰减也好,L2 正则也好,它本身是跟训练目标无关的,它只是一个辅助项,目的是提高模型的泛化能力。既然是辅助,那么一个基本的要求就是它不应该“喧宾夺主”,为此,我们不妨加入一个限制:

也就是说,在整个更新过程中,权重衰减带来的更新量始终要比目标相关的更新量高一阶,由于 基本上都是小于 1 的,所以更高阶意味着更小。
设优化的参数终点是 ,我们记 ,根据更新规则可以得到

最后的近似只保留了不超过 的项。
很明显, 是当前结果与终点的距离,它自然是越小越好,因此我们自然也希望每一步的更新都能缩小这个距离,即
而我们看式 (4), 可正可负,如果它为负就有助于实现 ,但是 必然是正的,它是不利于实现 ,不过在引入权重衰减后,多出了一项 ,如果这一项能抵消掉 的负面作用,那么权重衰减的引入就不仅能增强泛化能力,还有利于模型收敛了。


可行分析
所以,接下来的事情,我们就是要考察

的可行性。所谓可行性,就是 能否大于 0,只有它大于 0,左右两端才有可能相等。利用 的定义我们得到 ,于是

注意 是我们的目标,是一个固定的点,而 是当前时刻与目标的差异向量,两者一般来说没什么必然的相关性,于是我们可以近似认为它们是高维空间中两个随机向量。根据《n维空间下两个随机向量的夹角分布》[5],我们知道高维空间中两个随机向量几乎都是垂直的,于是 。当然,如果不放心,还可以引入一个参数

此时式 (5) 就变成了

两端都大于 0,因此式 (5) 是有可能成立的。


渐近估计

如果式 (5) 成立,那么式 (4) 就简化为了

我们说了 代表的是任务相关的更新量,平均来说它必然是有利于任务的(否则原来的优化器就是有缺陷的了),所以平均来说应该有 。这里我们进一步假设,存在一个 ,使得 ,于是我们有

根据近似 (8) 我们有 ,代入上式得到

一步一步往前递推,可以得到

可以看出右端的指数必然是单调递减的,它是一个衰减函数。现在我们再看近似 (8) 它有两个参数 要调,但只有一个(近似)等式。为了使 能够同等程度地衰减,我们设 ,于是解得

这就是本文推出的 的变化规律。当然,变化规律是有了,可是还有四个参数 要确定,其中 相对来说比较简单,直接设 问题也不大,但即便这样还有三个参数要确定。


尺度预判

根据定义,,也就是初始化参数与目标参数的距离,可以理解为参数的变化尺度,它有几种不同的情况。
第一种,参数是矩阵乘法核,比如全连接层、卷积层的 kernel 矩阵,它们的初始化一般是“均值为 0、方差为 ”( 取决于 shape)的随机初始化,这样如果 ,那么我们就可以估算出
另外,这类参数有一个特点,就是在合理的初始化下,训练完成后参数的均值方差也不会有太大变化,至少量级是一致的,因此也可以认为 ,而因为初始化是随机的,所以 ,因此

第二种,参数是加性偏置项,比如全连接层、卷积层的 bias 向量,以及 Normalization 层的 向量,这些参数一般是“全零初始化”,所 如果我们根据经验预测训练好的模型偏置项都在 附近,那么也可以估计出 ,Amos 原论文取了 。最后还有 Normalization 层的 向量,它一般是“全 1 初始化”,训练完成后也是在 1 附近,不妨假设误差为 ,那么也可以估算出 。这里的 k 都是指向量维度。
可以看出, 的结果都有一个共性,那就是都可以写成 ,其中 是我们对参数变化尺度的一个预判。乘性矩阵的 可以直接取初始化的标准差,加性偏置或者 向量可以直接简单地取 ,或者有其他特殊参数的再做特殊处理。


分离尺度
现在我们来看完整的更新量,根据式 (13),有

其中 是一个单位向量,控制更新方向, 部分是一个衰减项,我们可以先不管它,所以更新量的模长由 控制。
回到文章开头的第一个问题“学习率如何适应不同初始化和参数化?”,很明显,直观想法应该就是变化尺度大的参数每一步的更新量应该更大,或者直接简单地正比于变化尺度,而变化尺度我们刚才估计了,可以用 来描述,所以我们认为应该有 ,其中 是全局的初始学习率。反过来解得 ,代入式 (13) 得到

其中 代表了每一步的相对更新幅度(全局学习率),这一步没啥推导空间了,一般取 左右就行,如果任务简单也可以取到 在上一节已经做了估计,大概是 代表参数尺度,不同参数不一样,我们正是通过它把参数尺度显式地分离了出来,从而达到了自适应参数尺度的效果(更新量正比 )。特别地,如果将上式的 换成 ,那么就是 LAMB 优化器。
从这里也可以看出,如果 的初始化均值不是 0(像 向量),用 替代 是会有问题的,所以 LAMB 的做法是直接不对这些参数的更新量进行变换(即保留原来的更新规则)。


解析近似


其实目前的结果已经适合编程实现了,只是参数 p 不好调罢了。为了进一步看出参数 p 是怎么影响衰减函数的,我们可以进一步求出 的解析近似!
在式 (16) 的 两边乘以 ,然后两边开平方,得到

将指数的求和 记为 ,那么上式就对应差分方程

此时衰减函数就是 。为了求渐近近似,我们用导数代替差分(参考《差分方程的摄动法》[6]),得到

这是个简单的微分方程,可以解得(结合

这就是衰减函数的显式解,表明超参数应该按照步数的平方反比衰减,具体代入式 (16) 后得到

这个显式解不但能让编程实现更方便,还使得 p 的含义更为清晰。比如我们希望学习率在 T 步后就降低为原来的一半,那么就有 ,从中解得

至于 T 应该是多少,这依赖于任务难度和数据量,也没有太大推导空间了。


文章小结

本文借鉴了 Amos 优化器的思路,推导了一些关于学习率和权重衰减率的结果 (21),这些结果可以即插即用地应用到现有优化器中,能一定程度上简化调参难度。


参考文献

[1] https://arxiv.org/abs/2210.11693

[2] https://kexue.fm/archives/7180

[3] https://kexue.fm/archives/8620

[4] https://kexue.fm/archives/7180

[5] https://kexue.fm/archives/7076

[6] https://kexue.fm/archives/3889



更多阅读



#投 稿 通 道#

 让你的文字被更多人看到 



如何才能让更多的优质内容以更短路径到达读者群体,缩短读者寻找优质内容的成本呢?答案就是:你不认识的人。


总有一些你不认识的人,知道你想知道的东西。PaperWeekly 或许可以成为一座桥梁,促使不同背景、不同方向的学者和学术灵感相互碰撞,迸发出更多的可能性。 


PaperWeekly 鼓励高校实验室或个人,在我们的平台上分享各类优质内容,可以是最新论文解读,也可以是学术热点剖析科研心得竞赛经验讲解等。我们的目的只有一个,让知识真正流动起来。


📝 稿件基本要求:

• 文章确系个人原创作品,未曾在公开渠道发表,如为其他平台已发表或待发表的文章,请明确标注 

• 稿件建议以 markdown 格式撰写,文中配图以附件形式发送,要求图片清晰,无版权问题

• PaperWeekly 尊重原作者署名权,并将为每篇被采纳的原创首发稿件,提供业内具有竞争力稿酬,具体依据文章阅读量和文章质量阶梯制结算


📬 投稿通道:

• 投稿邮箱:hr@paperweekly.site 

• 来稿请备注即时联系方式(微信),以便我们在稿件选用的第一时间联系作者

• 您也可以直接添加小编微信(pwbot02)快速投稿,备注:姓名-投稿


△长按添加PaperWeekly小编



🔍


现在,在「知乎」也能找到我们了

进入知乎首页搜索「PaperWeekly」

点击「关注」订阅我们的专栏吧


·

您可能也对以下帖子感兴趣

文章有问题?点此查看未经处理的缓存