Large Batch Optimization for Deep Learning: Training BERT in 76 minutes

By Yang You, Jing Li, Sashank J. Reddi, Jonathan Hseu, Sanjiv Kumar, Srinadh Bhojanapalli, Xiaodan Song, J. Demmel, Kurt Keutzer, Cho-Jui Hsieh. 2019

1,222 citations recorded by the source corpus.

Training large deep neural networks on massive datasets is computationally very challenging. There has been recent surge in interest in using large batch stochastic optimization methods to tackle this issue. The most prominent algorithm in this line of research is LARS, which by employing layerwise adaptive learning rates trains ResNet on ImageNet in a few minutes. However, LARS performs poorly for attention models like BERT, indicating that its performance gains are not consistent across tasks. In this paper, we first study a principled layerwise adaptation strategy to accelerate training of deep neural networks using large mini-batches. Using this strategy, we develop a new layerwise adaptive large batch optimization technique called LAMB; we then provide convergence analysis of LAMB as well as LARS, showing convergence to a stationary point in general nonconvex settings. Our empirical results demonstrate the superior performance of LAMB across various tasks such as BERT and ResNet-50 training with very little hyperparameter tuning. In particular, for BERT training, our optimizer enables use of very large batch sizes of 32868 without any degradation of performance. By increasing the batch size to the memory limit of a TPUv3 Pod, BERT training time can be reduced from 3 days to just 76 minutes (Table 1). The LAMB implementation is available at this https URL

Subfields: Artificial Intelligence, Computer Vision and Pattern Recognition, Large Language Models, Optimization & Training

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Computer Vision and Pattern Recognition

In Large Language Models

In Optimization & Training

Citations in the corpus

An Empirical Study of Training Self-Supervised Vision Transformers (2021, methodology): “3, we study the LAMB optimizer [47], which is an AdamW-counterpart of LARS.”
Perceiver IO: A General Architecture for Structured Inputs & Outputs (2021, positive): “Training steps 500,000 Batch size 2048 Masking strategy Words Optimizer LAMB (You et al., 2021) Learning rate 0.”
Perceiver: General Perception with Iterative Attention (2021, methodology): “Although it is typical to train convolutional networks on ImageNet using SGD, we found it easier to optimize Perceiver models using the LAMB optimizer (You et al., 2020), which was developed for optimizing Transformer-based models.”
Initialization and Regularization of Factorized Neural Layers (2021, methodology): “Table 4: Comparison of BERT unsupervised pretraining using LAMB (You et al., 2020) and our Frobenius decay scheme FLAMBé.”
Enabling Multimodal Generation on CLIP via Vision-Language Knowledge Distillation (2022, methodology): “Within each task, we use the same setting796 for multiple datasets.797 For the GLUE benchmark, we use the LAMB op-798 timizer (You et al., 2020) to train for 10 epochs.”
FILIP: Fine-grained Interactive Language-Image Pre-Training (2021, methodology): “We train both models using the LAMB optimizer (You et al., 2020) and cosine learning rate schedule (Loshchilov & Hutter, 2016) with a linear warmup.”
Scaling Up Visual and Vision-Language Representation Learning With Noisy Text Supervision (2021, methodology): “We use LAMB optimizer (You et al., 2020)1 with weight decay ratio 1e-5.”
High-Performance Large-Scale Image Recognition Without Normalization (2021, background): “AGC is closely related to a recent line of work studying “normalized optimizers” (You et al., 2017; Bernstein et al., 2020; You et al., 2019), which ignore the scale of the gradient by choosing an adaptive learning rate inversely proportional to the gradient norm.”
Decoupling the Role of Data, Attention, and Losses in Multimodal Transformers (2021, methodology): “We use the LAMB optimizer (You et al., 2019) with an initial learning rate of 0.00176 and 20,000 warm up steps.”
A Primer in BERTology: What We Know About How BERT Works (2020, background): “You et al. (2019) report that with a batch size of 32k BERT’s training time can be sig-niﬁcantly reduced with no degradation in performance.”
Mime: Mimicking Centralized Stochastic Algorithms in Federated Learning (2020, background): “[39] Yang You, Jing Li, Sashank Reddi, Jonathan Hseu, Sanjiv Kumar, Srinadh Bhojanapalli, Xiaodan Song, James Demmel, Kurt Keutzer, and Cho-Jui Hsieh.”
Evolutionary-scale prediction of atomic level protein structure with a language model (2022, background): “[67] Yang You, Jing Li, Sashank Reddi, Jonathan Hseu, Sanjiv Kumar, Srinadh Bhojanapalli, Xiaodan Song, James Demmel, Kurt Keutzer, and Cho-Jui Hsieh.”
Class-Incremental Learning with Cross-Space Clustering and Controlled Transfer (2022, neutral): “Prior work has observed that increasing the batch size tends to slightly reduce the test accuracy, and have proposed various methods to reduce the drop in accuracy [4, 2, 5].”