LAMB（Layer-wise Adaptive Moments for Batch training，逐层自适应矩批量训练）是一种用于训练深度神经网络的优化算法，尤其适用于大批量分布式计算环境。该算法于2019年由谷歌和多伦多大学的研究人员提出，旨在解决在数千个加速器上扩展训练时，同时保持模型准确性和收敛速度的挑战。
该算法扩展了Adam优化器，通过为每一层而非每个参数计算自适应学习率。这种逐层自适应机制使LAMB能够更有效地处理不同网络层之间梯度尺度的差异，这在Transformer和残差网络等深层架构中尤为重要。通过基于层的权重和梯度范数对更新进行归一化，LAMB即使在使用非常大的批量大小时也能确保稳定高效的训练。
背景与动机
训练大规模神经网络通常需要大量的计算资源，这些资源往往分布在许多GPU或TPU上。增加批量大小是高效利用这些资源的常见策略，但这通常会导致模型性能下降或收敛速度变慢。传统的优化器如随机梯度下降（SGD）和Adam在处理大批量时表现不佳，因为它们依赖全局学习率，无法考虑各层梯度尺度的异质性。
LAMB优化器的开发就是为了克服这些限制。其设计灵感来源于此前用于卷积网络大批量训练的逐层自适应速率缩放（LARS）算法。LAMB将该概念推广到自适应矩估计中，结合了两种方法的优点。
算法细节
LAMB根据每层的权重范数与其梯度范数的比值来计算该层的更新。参数张量在步骤t的核心更新规则如下：
- 如Adam一样，计算第一和第二矩估计（梯度的均值和方差）。
- 将更新方向计算为矩校正后的梯度除以第二矩的平方根加上一个小epsilon。
- 将该方向按层的权重范数与更新方向范数的比值进行缩放。
- 乘以全局学习率并应用更新。
这种逐层缩放确保了权重范数较大的层获得相应较大的更新，而范数较小的层则保守更新。该算法还包含一个信任比率，可以像梯度裁剪技术一样进行裁剪，以防止极端更新。
作者证明，LAMB可以在批量大小为32,768的情况下训练ResNet-50在ImageNet上，达到与批量大小为256的基线相同的准确率，但所需步骤显著减少。这使得它非常适合在数百或数千个加速器上进行分布式训练。
应用与影响
LAMB已被广泛用于训练大型语言模型和其他深度学习模型。例如，它被用于大规模训练BERT和其他基于Transformer的模型，将训练时间从数天缩短到数小时。该算法在硬件资源丰富的环境中尤其有价值，例如AWS和Microsoft Azure云平台，以及AWS Trainium和Graphcore IPU等专用AI硬件。
许多后续优化器，如LAMB的变体和后继者，都基于其原理构建。它还影响了学习率调度和自适应优化方法的研究。在TensorFlow和PyTorch等框架中的开源实现使其对更广泛的机器学习社区可用。
与其他优化器的比较
与Adam相比，LAMB在使用大批量时通常能实现更快的收敛和更好的最终性能。Adam的全局学习率通常需要仔细调整，并且在大批量下可能导致不稳定。LAMB的逐层自适应缓解了这些问题，允许更激进的扩展。
与为带动量的SGD设计的LARS相比，LAMB结合了自适应矩估计，使其对噪声梯度和稀疏特征更加鲁棒。这使得LAMB成为适用于各种架构的更通用选择，包括序列到序列模型和编码器-解码器框架。
局限性与考虑
尽管有其优势，LAMB并非没有局限性。该算法引入了额外的超参数，如信任比率裁剪阈值和epsilon项，这些可能需要针对特定任务进行调整。它还假设逐层缩放是有益的，但对于层间高度相关或使用某些权重初始化方案的架构，这一假设可能并不总是成立。
此外，虽然LAMB在大批量设置中表现出色，但在小批量情况下其优势会减弱，此时像Adam这样更简单的优化器可能就足够了。研究人员还指出，该算法的性能可能对全局学习率的选择敏感，并且可能需要学习率预热才能达到最佳效果。
结论
LAMB代表了大规模深度学习训练优化领域的重大进步。通过结合逐层自适应和自适应矩，它实现了大批量下的高效稳定训练，使其成为生成式AI和人工智能研究时代的基础技术。其影响超越了最初的应用，塑造了后续优化器和训练方法的发展。