本文将解释随机梯度下降(SGD)等必要背景,是LARS建立的基础。
在现代深度学习中,周期时间通常为数小时或数天。加速训练,数据并行性的最简单方法是在GPU和机器之间分发模型的副本,并让每个副本计算训练数据的分片上的损失。然后可以使用单个参数服务器或诸如ring all-reduce之类的更高级的东西来累积来自这些损失的梯度。
反向传播梯度后,重复更新的模型。但随着批量大小的增加,梯度会变得不稳定,导致训练发散。
SGD
对于SGD,标准更新规则是

打破这一点,它说新模型权重是现有的减去学习率乘以给定示例的现有权重的梯度或批次的平均值。问题:
- 我们不知道真正的梯度,我们只有嘈杂的估计。噪音可能会导致虚假更新
- 梯度比例未标准化,因此我们的更新可能会随着时间的推移而消失或爆炸
Adam
在Adam,我们保持提督的移动平均值及其变化:

其中是移动平均值,是移动的未中心方差,β1是平均值的插值常数,β2是未中心方差的插值常数,∇L是损失的梯度。指数中的括号表示它实际上不是指数,而是时间步长。这看起来有些可怕,但需要注意的重要一点是,和都只是梯度的线性插值(β* x0 +(1 – β)* x1)及其方差,这为我们提供了移动平均线每个。β越高,我们更新每个新样本的移动平均值越少,从而平滑我们对批次间梯度的均值和方差的估计。这里是我们在不同测试数据集的噪声数据集上获得多少平滑的可视化。

各种β强度的线性插值
我们需要更高的β,但如果我们使用具有一致数据的大批量,我们可能需要更少的β。
上面的移动平均值的问题是,当算法首次初始化时,移动平均值为0。这导致汇总统计数据比前几个时间步长更接近0,如果β接近1,因为我们是从上一步中获取大部分权重。这种效果在上面的β= 0.99图中特别明显。
我们通过去偏差来解决这个问题

结果如下:

具有去偏移的线性插值
问题不会消失,但它会好得多。
为了插入一些数字,如果β=0.9,在第一次迭代时,去偏差将该值乘以1 /1-0.9¹= 10.然后当我们线性插值时,βx0 +(1 – β)x1,第一项为βx0=0 ,去偏差因子10将在第二项中抵消掉(1-β)= 0.1,所以我们完全使用新值x1。在几步之后,去偏差因子将收敛到1。下图显示了去偏差项消失所需的步数(注意y轴的差异):

随着时间的推移可视化去偏差因子
现在最后的参数更新是

分子说”对于每个参数,在该参数的梯度方向上迈出一步。”分母说”用标准偏差将步骤标准化”。
直观的解释是:
当我们第一次开始更新参数时,我们可能会偏离。如果梯度都指向不同的方向(高方差),我们将采取一个小而谨慎的步骤。相反,如果所有的梯度都告诉我们向同一方向移动,那么方差就会很小,所以我们会朝这个方向迈出更大的一步。
无论哪种方式,如果所有梯度的比例都很大,那么当我们使用未中心化的方差时,常数因子将在我们分开时抵消。随着训练稳定并且损失接近0,平均值将接近0,因此更新将自动变得更精细。
分母中的ε表示有效地设定了步长大小的上限当噪声方差接近零时你采取。
该比率m / sqrt(v)可能看起来像μ/σ,这是信噪比,但该解释仅适用于标量。
LARS
随着批量大小的增加,每个时期的迭代次数减少。为了收敛相同数量的数据集迭代,我们可以通过提高学习率来弥补。然而,随着学习率的提高,训练变得更加不稳定。SOTA是使用学习率热身,但这只能帮助到某一点,无论如何学习将开始分化。热身是真正问题的补丁:梯度必须是嘈杂的。
分层自适应速率调整(LARS)的作者解释了他们解决这个问题的技巧:
为了分析大LR的训练稳定性,我们测量了层权重的范数与梯度更新的范数之间的比率。我们观察到如果这个比例太高,训练可能会变得不稳定。另一方面,如果比率太小,则权重变化不够快。
他们将这个比率称为”信任比率”。当它更高时,梯度变化更快,反之亦然。由于我们现在可以对每个步骤更加自信,因此不再需要在学习速率计划中经常使用的警示预热,并且我们可以扩展到更大的批量大小而不会发散。

逐层学习率λ是全局学习率η乘以层权重的范数与层梯度的范数的比率。如果我们使用权重衰减,我们可以将它添加到分母中。当我们将其插入SGD时,分母最终将梯度标准化为单位范数,这有助于避免分歧。
分子是权重的范数,因为随着网络的加深,具有零均值单位方差(ZMUV)权重变得很重要。这是因为在每一层,这些权重相乘,所以如果它与ZMUV不同,则值可能会爆炸或消失。当权重很小时,我们采取一小步。当权重很大时,我们会迈出更大的一步。结合权重衰减,这有助于我们稳定地迈向ZMUV权重。
在训练开始时,层应该输出ZMUV,因此上面的分子将为0或接近它。我们采取的任何步骤都可能很小。相比之下,分母可能会很大,因为当一切都错了时,梯度很大。这样,随着权重的增加,我们自然会热身。当我们接近0损失时,梯度将很小,因此信任比率将使学习率保持高达10倍(由于削波)高于没有信任比率,从而使我们不会过早放弃达到最佳状态。
LAMB
LAMB代表”用于批量训练的分层自适应时刻优化器“。它对LARS进行了一些小的改动
- 如果信任比的分子(r 1以下)或分母(r 2以下)为0,则使用1代替。这部分很难阅读,利用代码讲解。
- 固定权重衰减:在LARS中,信任比率的分母是|∇L| +β| w |,而在LAMB中,它是|∇L+βw|。这保留了更多信息。
- 他们使用Adam更新规则而不是使用SGD更新规则。
- 将信任比率限制为10。
总而言之,LAMB的信任比率是

最后一行是逐层LAMB更新规则。2是具有权重衰减的Adam更新规则的范数,ηᴸ是通过信任比率调整的分层学习率。总的来说,这个方法可以概括为应用于Adam的LARS,因为它只是将旧的更新步骤乘以信任比率。
作者没有报告LAMB是否提高了LARS的ImageNet训练性能,并且他们没有将BAR与LARS进行比较,因此很难说这些变化有多大差异,但实现非常简单。
实验
为了更好地了解正在发生的事情,本文在Pytorch中实现了 LAMB。本文在MNIST上进行了一系列实验,发现在Adam发散的地方,LAMB一直在运行。本文之所以选择MNIST是因为它足够小,可以试用CPU,但这意味着我们无法看到任何融合改进。
下面,本文将Adam(下面的蓝色)和LAMB(下面的红色)与学习率0.01和beta 0.9,0.99进行比较。它们非常相似,但LAMB更好地推广了测试精度。

对于Adam(蓝色)与LAMB(红色),MNIST训练损失和测试精度随时间推移
为了弄清楚幕后发生了什么,本文想看一下信任比的分层组成部分,所以本文在每几批后记录r,r1和r2的每个值作为直方图。对于Adam,本文计算了值,但不要在任何地方使用它们。
如何解读下面的图表:Y轴显示时间步长,第一个在顶部,X轴是直方图桶,Z轴是直方图频率。


在MNIST上的Adam(蓝色)与LAMB(红色)上的LAMB参数直方图
你可以看到,在两侧,r开始明显低于1。在LAMB方面,这会在所有层上创建一个自然的预热期。然后,随着一些层开始获得更大的权重和更稳定的梯度,r鼓励他们采取更大的步骤。这夸大了与Adam基线相关的规范。
对于下一个实验,本文将LAMB与学习率0.1和0.01进行了比较。Adam通常以学习率.01和0.1完全收敛,所以我们不会在这里进行比较。左侧(蓝色)学习率=0 .01,右侧(绿色)学习率= 0.1。在右边,它在预热期间几乎立即收敛,但随后几层重量开始爆炸(参见X轴刻度的差异)并且它发散。

为了解决权重失控的问题,本文在右下方添加了0.01的权重衰减。训练没有分歧!一般来说,信任比率在低于1时保持学习缓慢,而在左上方更舒适的政权中,它高达4.5。

总结
- 随着学习率的提高,vanilla SGD变得不稳定。
- LARS通过逐层信任比来调整SGD学习率,该信任比将梯度和权重标准化。
- Adam用脱离方法调整更新,通过去偏差的方差归一化。
- LAMB通过更准确的分层,削减信任比来调整Adam的学习率。
结合所有这些技术,我们可以训练大批量学习率,将BERT的训练时间减少100倍!