LAMB Optimizer/zh
| Article | |
|---|---|
| Topic area | Machine Learning |
| Prerequisites | Adam Optimizer, Stochastic Gradient Descent, Gradient Descent |
概述
LAMB 优化器(Layer-wise Adaptive Moments for Batch training,分层自适应矩批量训练优化器)是一种一阶随机优化算法,旨在使深度神经网络能够在非常大的小批量大小下训练而不损失泛化能力。由 You 等人于 2019 年提出,LAMB 将 Adam 的逐参数自适应矩估计与受 LARS(Layer-wise Adaptive Rate Scaling,分层自适应率缩放)启发的逐层信任比相结合。该算法因被用于将 BERT 的预训练时间从单台 TPU pod 上的约三天缩短至 1024 个 TPUv3 芯片上的 76 分钟而声名鹊起,同时在 SQuAD 基准上达到或超过了已公布的 F1 分数。[1]
LAMB 处于深度学习中两个长期关注问题的交汇点:如何利用现代加速器的并行能力(这有利于大批量),以及如何保留人们认为小批量 SGD 所提供的隐式正则化。其核心洞察是:在全局学习率固定的情况下,权重范数差异很大的不同层会经历差异极大的有效更新幅度,而通过权重范数与更新范数之比的信任比对每层的更新进行重缩放,可以在大批量下恢复稳定的训练。
动机:大批量训练
分布式数据并行训练使有效批量大小与工作节点数量成正比扩展。原则上,将批量大小与学习率同时加倍可以保持梯度下降每轮的轨迹,这一启发式称为线性缩放规则。但在实际中,该规则在某个依赖问题的临界批量大小之外便会失效:训练要么发散,要么在更差的损失上停滞,要么泛化效果不佳。Goyal 等人证明,对 ImageNet 上的 ResNet-50,使用带动量、预热和精细归一化的 SGD 可以将批量大小线性扩展到 8192,[2] 但同样的方案对使用 Adam 训练的 Transformer 模型却行不通。
由 You 等人较早提出的 LARS 通过引入逐层信任比来处理 SGD 情形,使学习率自适应于每层的权重范数。[3] LARS 将 ResNet-50 的训练批量大小推进到 32K。然而 LARS 建立在带动量的 SGD之上,直接将其应用于使用自适应优化器的 Transformer 训练效果不佳。LAMB 将信任比的思想扩展到基于自适应矩的方法,而后者是 Transformer 预训练的事实标准选择。
算法表述
设 $ \theta_t \in \mathbb{R}^d $ 表示第 $ t $ 步的参数向量,按 $ L $ 层划分为 $ \theta_t = (\theta_t^{(1)}, \dots, \theta_t^{(L)}) $。给定小批量 $ \xi_t $ 上的随机梯度 $ g_t = \nabla_\theta \ell(\theta_t; \xi_t) $,LAMB 执行如下更新。
首先,它维护梯度及其平方的指数移动平均,这与 Adam 完全一致:
$ {\displaystyle m_t = \beta_1 m_{t-1} + (1 - \beta_1)\, g_t,} $
$ {\displaystyle v_t = \beta_2 v_{t-1} + (1 - \beta_2)\, g_t \odot g_t,} $
并通过偏差修正得到 $ \hat m_t = m_t / (1 - \beta_1^t) $ 与 $ \hat v_t = v_t / (1 - \beta_2^t) $。
其次,它构造一个按坐标的 Adam 风格更新,并加入解耦的权重衰减 $ \lambda $:
$ {\displaystyle r_t = \frac{\hat m_t}{\sqrt{\hat v_t} + \epsilon} + \lambda\, \theta_t.} $
第三,也是最关键的一步,它通过逐层信任比对每一层 $ i $ 的更新进行重缩放:
$ {\displaystyle \theta_{t+1}^{(i)} = \theta_t^{(i)} - \eta_t \cdot \frac{\phi\!\left(\| \theta_t^{(i)} \|\right)}{\| r_t^{(i)} \|}\, r_t^{(i)},} $
其中 $ \| \cdot \| $ 表示限制在第 $ i $ 层参数上的 L2 范数,$ \eta_t $ 为全局学习率,$ \phi: \mathbb{R}_{\ge 0} \to \mathbb{R}_{>0} $ 是一个缩放函数。在标准实现中取 $ \phi(x) = x $(恒等函数),并可选择性地裁剪到诸如 $ [\phi_{\min}, \phi_{\max}] $ 之类的范围,以防止权重范数过小或过大的层产生极端更新。
信任比 $ \| \theta^{(i)} \| / \| r^{(i)} \| $ 确保在期望意义上,相对变化 $ \| \Delta \theta^{(i)} \| / \| \theta^{(i)} \| $ 等于 $ \eta_t $,与该层的绝对尺度无关。Adam 更新相对于当前权重显得过大的层会被抑制,而过于保守的小更新则会被放大。
伪代码
input: learning rate eta, betas (b1, b2), epsilon, weight decay lambda
init: theta_0, m_0 = 0, v_0 = 0
for t = 1, 2, ... do
sample mini-batch, compute g_t
m_t = b1 * m_{t-1} + (1 - b1) * g_t
v_t = b2 * v_{t-1} + (1 - b2) * g_t * g_t
m_hat = m_t / (1 - b1**t)
v_hat = v_t / (1 - b2**t)
r_t = m_hat / (sqrt(v_hat) + epsilon) + lambda * theta_{t-1}
for each layer i do
w_norm = ||theta_{t-1}^(i)||
g_norm = ||r_t^(i)||
if w_norm > 0 and g_norm > 0:
trust = phi(w_norm) / g_norm
else:
trust = 1
theta_t^(i) = theta_{t-1}^(i) - eta * trust * r_t^(i)
end for
end for
实践考量
LAMB 在搭配包含预热阶段的学习率调度时最为有效。最初的 BERT 方案在前数千步采用线性预热,随后采用多项式衰减;峰值学习率明显高于典型的 Adam 范围,在 Transformer 预训练中常常达到 $ 10^{-2} $ 或更高,因为信任比已经吸收了更新的绝对幅度。
权重衰减以 AdamW 的方式与梯度解耦;若将衰减折入 $ g_t $,会与二阶矩归一化耦合,从而抹去其大部分收益。偏置和归一化参数(例如 batch norm 或 layer-norm 的缩放与平移系数)按惯例同时被排除在权重衰减和信任比重缩放之外,因为它们的范数很小,信任比在数值上可能变得不稳定。
数值稳定性还要求防范初始化或剪枝后出现零范数的层。参考实现会在权重范数或更新范数为零时退化为单位信任比。
比较
与 Adam 相比,LAMB 每步仅引入适度的常数倍开销,主要来自逐层范数计算。在小批量下,LAMB 相对 Adam 的收益可忽略不计,此时 Adam 已能很好收敛;一旦批量大小超过数千样本,差距才会显现。与 LARS 相比,LAMB 继承了 LARS 的逐层自适应,但使用自适应矩估计而非动量,这使其更适合 Transformer 类架构——在这类架构中,各层之间的梯度幅度可能相差多个数量级。
原始论文在标准的光滑性与有界方差假设下给出了收敛性分析,证明对非凸目标可在 $ O(1/\sqrt{T}) $ 速率下收敛到驻点。该分析还表明,LAMB 的更新在某种意义上是使逐层更新范数与权重范数成比例的唯一缩放方式,这为该算法提供了超越纯经验的理论依据。
局限性
LAMB 训练出的模型并不等同于小批量 Adam 所训练的模型。小批量与大批量之间的隐式正则化差异并未被完全消除,即使预训练损失相同,在非常大的批量下下游任务的准确率也可能略有降低。该算法还引入了额外的超参数:信任比的裁剪范围、按参数组的豁免设置以及 $ \phi $ 的选择。在实践中,源自 BERT 方案的默认值能很好地迁移到其他 Transformer 预训练任务,但在视觉或强化学习任务上需要重新调参。
从经验上看,LAMB 在 Transformer 预训练和大规模监督学习中最为成功。在微调阶段,由于批量通常较小且预训练初始化附近的优化曲面条件良好,它相对于精心调参的 Adam 的优势较小。对于非常小的模型或表格数据,通常使用更简单的优化器即可。