这段时间笔者一直在实验《Google新搜出的优化器Lion:效率与效果兼得的“训练狮”》所介绍的Lion优化器。之所以对Lion饶有兴致,是因为它跟笔者之前的关于理想优化器的一些想法不谋而合,但当时笔者没有调出好的效果,而Lion则做好了。

相比标准的Lion,笔者更感兴趣的是它在$\beta_1=\beta_2$时的特殊例子,这里称之为“Tiger”。Tiger只用到了动量来构建更新量,根据《隐藏在动量中的梯度累积:少更新几步,效果反而更好?》的结论,此时我们不新增一组参数来“无感”地实现梯度累积!这也意味着在我们有梯度累积需求时,Tiger已经达到了显存占用的最优解,这也是“Tiger”这个名字的来源(Tight-fisted Optimizer,抠门的优化器,不舍得多花一点显存)。

此外,Tiger还加入了我们的一些超参数调节经验,以及提出了一个防止模型出现NaN(尤其是混合精度训练下)的简单策略。我们的初步实验显示,Tiger的这些改动,能够更加友好地完成模型(尤其是大模型)的训练。

基本形式 #

Tiger的更新规则为
\begin{equation}\text{Tiger}:=\left\{\begin{aligned}
&\boldsymbol{m}_t = \beta \boldsymbol{m}_{t-1} + \left(1 - \beta\right) \boldsymbol{g}_t \\
&\boldsymbol{\theta}_t = \boldsymbol{\theta}_{t-1} - \eta_t \left[\text{sign}(\boldsymbol{m}_t) \color{skyblue}{ + \lambda_t \boldsymbol{\theta}_{t-1}}\right] \\
\end{aligned}\right.\end{equation}
相比Lion,它就是选择了参数$\beta_1 = \beta_2 = \beta$;相比SignSGD,它则是新增了动量和权重衰减。

参考实现:

下表对比了Tiger、Lion和AdamW的更新规则:
\begin{array}{c|c|c}
\hline
\text{Tiger} & \text{Lion} & \text{AdamW} \\
\hline
{\begin{aligned}
&\boldsymbol{m}_t = \beta \boldsymbol{m}_{t-1} + \left(1 - \beta\right) \boldsymbol{g}_t \\
&\boldsymbol{\theta}_t = \boldsymbol{\theta}_{t-1} - \eta_t \left[\text{sign}(\boldsymbol{m}_t) \color{skyblue}{ + \lambda_t \boldsymbol{\theta}_{t-1}}\right] \\
\end{aligned}} &
{\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}} &
{\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}} \\
\hline
\end{array}

可见Tiger是三者之中的极简者。

超参设置 #

尽管Tiger已经相当简化,但仍有几个超参数要设置,分别是滑动平均率$\beta$、学习率$\eta_t$以及权重衰减率$\lambda_t$,下面我们分别讨论这几个参数的选择。

滑动率 #

比较简单的是滑动平均的衰减率$\beta$。我们知道,在基本形式上Tiger相当于Lion取$\beta_1 = \beta_2 = \beta$的特殊情形,那么一个直觉是Tiger应当取$\beta=\frac{1}{2}(\beta_1 + \beta_2)$。在Lion的原论文中,对于CV任务有$\beta_1=0.9,\beta_2=0.99$,所以我们建议CV任务取$\beta = 0.945$;而对于NLP任务则有$\beta_1=0.95,\beta_2=0.98$,所以我们建议NLP任务取$\beta = 0.965$。

学习率 #

对于学习率,Tiger参考了AmosLAMB等工作,将学习率分两种情况设置。第一种是线性层的bias项和Normalization的beta、gamma参数,这类参数的特点是运算是element-wise,我们建议学习率选取为全局相对学习率$\alpha_t$的一半;第二种主要就是线性层的kernel矩阵,这类参数的特点是以矩阵的身份去跟向量做矩阵乘法运算,我们建议学习率选取为全局相对学习率$\alpha_t$乘以参数本身的$\text{RMS}$(Root Mean Square):
\begin{equation}\eta_t = \left\{\begin{aligned}
&\alpha_t \times 0.5, &\boldsymbol{\theta} \in \{bias, beta, gamma\}\\[5pt]
&\alpha_t \times \text{RMS}(\boldsymbol{\theta}_{t-1}), &\boldsymbol{\theta} \not\in \{bias, beta, gamma\}
\end{aligned}\right.\end{equation}
其中
\begin{equation}\text{RMS}(\boldsymbol{\theta})=\sqrt{\frac{1}{k}\sum_{i=1}^k \theta_i^2},\quad \boldsymbol{\theta}=(\theta_1,\theta_2,\cdots,\theta_k)\end{equation}
这样设置的好处是我们把参数的尺度分离了出来,使得学习率的调控可以交给一个比较通用的“全局相对学习率”$\alpha_t$——大致可以理解为每一步的相对学习幅度,是一个对于模型尺度不是特别敏感的量。

换句话说,我们在base版模型上调好的$\alpha_t$,基本上可以不改动地用到large版模型。注意$\alpha_t$带有下标$t$,所以它包含了整个学习率的schedule,包括Wamrup以及学习率衰减策略等,笔者的设置经验是$\max(\alpha_t)\in[0.001,0.002]$,至于怎么Warmup和衰减,那就是大家根据自己的任务而设了,别人无法代劳。笔者给的tiger实现,内置了一个分段线性学习率策略,理论上可以用它模拟任意的$\alpha_t$。

衰减率 #

最后是权重衰减率$\lambda_t$,这个Lion论文最后一页也给出了一些参考设置,一般来说$\lambda_t$也就设为常数,笔者常用的是0.01。特别的是,不建议对前面说的bias、beta、gamma这三类参数做权重衰减,或者即便要做,$\lambda_t$也要低一个数量级以上。因为从先验分布角度来看,权重衰减是参数的高斯先验,$\lambda_t$跟参数方差是反比关系,而bias、beta、gamma的方差显然要比kernel矩阵的方差大,所以它们的$\lambda_t$应该更小。
\begin{equation}\lambda_t = \left\{\begin{aligned}
&0, &\boldsymbol{\theta} \in \{bias, beta, gamma\}\\[5pt]
&constant > 0, &\boldsymbol{\theta} \not\in \{bias, beta, gamma\}
\end{aligned}\right.\end{equation}

梯度累积 #

对于很多算力有限的读者来说,通过梯度累积来增大batch_size是训练大模型时不可避免的一步。标准的梯度累积需要新增一组参数,用来缓存历史梯度,这意味着在梯度累积的需求下,Adam新增的参数是3组,Lion是2组,而即便是不加动量的AdaFactor也有1.x组(但说实话AdaFactor不加动量,收敛会慢很多,所以考虑速度的话,加一组动量就变为2.x组)。

而对于Tiger来说,它的更新量只用到了动量和原参数,根据《隐藏在动量中的梯度累积:少更新几步,效果反而更好?》,我们可以通过如下改动,将梯度累积内置在Tiger中:
\begin{equation}\text{Tiger}:=\left\{\begin{aligned} &\boldsymbol{m}_t = \big[(\beta - 1)\chi_{(t-1)/k} + 1\big] \boldsymbol{m}_{t-1} + \frac{1}{k}\left(1 - \beta\right) \boldsymbol{g}_t \\ &\boldsymbol{\theta}_t = \boldsymbol{\theta}_{t-1} - \chi_{t/k}\eta_t \left[\text{sign}(\boldsymbol{m}_t) \color{skyblue}{ + \lambda_t \boldsymbol{\theta}_{t-1}}\right] \\ \end{aligned}\right.\end{equation}
这里的$\chi_{t/k}$是判断$t$能否被$k$整除的示性函数
\begin{equation}\chi_{t/k} = \left\{ \begin{aligned}&1,\quad t \equiv 0\,(\text{mod}\, k) \\
&0,\quad t \not\equiv 0\,(\text{mod}\, k)
\end{aligned}\right.\end{equation}
可以看到,这仅仅相当于修改了滑动平均率$\beta$和学习率$\eta_t$,几乎不增加显存成本,整个过程是完全“无感”的,这是笔者认为的Tiger的最大的魅力。

需要指出的是,尽管Lion跟Tiger很相似,但是Lion并不能做到这一点,因为$\beta_1\neq\beta_2$时,Lion的更新需要用到动量以及当前批的梯度,这两个量需要用不同的参数缓存,而Tiger的更新只用到了动量,因此满足这一点。类似滴,SGDM优化器也能做到一点,但是它没有$\text{sign}$操作,这意味着学习率的自适应能力不够好,在Transformer等模型上的效果通常不如意(参考《Why are Adaptive Methods Good for Attention Models?》)。

全半精度 #

对于大模型来说,混合精度训练是另一个常用的“利器”(参考《在bert4keras中使用混合精度和XLA加速训练》)。混合精度,简单来说就是模型计算部分用半精度的FP16,模型参数的储存和更新部分用单精度的FP32。之所以模型参数要用FP32,是因为担心更新过程中参数的更新量过小,下溢出了FP16的表示范围(大致是$6\times 10^{-8}\sim 65504$),导致某些参数长期不更新,模型训练进度慢甚至无法正常训练。

然而,Tiger(Lion也一样)对更新量做了$\text{sign}$运算,这使得理论上我们可以全用半精度训练!分析过程并不难。首先,只要对Loss做适当的缩放,那么可以做到梯度$\boldsymbol{g}_t$不会溢出FP16的表示范围;而动量$\boldsymbol{m}_t$只是梯度的滑动平均,梯度不溢出,它也不会溢出,$\text{sign}(\boldsymbol{m}_t)$只能是$\pm 1$,更加不会溢出了;之后,我们只需要保证学习率不小于$6\times 10^{-8}$,那么更新量就不会下溢了,事实上我们也不会将学习率调得这么小。因此,Tiger的整个更新过程都是在FP16表示范围内的,因此理论上我们可以直接用全FP16精度训练而不用担心溢出问题。

防止NaN #

然而,笔者发现对于同样的配置,在FP32下训练正常,但切换到混合精度或者半精度后有时会训练失败,具体表现后Loss先降后升然后NaN,这我们之前在《在bert4keras中使用混合精度和XLA加速训练》也讨论过。虽然有一些排查改进的方向(比如调节epsilon和无穷大的值、缩放loss等),但有时候把该排查的都排查了,还会出现这样的情况。

经过调试,笔者发现出现这种情况时,主要是对于某些batch梯度变为NaN,但此时模型的参数和前向计算还是正常的。于是笔者就想了个简单的应对策略:对梯度出现NaN时,跳过这一步更新,并且对参数进行轻微收缩,如下
\begin{equation}\text{Tiger}:=\left\{\begin{aligned}
&\boldsymbol{m}_t = \boldsymbol{m}_{t-1} \\
&\boldsymbol{\theta}_t = (\boldsymbol{\theta}_{t-1} - c)\times s+ c \\
\end{aligned}\right. \quad if\,\,\boldsymbol{g}_t = \text{NaN}\end{equation}
其中$s\in(0, 1)$代表收缩率,笔者取$s=0.99$,$c$则是参数的初始化中心,一般就是gamma取1,其他参数都是0。经过这样处理后,模型的loss会有轻微上升,但一般能够恢复正常训练,不至于从头再来。个人的实验结果显示,这样处理能够缓解一部分NaN的问题。

当然,该技巧一般的使用场景是同样配置下FP32能够正常训练,并且已经做好了epsilon、无穷大等混合精度调节,万般无奈之下才不得已使用的。如果模型本身超参数设置有问题(比如学习率过大),连FP32都会训练到NaN,那么就不要指望这个技巧能够解决问题了。此外,有兴趣的读者,还可以尝试改进这个技巧,比如收缩之后可以再加上一点噪声来增加参数的多样性,等等。

实验结果 #

不考虑梯度累积带来的显存优化,Tiger就是Lion的一个特例,可以预估Tiger的最佳效果肯定是不如Lion的最佳效果的,那么效果下降的幅度是否在可接受范围内呢?综合到目前为止多方的实验结果,笔者暂时得出的结论是:
$$\begin{aligned}
&\text{效果}\color{red}{(\uparrow)}\text{:}\quad\text{Lion} \geq \text{Tiger} \geq \text{AdamW} \approx \text{LAMB} \\
&\text{显存}\color{red}{(\downarrow)}\text{:}\quad\text{Tiger} < \text{Lion} < \text{AdamW} = \text{LAMB} \\
\end{aligned}$$
也就是说,考虑效果Lion最优,考虑显存占用Tiger最优(启用梯度累积时),效果上Tiger不逊色于AdamW,所以Tiger替代AdamW时没有太大问题的。

具体实验结果包括几部分。第一部分实验来自Lion的论文《Symbolic Discovery of Optimization Algorithms》,论文中的Figure 12对比了Lion、Tiger、AdamW在不同尺寸的语言模型上的效果:

Lion、Tiger(Ablation)、AdamW在语言模型任务上的对比

Lion、Tiger(Ablation)、AdamW在语言模型任务上的对比

这里的Ablation0.95、Ablation0.98,就是Tiger的$\beta$分别取0.95、0.98。可以看到,对于small级别模型,两个Tiger持平AdamW,而在middle和large级别上,两个Tiger都超过了AdamW。但正如前面所说,$\beta$取两者的均值0.965,有可能还会有进一步的提升。

至于在CV任务上,原论文给出了Table 7:

Lion、Tiger(Ablation)、AdamW在图像分类任务上的对比

Lion、Tiger(Ablation)、AdamW在图像分类任务上的对比

同样地,这里的Ablation0.9、Ablation0.99,就是Tiger的$\beta$分别取0.9、0.99。在这个表中,Tiger跟AdamW有明显差距。但是考虑到作者只实验了0.9、0.99两个$\beta$,而笔者推荐的是$\beta=0.945$,所以笔者跟原作者取得了联系,请他们做了补充实验,他们回复的结果是“$\beta$分别取0.92、0.95、0.98时,在ViT-B/16上ImageNet的结果都是80.0%左右”,那么对比上图,就可以确定在精调$\beta$时,在CV任务上Tiger应该也可以追平AdamW的。

最后是笔者自己的实验。笔者常用的是LAMB优化器,它的效果基本跟AdamW持平,但相对更稳定,而且对不同的初始化适应性更好,因此笔者更乐意使用LAMB。特别地,LAMB的学习率设置可以完全不改动地搬到Tiger中。笔者用Tiger重新训练了之前的base版GAU-α模型,训练曲线跟之前的对比如下:

笔者在GAU-α上的对比实验(loss曲线)

笔者在GAU-α上的对比实验(loss曲线)

笔者在GAU-α上的对比实验(accuracy曲线)

笔者在GAU-α上的对比实验(accuracy曲线)

可以看到,Tiger确实可以取得比LAMB更优异的表现。

未来工作 #

Tiger还有改进空间吗?肯定有,想法其实有很多,但都没来得及一一验证,大家有兴趣的可以帮忙继续做下去。

《Google新搜出的优化器Lion:效率与效果兼得的“训练狮”》中,笔者对$\text{sign}$运算的评价是

Lion通过$\text{sign}$操作平等地对待了每一个分量,使得模型充分地发挥了每一个分量的作用,从而有更好的泛化性能。如果是SGD,那么更新的大小正比于它的梯度,然而有些分量梯度小,可能仅仅是因为它没初始化好,而并非它不重要,所以Lion的$\text{sign}$操作算是为每个参数都提供了“恢复活力”甚至“再创辉煌”的机会。

然而,细思之下就会发现,这里其实有一个改进空间。“平等地对待了每一个分量”在训练的开始阶段是很合理的,它保留了模型尽可能多的可能。然而,如果一个参数长时间的梯度都很小,那么很有可能这个参数真的是“烂泥扶不上墙”,即已经优化到尽头了,这时候如果还是“平等地对待了每一个分量”,那么就对那些梯度依然较大的“上进生”分量不公平了,而且很可能导致模型震荡。

一个符合直觉的想法是,优化器应该随着训练的推进,慢慢从Tiger退化为SGD。为此,我们可以考虑将更新量设置为
\begin{equation}\boldsymbol{u}_t = \text{sign}(\boldsymbol{m}_t) \times |\boldsymbol{m}_t|^{1-\gamma_t}\end{equation}
这里的绝对值和幂运算都是element-wise的,$\gamma_t$是从1到0的单调递减函数,当$\gamma_t=1$时对应Tiger,当$\gamma_t = 0$时对应SGDM。

可能读者会吐槽这里多了$\gamma_t$这个schedule要调整,问题变得复杂很多。确实如此,如果将它独立地进行调参,那么确实会引入过多的复杂度了。但我们不妨再仔细回忆一下,抛开Warmup阶段不算,一般情况下相对学习率$\alpha_t$不正是一个单调递减至零的函数?我们是否可以借助$\alpha_t$来设计$\gamma_t$呢?比如$\alpha_t/\alpha_0$不正好是一个从1到0的单调递减函数?能否用它来作为$\gamma_t$?当然也有可能是$(\alpha_t/\alpha_0)^2$、$\sqrt{\alpha_t/\alpha_0}$更好,调参空间还是有的,但至少我们不用重新设计横跨整个训练进程的schedule了。

更发散一些,既然有时候学习率我们也可以用非单调的schedule(比如带restart的cosine annealing),那么$\gamma_t$我们是否也可以用非单调的(相当于Tiger、SGDM反复切换)?这些想法都有待验证。

文章小结 #

在这篇文章中,我们提出了一个新的优化器,名为Tiger(Tight-fisted Optimizer,抠门的优化器),它在Lion的基础上做了一些简化,并加入了我们的一些超参数经验。特别地,在需要梯度累积的场景下,Tiger可以达到显存占用的理论最优(抠)解!

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

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

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

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

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

苏剑林. (Mar. 07, 2023). 《Tiger:一个“抠”到极致的优化器 》[Blog post]. Retrieved from https://www.kexue.fm/archives/9512

@online{kexuefm-9512,
        title={Tiger:一个“抠”到极致的优化器},
        author={苏剑林},
        year={2023},
        month={Mar},
        url={\url{https://www.kexue.fm/archives/9512}},
}