LAMB优化器（逐层自适应矩批量训练）是一种用于训练深度神经网络的随机优化算法，它通过逐层归一化步骤扩展了Adam优化器。该算法于2019年由谷歌的研究人员（具体为杨宇、李静、Jonathan Hseu等人）提出，旨在实现高效使用非常大的小批量大小（例如32,768或更大），同时不降低模型精度或需要详尽的超参数调优。它通过根据每层权重范数与其更新范数的比率来缩放每层的更新幅度，从而有效地将学习率与每层参数的尺度解耦。
LAMB在训练基于Transformer (architecture)的模型方面尤其具有影响力，包括早期的大型语言模型（LLM）以及ResNet等视觉架构。它在分布式训练框架中的采用，例如TensorFlow（通过tf.keras.optimizers.LAMB）和PyTorch（通过NVIDIA、Hugging Face和FairScale等库中的LAMB实现），使其成为扩展训练运行规模的标准工具。通过允许更大的批量，LAMB减少了训练最先进模型所需的挂钟时间，这对于依赖大规模计算集群的OpenAI、Anthropic和Google DeepMind等组织至关重要。
动机与背景
使用大批量训练深度神经网络减少了每个训练周期内的参数更新次数，但简单地扩大批量大小通常会导致泛化性能差和收敛不稳定。这被称为“大批量训练问题”。标准的优化器，如带动量的随机梯度下降（SGD）或Adam，在批量大小增加时需要仔细调整学习率，即便如此，精度也常常下降。LAMB优化器的开发正是为了解决这个问题，使优化器对批量大小的变化更加稳健。
LAMB背后的关键洞察是，深度网络中的不同层表现出截然不同的梯度和权重范数尺度。例如，早期的卷积层权重较小，而较晚的全连接层权重较大。Adam中的单一全局学习率可能导致某些层的更新过大（导致发散）或另一些层的更新过小（减慢收敛）。LAMB引入了一种逐层自适应速率，根据该层权重和梯度范数的比率对更新进行归一化，确保每层相对于其幅度以稳定的速度移动。
算法描述
LAMB可以看作是Adam的一种变体，增加了一个额外的归一化步骤。设\(\theta_t\)为第\(t\)次迭代时的参数，\(g_t\)为损失相对于\(\theta_t\)的梯度。LAMB维护梯度的第一矩和第二矩（\(m_t\)和\(v_t\)），类似于Adam，具有指数衰减率\(\beta_1\)和\(\beta_2\)（通常为0.9和0.999）。在偏差校正之后，它计算Adam更新\(\frac{m_t}{\sqrt{v_t} + \epsilon}\)。
关键的区别在于信任比率 \(\phi\)：对于每一层\(i\)，\(\phi_i = \frac{||\theta_{t,i}||}{||r_{t,i}||}\)，其中\(r_{t,i} = \frac{m_{t,i}}{\sqrt{v_{t,i}} + \epsilon}\)是该层的Adam更新（不含学习率），\(||\cdot||\)表示L2范数。第\(i\)层的最终更新为\(\theta_{t+1,i} = \theta_{t,i} - \eta \cdot \phi_i \cdot r_{t,i}\)，其中\(\eta\)是全局学习率。这个信任比率按比例缩放更新到该层的权重范数，因此小层获得较小（但不可忽略）的更新，大层获得更大但稳定的更新。
在实践中，会在分母中添加一个小常数（例如1e-6）以避免除以零。该算法还包括可选的权重衰减（L2正则化），融合到更新中，遵循AdamW中使用的解耦权重衰减方法。当所有层的信任比率设置为1时，LAMB退化为标准Adam（带偏差校正）。
超参数与调优
LAMB继承了Adam的大部分超参数：\(\beta_1\)（动量）、\(\beta_2\)（方差衰减）、\(\epsilon\)（数值稳定性）和权重衰减率。主要的新超参数是全局学习率 \(\eta\)，对于大批量训练，通常设置在0.01-0.1范围内，显著高于Adam的典型值（例如1e-3）。作者发现，对于非常大的批量（例如BERT的32,768），在前10%的步骤中使用0.01的学习率并进行线性预热效果良好，他们还建议在剩余步骤中使用余弦衰减学习率调度。
此外，\(\beta_2\)的选择会影响稳定性；对于梯度稀疏的模型，可以使用更高的\(\beta_2\)（例如0.99）。作者还建议，批量大小可以随学习率按比例缩放（线性缩放规则），这一准则与LAMB配合良好。例如，如果批量大小加倍，学习率也可以加倍而不会损失精度。
性能与基准测试
在原始论文中，LAMB在两个主要任务上进行了评估：在ImageNet上训练ResNet-50（图像分类）和BERT（基于transformer的语言模型）进行掩码语言建模。使用LAMB，作者在批量大小为32,768的情况下，仅用2,048次迭代就达到了ImageNet top-1准确率76.0%，与使用较小批量（例如256）在更少周期内达到的最先进准确率相匹配。对于BERT，他们使用1,024个TPU在约3.5分钟内将模型训练到相同的准确率（例如在SQUAD上F1分数为1.0），比先前方法提速10倍。
随后，LAMB成为谷歌内部工作流中训练基于BERT模型默认优化器。论文报告称，当批量大小从1,024扩展到65,536时，LAMB的性能优于Adam和带动量的SGD。作者还表明，LAMB与Gradient Clipping（用于防止梯度爆炸）配合良好，并且与混合精度训练兼容，正如在NVIDIA GPU和谷歌TPU等现代硬件上使用的那样。
在大规模训练中的应用
LAMB的主要应用场景是分布式训练，其中批量大小太大，无法容纳在单个设备的内存中。在这种设置中，梯度通过数据并行在多个GPU或TPU之间进行平均。例如，OpenAI和Google DeepMind在训练序列长度达数千的大型Transformer (architecture)模型时，使用类似于LAMB的优化器。尽管已经提出了更新的优化器，如LAMB（及其后继者LAMB2），但LAMB在许多开源工作中仍然是可靠的选择，包括AI21 Labs和SambaNova等研究团体和公司训练视觉transformer和LLM。
在Amazon Web Services（使用AWS Trainium硬件）上的Machine learning环境中，LAMB在自定义内核中得到支持以提高效率。类似地，Intel和AMD在其加速器上对LAMB进行了基准测试。该优化器处理极端批量大小的能力使其对于在大型数据集上预训练模型非常有价值，因为单个周期的成本很高，减少周期数至关重要。
与其他优化器的关系
LAMB是自适应优化器家族的一部分，该家族包括SGD变体、Adam及其后继者，如AdamW（解耦权重衰减）和LARS（逐层自适应速率缩放）。LARS由You等人于2017年为CNN的大批量训练提出，使用类似的逐层信任比率，但不维护第二矩；它依赖于第一矩（动量）和梯度范数。LAMB结合了LARS的优点（逐层缩放）和Adam的自适应逐参数学习率，使其对于梯度稀疏的模型（如transformer）更加稳健。
另一个密切相关的优化器是NVLAMB（来自Nvidia），它结合了方差缩减。然而，LAMB仍然更简单且被广泛使用。对于使用Beam Search的Sequence-to-Sequence (Seq2Seq)任务，LAMB不直接影响推理，但它有助于训练收敛，从而间接改善序列解码。
扩展与变体
自推出以来，已经提出了几种变体。LAMB2（也来自谷歌）添加了基于梯度方差的归一化因子，提高了某些问题的稳定性。带偏差校正的LARS和其他修改也很常见。在实践中，许多框架实现了带有可选偏差校正的LAMB，这在最初的几步中是有益的。一些实现，例如PyTorch的torch.optim.Lamb（在torch_optimizer包中），允许调整信任比率参数或使用自定义的逐层学习率。
尽管出现了新的优化器，如具有不同缩放策略（例如1cycle调度）的AdamW，但LAMB仍然是大批量训练的强有力基线。研究界探索了将LAMB与Data Augmentation和Gradient Clipping相结合，以进一步提高泛化性能。
实际考虑与局限性
虽然LAMB在大批量设置中表现出色，但它并不总是小批量大小（例如低于1,024）的最佳选择。在这种场景下，标准Adam或带动量的SGD可能更简单且同样有效。LAMB增加了计算每层范数的计算开销，这在现代硬件上可以忽略不计，但对于具有许多小层的模型（例如U-Net架构）可能不可忽略。
另一个局限性是，如果某些层的权重范数非常小（接近零），LAMB的信任比率偶尔会导致训练不稳定。这通常通过向分母添加epsilon项以及使用权重衰减来缓解，权重衰减可以防止权重漂移到零。此外，LAMB需要仔细调整学习率和预热步骤；不合适的调度可能导致发散。
内存使用与Adam类似（每个参数两个矩向量），因此不会更消耗内存。对于非常大的模型，可以结合使用Model Pruning或Gradient Clipping与LAMB，但这些是正交技术。
结论
LAMB已成为大规模深度学习优化算法工具箱中的基石。通过实现使用大规模小批量的有效训练，它加速了许多基准模型的发展，并降低了实验成本。其逐层自适应原理影响了后续优化器的设计，并且对于面临分布式训练挑战的从业者来说，它仍然是一个实用且广为人知的解决方案。随着Artificial intelligence的持续发展，像LAMB这样的优化器可能会不断演变，但其逐层信任和自适应矩的核心思想将持续存在。