昨天在Arixv上发现了Google新发的一篇论文《Symbolic Discovery of Optimization Algorithms》,主要是讲自动搜索优化器的,咋看上去没啥意思,因为类似的工作也有不少,大多数结果都索然无味。然而,细读之下才发现别有洞天,原来作者们通过数千TPU小时的算力搜索并结合人工干预,得到了一个速度更快、显存更省的优化器Lion(EvoLved Sign Momentum,不得不吐槽这名字起得真勉强),并在图像分类、图文匹配、扩散模型、语言模型预训练和微调等诸多任务上做了充分的实验,多数任务都显示Lion比目前主流的AdamW等优化器有着更好的效果。

更省显存还更好效果,真可谓是鱼与熊掌都兼得了,什么样的优化器能有这么强悍的性能?本文一起来欣赏一下论文的成果。

先说结果 #

本文主要关心搜索出来的优化器本身,所以关于搜索过程的细节就不讨论了,对此有兴趣读者自行看原论文就好。Lion优化器的更新过程为
\begin{equation}\text{Lion}:=\left\{\begin{aligned}
&\boldsymbol{u}_t = \text{sign}\big(\beta_1 \boldsymbol{m}_{t-1} + \left(1 - \beta_1\right) \boldsymbol{g}_t\big) \\
&\boldsymbol{\theta}_t = \boldsymbol{\theta}_{t-1} - \eta_t (\boldsymbol{u}_t \color{skyblue}{ + \lambda_t \boldsymbol{\theta}_{t-1}}) \\
&\boldsymbol{m}_t = \beta_2 \boldsymbol{m}_{t-1} + \left(1 - \beta_2\right) \boldsymbol{g}_t
\end{aligned}\right.\end{equation}
其中$\boldsymbol{g}_t = \nabla_{\boldsymbol{\theta}} L(\boldsymbol{\theta}_{t-1})$是损失函数的梯度,$\text{sign}$是符号函数,即正数变为1、负数变为-1。我们可以对比一下目前的主流优化器AdamW的更新过程
\begin{equation}\text{Adam}\color{skyblue}{\text{W}}:=\left\{\begin{aligned}
&\boldsymbol{m}_t = \beta_1 \boldsymbol{m}_{t-1} + \left(1 - \beta_1\right) \boldsymbol{g}_t\\
&\boldsymbol{v}_t = \beta_2 \boldsymbol{v}_{t-1} + \left(1 - \beta_2\right) \boldsymbol{g}_t^2\\
&\hat{\boldsymbol{m}}_t = \boldsymbol{m}_t\left/\left(1 - \beta_1^t\right)\right.\\
&\hat{\boldsymbol{v}}_t = \boldsymbol{v}_t\left/\left(1 - \beta_2^t\right)\right.\\
&\boldsymbol{u}_t =\hat{\boldsymbol{m}}_t\left/\left(\sqrt{\hat{\boldsymbol{v}}_t} + \epsilon\right)\right.\\
&\boldsymbol{\theta}_t = \boldsymbol{\theta}_{t-1} - \eta_t (\boldsymbol{u}_t \color{skyblue}{ + \lambda_t \boldsymbol{\theta}_{t-1}})
\end{aligned}\right.\end{equation}
对比很明显,Lion相比AdamW参数更少(少了个$\epsilon$),少缓存了一组参数$\boldsymbol{v}$(所以更省显存),并且去掉了AdamW更新过程中计算量最大的除法和开根号运算(所以更快)。

在此之前,跟Lion最相似的优化器应该是SIGNUM,其更新过程为
\begin{equation}\text{SIGNUM}:=\left\{\begin{aligned}
&\boldsymbol{m}_t = \beta \boldsymbol{m}_{t-1} + \left(1 - \beta\right) \boldsymbol{g}_t \\
&\boldsymbol{u}_t = \text{sign}\big(\boldsymbol{m}_t\big) \\
&\boldsymbol{\theta}_t = \boldsymbol{\theta}_{t-1} - \eta_t \boldsymbol{u}_t \end{aligned}\right.\end{equation}
跟Lion一样,SIGNUM也用到了符号函数处理更新量,而且比Lion更加简化(等价于Lion在$\beta_1=\beta_2$和$\lambda_t=0$的特例),但是很遗憾,SIGNUM并没有取得更好的效果,它的设计初衷只是降低分布式计算中的传输成本。Lion的更新规则有所不同,尤其是动量的更新放在了变量的更新之后,并且在充分的实验中显示出了它在效果上的优势。

论文实验 #

本文开头就说了,Lion在相当多的任务上都做了实验,实验结果很多,下面罗列一些笔者认为比较关键的结果。

Lion在NLU和NLG任务上的结果,大部分都比AdamW、Adafactor优秀

Lion在NLU和NLG任务上的结果,大部分都比AdamW、Adafactor优秀

在视觉Transformer上Lion与众多优化器的对比

在视觉Transformer上Lion与众多优化器的对比

在CV的分类任务上,Lion收敛速度更快

在CV的分类任务上,Lion收敛速度更快

在NLP的自回归生成上,Lion的收敛速度更快

在NLP的自回归生成上,Lion的收敛速度更快

上右图是ImageNet上的训练曲线,显示Lion尽管验证集效果更好,但训练集上的效果未必会优于AdamW

上右图是ImageNet上的训练曲线,显示Lion尽管验证集效果更好,但训练集上的效果未必会优于AdamW

超参设置 #

看到论文效果如此惊人,笔者也跃跃欲试。在跑实验之前,自然需要了解一下各个超参的设置。首先是$\beta_1,\beta_2$,原论文自动搜索出来的结果是$\beta_1=0.9,\beta=0.99$,并在大部分实验中复用了这个组合,但是在NLP的任务上则使用了$\beta_1=0.95,\beta_2=0.98$这个组合(论文的详细实验配置在最后一页的Table 12)。

比较关键的学习率$\eta$和权重衰减率$\lambda$,由于Lion的更新量$\boldsymbol{u}$每个分量的绝对值都是1,这通常比AdamW要大,所以学习率要缩小10倍以上,才能获得大致相同的更新幅度;而由于学习率降低了,那么为了使权重衰减的幅度保持不变,权重衰减率应该要放大相应的倍数。原论文的最后一页给出了各个实验的超参数参考值,其中小模型(Base级别)上使用的是$\eta = 3\times 10^{-4}$和$\lambda=0.01$,大模型(参数10亿以上)则适当降低了学习率到$\eta = 2\times 10^{-4}$甚至$\eta = 10^{-4}$。

事实上,之前我们在《基于Amos优化器思想推导出来的一些“炼丹策略”》就推导过学习率和权重衰减率的一个组合方案,参考这个方案来设置是最方便的。在该方案中,更新量写为(记号跟前面的描述略有不同,但不至于混淆,应该就不强行统一了)
\begin{equation}\boldsymbol{\theta}_{t+1} = \boldsymbol{\theta}_t - (\alpha_t \boldsymbol{u}_t + \rho_t\boldsymbol{\theta}_t)\end{equation}
其中
\begin{equation}\alpha_t \approx \frac{\alpha_0\Vert\boldsymbol{\varepsilon}_0\Vert}{\Vert\boldsymbol{u}_t\Vert} \frac{1}{\kappa t + 1},\quad \rho_t \approx \frac{\alpha_0^2}{2q} \frac{1}{\kappa t + 1}\end{equation}
其中$\boldsymbol{u}_t$是原本的更新量;$\alpha_0$是(初始阶段)参数变化的相对大小,一般是$10^{-3}$级别,表示每步更新后参数模长的变化幅度大致是千分之一;$q$是一个超参数,没什么特殊情况可以设为1;$\kappa$是控制学习率衰减速度的超参数,可以根据训练数据大小等设置。

由于$\boldsymbol{u}_t$经过了$\text{sign}$运算,因此$\Vert\boldsymbol{u}_t\Vert=\sqrt{k}$,$k$是参数的维度;$\Vert\boldsymbol{\varepsilon}_0\Vert\approx\sqrt{k}\sigma$,这我们在《基于Amos优化器思想推导出来的一些“炼丹策略”》已经推导过了,其中$\sigma$是参数的变化尺度,对于乘性矩阵,$\sigma^2$就是它的初始化方差。所以,经过一系列简化之后,有
\begin{equation}\alpha_t \approx \frac{\alpha_0\sigma}{\kappa t + 1},\quad \rho_t \approx \frac{\alpha_0^2}{2(\kappa t + 1)}\end{equation}
这里的$\alpha_t$就是前面的$\eta_t$,而$\lambda_t = \rho_t / \alpha_t = \alpha_0 / 2\sigma$。按照BERT base的$d=768$来算,初始化方差的量级大致在$1/d$左右,于是$\sigma = \sqrt{1/d}\approx 0.036$,假设$\alpha_0$取$1.11 \times 10^{-3}$(为了给结果凑个整),那么按照上式学习率大约是$4\times 10^{-5}$、衰减率大约是$0.015$。在笔者自己的MLM预训练实验中,选取这两个组合效果比较好。

延伸思考 #

总体来看,Lion表现可圈可点,不管是原论文还是笔者自己的实验中,跟AdamW相比都有一战之力,再加上Lion更快以及更省显存的特点,或者可以预见未来的主流优化器将有它的一席之地。

自Adam提出以来,由于其快速收敛的特性成为了很多模型的默认优化器。甚至有学者提出,这个现象将反过来导致一个进化效应:所有的模型改进都在往Adam有利的方向发展,换句话说,由于我们选择了Adam作为优化器,那么就有可能将很多实际有效、但是在Adam优化器上无效的改动都抛弃了,剩下的都是对Adam有利的改进,详细的评价可以参考《NEURAL NETWORKS (MAYBE) EVOLVED TO MAKE ADAM THE BEST OPTIMIZER》。所以,在此大背景之下,能够发现比Adam更简单且更有效的优化器,是一件很了不起的事情,哪怕它是借助大量算力搜索出来的。

可能读者会有疑问:Lion凭啥可以取得更好的泛化性能呢?原论文的解释是$\text{sign}$这个操作引入了额外的噪声(相比于准确的浮点值),它使得模型进入了Loss更平坦(但未必更小)的区域,从而泛化性能更好。为了验证这一点,作者比较了AdamW和Lion训练出来的模型权重的抗干扰能力,结果显示Lion的抗干扰能力更好。然而,理论上来说,这只能证明Lion确实进入到了更平坦的区域,但无法证明该结果是$\text{sign}$操作造成的。不过,Adam发表这么多年了,关于它的机理也还没有彻底研究清楚,而Lion只是刚刚提出,就不必过于吹毛求疵了。

笔者的猜测是,Lion通过$\text{sign}$操作平等地对待了每一个分量,使得模型充分地发挥了每一个分量的作用,从而有更好的泛化性能。如果是SGD,那么更新的大小正比于它的梯度,然而有些分量梯度小,可能仅仅是因为它没初始化好,而并非它不重要,所以Lion的$\text{sign}$操作算是为每个参数都提供了“恢复活力”甚至“再创辉煌”的机会。事实上可以证明,Adam早期的更新量也接近于$\text{sign}$,只是随着训练步数的增加才逐渐偏离。

Lion是不是足够完美呢?显然不是,比如原论文就指出它在小batch_size(小于64)的时候效果不如AdamW,这也不难理解,本来$\text{sign}$已经带来了噪声,而小batch_size则进一步增加了噪声,噪声这个东西,必须适量才好,所以两者叠加之下,很可能有噪声过量导致效果恶化。另外,也正因为$\text{sign}$加剧了优化过程的噪声,所以参数设置不当时容易出现损失变大等发散情况,这时候可以尝试引入Warmup,或者增加Warmup步数。还有,Lion依旧需要缓存动量参数,所以它的显存占用多于AdaFactor,能不能进一步优化这部分参数量呢?暂时还不得而知。

文章小结 #

本文介绍了Google新提出的优化器Lion,它通过大量算力搜索并结合人工干预得出,相比主流的AdamW,有着速度更快且更省内存的特点,并且大量实验结果显示,它在多数任务上都有着不逊色于甚至优于AdamW的表现。

转载到请包括本文地址:https://kexue.fm/archives/9473

更详细的转载事宜请参考:《科学空间FAQ》

如果您还有什么疑惑或建议,欢迎在下方评论区继续讨论。

如果您觉得本文还不错,欢迎分享/打赏本文。打赏并非要从中获得收益,而是希望知道科学空间获得了多少读者的真心关注。当然,如果你无视它,也不会影响你的阅读。再次表示欢迎和感谢!

如果您需要引用本文,请参考:

苏剑林. (Feb. 16, 2023). 《Google新搜出的优化器Lion:效率与效果兼得的“训练狮” 》[Blog post]. Retrieved from https://kexue.fm/archives/9473

@online{kexuefm-9473,
        title={Google新搜出的优化器Lion:效率与效果兼得的“训练狮”},
        author={苏剑林},
        year={2023},
        month={Feb},
        url={\url{https://kexue.fm/archives/9473}},
}