优化算法的符号发现.

本文作者通过数千TPU小时的算力搜索并结合人工干预,得到了一个速度更快、显存更省的优化器LionEvoLved Sign Momentum),并在计算机视觉、自然语言处理等诸多任务上做了充分的实验,表明Lion比目前主流的AdamW等优化器有着更好的效果。

Lion的更新形式如下:

\[\begin{aligned} u_t &= \text{sign}( \beta_1m_{t-1}+(1-\beta_1)g_t ) + \lambda_t\theta_{t} \\ θ_t&=θ_{t-1}-\eta_t u_t \\ m_t &= \beta_2m_{t-1}+(1-\beta_2)g_t \end{aligned}\]

相比于主流的AdamW等优化器,Lion少缓存了一组参数(梯度平方值),所以更省显存;并且去除了更新过程中计算量较大的除法和根号运算,所以计算速度更快。Lion把动量的更新放在变量的更新之后,在大量实验中显示出优越性。

Lion通过使用符号函数sign这个操作引入了额外的噪声(相比于准确的浮点值),它使得模型进入了损失更平坦(但未必更小)的区域,从而泛化性能更好。另一方面,额外引入的噪声加剧了优化过程的噪声,在batch size较小时(小于64),可能噪声过量导致效果恶化,参数设置不当时容易出现损失变大等发散情况。

Lion需要设置的超参数为$\beta_1,\beta_2,\lambda,\eta$。在计算机视觉任务中设置$\beta_1=0.9,\beta_2=0.99$,在自然语言任务中设置$\beta_1=0.95,\beta_2=0.98$。对于学习率$\eta$和权重衰减率$\lambda$,由于参数更新量$u$的每个分量的绝对值都是$1$,更新幅度比其他优化算法大,因此学习率要缩小一些;随着学习率降低,为使权重衰减的幅度保持不变,权重衰减率应该放大。作者给出了不同实验的超参数参考值: