The LAMB optimizer (Layer-wise Adaptive Moments for Batch training) is a stochastic optimization algorithm for training deep neural networks that extends the Adam optimizer with a per-layer normalization step. Introduced in 2019 by researchers at Google (specifically by Yang You, Jing Li, Jonathan Hseu, et al.), LAMB was designed to enable the efficient use of very large mini-batch sizes (e.g., 32,768 or more) without degrading model accuracy or requiring exhaustive hyperparameter tuning. It achieves this by scaling the update magnitude for each layer based on the ratio of the layer's weight norm to its update norm, effectively decoupling the learning rate from the scale of each layer's parameters.
LAMB has been particularly influential in training Transformer (architecture)-based models, including early large language models (LLMs) and vision architectures like ResNet. Its adoption in distributed training frameworks, such as TensorFlow (via tf.keras.optimizers.LAMB) and PyTorch (via the LAMB implementation in libraries like NVIDIA, Hugging Face, and FairScale), made it a standard tool for scaling up training runs. By allowing larger batches, LAMB reduces the wall-clock time required to train state-of-the-art models, which is critical for organizations like OpenAI, Anthropic, and Google DeepMind that rely on massive compute clusters.
Motivation and Background
Training deep neural networks with large mini-batches reduces the number of parameter updates per epochadian computational steps, but naive scaling of the batch size often leads to poor generalization and unstable convergence. This is known as the "large-batch training problem." Standard optimizers like stochastic gradient descent (SGD) with momentum or Adam require careful adjustment of the learning rate when batch size increases, and even then, accuracy often degrades. The LAMB optimizer was developed to address this by making the optimizer more robust to batch size changes.
The key insight behind LAMB is that different layers in a deep network exhibit vastly different scales of gradients and weight norms. For example, early convolutional layers have small weights, while later fully connected layers have large ones. A single global learning rate in Adam can lead to either too large updates for some layers (causing divergence) or too small updates for others (slowing convergence). LAMB introduces a layer-wise adaptive rate that normalizes the update based on the ratio of the layer's weight and gradient norms, ensuring each layer moves at a stable pace relative to its magnitude.
Algorithm Description
LAMB can be seen as a variant of Adam with an additional normalization step. Let \(\theta_t\) be the parameters at iteration \(t\), and \(g_t\) the gradient of the loss with respect to \(\theta_t\). LAMB maintains first and second moments (\(m_t\) and \(v_t\)) of the gradients, similar to Adam, with exponential decay rates \(\beta_1\) and \(\beta_2\) (typically 0.9 and 0.999). After bias correction, it computes an Adam update \(\frac{m_t}{\sqrt{v_t} + \epsilon}\).
The crucial difference is the trust ratio \(\phi\): for each layer \(i\), \(\phi_i = \frac{||\theta_{t,i}||}{||r_{t,i}||}\), where \(r_{t,i} = \frac{m_{t,i}}{\sqrt{v_{t,i}} + \epsilon}\) is the Adam update (without learning rate) for that layer, and \(||\cdot||\) denotes the L2 norm. The final update for layer \(i\) is \(\theta_{t+1,i} = \theta_{t,i} - \eta \cdot \phi_i \cdot r_{t,i}\), where \(\eta\) is the global learning rate. This trust ratio scales the update proportionally to the layer's weight norm, so small layers get small (but not negligible) updateships, and large layers get larger, yet stable, updates.
In practice, a small constant (e.g., 1e-6) is added to the denominator to avoid division by zero. The algorithm also includes optional weight decay (L2 regularization) fused into the update, following the decoupled weight decay approach used in AdamW. When the trust ratio is set to 1 for all layers, LAMB reduces to the standard Adam (with bias correction).
Hyperparameters and Tuning
LAMB inherits most hyperparameters from Adam: \(\beta_1\) (momentum), \(\beta_2\) (variance decay), \(\epsilon\) (numerical stability), and weight decay rate. The primary new hyperparameter is the global learning rate \(\eta\), which is often set in the range 0.01-0.1 for large-batch training, significantly higher than what is typical for Adam (e.g., 1e-3). The authors found that for very large batches (e.g., 32,768 for BERT), a learning rate of 0.01 with linear warmup over the first 10% of steps works well, and they also recommended using a cosine decay learning rate schedule over the remaining steps.
Additionally, the choice of \(\beta_2\) can affect stability; for models with sparse gradients, a higher \(\beta_2\) (e.g., 0.99) may be used. The authors also suggested that batch size can be scaled proportionally with the learning rate (linear scaling rule), a guideline that works well with LAMB. For instance, if batch size is doubled, the learning rate can also be doubled without loss of accuracy.
Performance and Benchmarks
In the original paper, LAMB was evaluated on two main tasks: training ResNet-50 on ImageNet (image classification) and BERT (a transformer-based language model) for masked language modeling. Using LAMB, the authors achieved ImageNet top-1 accuracy of 76.0% in only 2,048 iterations with a batch size of 32,768, matching the state-of-the-art accuracy achieved with smaller batches (e.g., 256) in far fewer epochs. For BERT, they trained the model to the same accuracy (e.g., F1-score of 1.0 on SQUAD) in about 3.5 minutes using 1,024 TPUs, a 10x speedup over prior methods.
Subsequently, LAMB became the default optimizer for training BERT-based models in Google's internal workflows. The paper reported that LAMB outperformed both Adam and SGD with momentum when scaling batch sizes from 1,024 to 65,536. The authors also showed that LAMB works well with Gradient Clipping (used to prevent exploding gradients) and is compatible with mixed-precision training, as used on modern hardware like NVIDIA GPUs and Google TPUs.
Applications in Large-Scale Training
LAMB's primary application is in distributed training where the batch size is too large to fit in a single device's memory. In such setups, gradients are averaged across multiple GPUs or TPUs using data parallelism. For example, OpenAI and Google DeepMind use optimizers analogous to LAMB when training large Transformer (architecture) models with sequence lengths in the thousands. Although newer optimizers like LAMB (and its successor LAMB2) have been proposed, LAMB remains a reliable choice in many open-source efforts, including the training of vision transformers and LLMs by research groups and companies like AI21 Labs and SambaNova.
In the context of Machine learning on Amazon Web Services (with AWS Trainium hardware), LAMB is supported in custom kernels for efficiency. Similarly, Intel and AMD have benchmarked LAMB on their accelerators. The optimizer's ability to handle extreme batch sizes makes it valuable for pre-training models on massive datasets, where the cost of a single epoch is high, and reducing epochs is paramount.
Relationship to Other Optimizers
LAMB is part of a family of adaptive optimizers that include SGD variants, Adam, and its successors like AdamW (decoupled weight decay) and LARS (Layer-wise Adaptive Rate Scaling). LARS, introduced by You et al. in 2017 for large-batch training of CNNs, uses a similar layer-wise trust ratio but does not maintain second moments; it relies on first moments (momentum) and gradient norms. LAMB combines the benefits of LARS (layer-wise scaling) with Adam's adaptive per-parameter learning rates, making it more robust for models with sparse gradients (like transformers).
Another closely related optimizer is NVLAMB (from Nvidia), which incorporates variance reduction. However, LAMB remains simpler and widely used. For Sequence-to-Sequence (Seq2Seq) tasks with Beam Search, LAMB does not directly affect inference, but it helps training convergence, which indirectly improves sequence decoding.
Extensions and Variants
Since its introduction, several variants have been proposed. LAMB2 (also from Google) adds a normalization factor based on gradient variance, improving stability for certain problems. LARS with bias correction and other modifications are also common. In practice, many frameworks implement LAMB with optional bias correction for the moments, which is beneficial during the first few steps. Some implementations, such as in PyTorch's torch.optim.Lamb (in the torch_optimizer package), allow for adjusting the trust ratio parameter or using a custom per-layer learning rate.
Despite the emergence of new optimizers like AdamW with different scaling strategies (e.g., 1cycle schedules), LAMB remains a strong baseline for large-batch training. The research community has explored combining LAMB with Data Augmentation and Gradient Clipping to further improve generalization.
Practical Considerations and Limitations
While LAMB excels in large-batch settings, it is not always the best choice for small batch sizes (e.g., below 1,024). In such regimes, standard Adam or SGD with momentum may be simpler and equally effective. LAMB adds a computational overhead of computing norms per layer, which is negligible on modern hardware but can be non-trivial for models with many small layers (e.g., U-Net architectures).
Another limitation is that LAMB's trust ratio can occasionally lead to unstable training if some layers have very small weight norms (near zero). This is typically mitigated by adding an epsilon term to the denominator and by using weight decay, which prevents weights from drifting to zero. Furthermore, LAMB requires careful tuning of the learning rate and warmup steps; an inappropriate schedule can lead to divergence.
Memory usage is similar to Adam (two moment vectors per parameter), so it is not more memory-hungry. For very large models, Model Pruning or Gradient Clipping may be used alongside LAMB, but these are orthogonal techniques.
Conclusion
LAMB has become a cornerstone in the toolbox of optimization algorithms for large-scale deep learning. By enabling effective training with massive mini-batches, it has accelerated the development of many benchmark modelsause and reduced the cost of experimentation. Its layer-wise adaptation principle has influenced subsequent optimizer designs and remains a practical, well-understood solution for practitioners facing the challenges of distributed training. As Artificial intelligence continues to grow, optimizers like LAMB will likely evolve, but its core ideas of per-layer trust and adaptive moments are here to stay.