PaperlayerView exact revision on arXiv

Source-linked AI summary

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

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

arXiv:1904.00962v5cs.LGcs.AIcs.CLstat.ML

TL;DR

Training large neural networks remains computationally challenging, and existing large-batch methods such as LARS do not perform consistently across tasks. The paper develops LAMB from a layerwise adaptation strategy, analyzes LAMB and LARS in nonconvex settings, and reports strong results across BERT and ResNet-50. LAMB scales BERT pre-training to 64K batches without losing accuracy and reduces training time from 3 days to around 76 minutes.

  • Problem

    Training large deep neural networks is computationally challenging, while LARS’s performance gains are not consistent across tasks and it performs poorly for BERT.

  • Method

    The paper develops LAMB using layerwise update normalization and parameter-norm-based learning-rate scaling, and provides convergence analyses for LAMB and LARS.

  • Results

    LAMB achieves strong empirical performance across challenging tasks, including state-of-the-art accuracy for ResNet-50 and BERT training scaled to 64K batches without losing accuracy.

  • Takeaways & Limitations

    LAMB reduces BERT training time from 3 days to around 76 minutes while supporting adaptive optimization for both small and large batches.

  • Takeaways & Limitations

    The analysis assumes bounded gradients and other conditions typical of stochastic first-order-method analyses, while large-batch SGD learning-rate tuning remains difficult in practice.

Abstract

from arXiv · show

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 https://github.com/tensorflow/addons/blob/master/tensorflow_addons/optimizers/lamb.py

1 INTRODUCTION

Large-batch synchronous optimization is motivated by the cost of training modern neural networks, but existing scaling and layerwise-adaptation methods have task-dependent limitations. The paper develops and studies LAMB to improve large-batch training across tasks.

  • Training BERT takes 3 days on 16 TPUv3 chips, while ResNet-50 takes 29 hours on 8 Tesla P100 GPUs.
  • Synchronous large-batch SGD enables parallel gradient computation and reduces stochastic-gradient variance, but learning-rate scaling requires warmup and becomes harmful beyond certain batch sizes.
  • LARS trains ResNet-50 on ImageNet in a few minutes, yet performs poorly for attention models such as BERT and lacks substantial theoretical understanding.
  • The paper investigates a general layerwise adaptation strategy for large-batch learning and develops LAMB as a new adaptive SGD optimizer.
  • LAMB scales BERT batches beyond 32k without performance degradation, reducing training time from 3 days to 76 minutes.
  • The paper also targets adaptive optimization for image classification, where adaptive solvers such as Adam had not matched SGD with momentum on ResNet-50.

2 PRELIMINARIES

The paper formulates nonconvex stochastic optimization with layerwise smoothness and gradient-variance assumptions, then reviews SGD’s large-batch behavior. It identifies learning-rate tuning and dependence on maximum smoothness as practical limitations.

  • The optimization problem uses model parameters x, a smooth possibly nonconvex loss ℓ, and an unknown data distribution P.
  • The analysis assumes layerwise Lipschitz smoothness with constants L_i and bounds on both per-layer and per-coordinate stochastic-gradient variance.
  • The paper additionally assumes bounded gradients, a condition described as typical in stochastic first-order-method analyses.
  • SGD updates parameters using a scaled negative minibatch-gradient direction, with the minibatch drawn from the data distribution.
  • For large-batch SGD with b = T and an appropriate learning rate, the paper states a convergence result for a uniformly random iterate.
  • Learning-rate tuning is difficult in large-batch SGD, and dependence on L∞ can significantly slow convergence.

3 ALGORITHMS

The paper proposes layerwise normalization and parameter-norm-based learning-rate scaling as a general large-batch adaptation strategy, instantiating it as LARS and LAMB. It analyzes their nonconvex convergence and compares their rates with SGD.

  • General Strategy: The general strategy modifies a base optimizer by normalizing each layer’s update to unit l2-norm and scaling its learning rate by φ of the layer’s parameter norm.
  • General Strategy: Layerwise normalization can improve robustness to exploding gradients and plateaus, while φ makes update and parameter norms comparable.
  • LARS: LARS instantiates the strategy with momentum as its base optimizer, and its analysis considers a simplified nonconvex setting while stating extension to the general case.
  • LARS: The paper provides a convergence bound for LARS under large-batch and bounded-scaling-function conditions.
  • LAMB: LAMB uses ADAM as its base algorithm, combining per-dimension second-moment normalization with layerwise normalization.
  • LAMB: The paper provides a general nonconvex convergence bound for LAMB and states that the analysis extends beyond the simplified parameter setting.
  • Convergence Discussion: LARS and LAMB can have better convergence rates than SGD when gradients are sufficiently dense, with dependence on L_avg rather than L∞ under comparable criteria.

4 EXPERIMENTS

Experiments evaluate LAMB on BERT and ResNet-50 under large-batch settings, using minimal tuning and comparisons with established optimizers. LAMB preserves or reaches target accuracy while substantially reducing BERT training time and achieving competitive ImageNet accuracy.

  • Experimental Setup: LAMB experiments compare BERT and ResNet-50 training against existing optimizers using very minimal hyperparameter tuning.The experiments tune only the learning rate while fixing β1 = 0.9 and β2 = 0.999.
  • BERT Training: 91.460 F1 was achieved for BERT with batch size 32768, reducing training time from 3 days to around 100 minutes.This configuration used 15625 iterations and achieved 49.1 times speedup with 76.7% efficiency.
  • BERT Training: 76 minutes was sufficient for BERT training with mixed-batch LAMB and re-warm-up, requiring only 8599 iterations and achieving 100.2% efficiency.The procedure re-warms the learning rate during the second stage before decaying it again.
  • BERT Training: LAMB outperformed LARS at every reported BERT batch size, while ADAMW stopped scaling at batch size 16K and failed to reach the target F1 score.For 64K/32K mixed-batch training, extensive ADAMW tuning still produced no reasonable result.
  • ImageNet Training with ResNet-50: 76.4% top-1 accuracy was achieved by LAMB at batch size 32K on ImageNet/ResNet-50, compared with 76.3% for LARS.At batch size 2K, LAMB achieved 77.11% top-1 accuracy while LARS achieved 76.6%.
  • Hyperparameters for Scaling the Batch Size: LAMB uses square root learning-rate scaling and linear-epoch warmup for untuned batch-size scaling in BERT and ResNet-50.For BERT at batch size 32K, 3125 of 15625 iterations are used for learning-rate warmup.

5 CONCLUSION

The paper proposes LAMB for adaptive elementwise updating and layerwise learning rates, with applicability to small and large batches. It reports improved performance across applications, including scaling BERT pre-training to 64K without accuracy loss and achieving state-of-the-art ImageNet accuracy with ResNet-50.

  • LAMB supports adaptive elementwise updating and layerwise learning rates for both small and large batches.
  • LAMB achieves better performance than existing optimizers across a wide range of applications.
  • 64K batch size scales BERT pre-training without losing accuracy and reduces training time from 3 days to around 76 minutes.
  • LAMB is the first large batch adaptive solver reported to achieve state-of-the-art accuracy on ImageNet with ResNet-50.

A PROOF OF THEOREM 2

The proof analyzes LARS updates under general minibatch sizes using smoothness, Cauchy-Schwarz, expectation, and telescoping-sum arguments. Rearrangement yields a convergence inequality involving the initial and terminal objective values and optimization parameters.

  • The proof starts from the LARS update and applies the L-smoothness property of the objective.
  • The analysis rewrites the smoothness inequality using a gradient-error term and then applies Cauchy-Schwarz.
  • Taking expectation produces a bound that is summed from t = 1 to T using a telescoping sum.
  • The final rearrangement divides by ηTα_l and expresses the resulting convergence bound in terms of objective decrease and optimization parameters.

B PROOF OF THEOREM 3

The proof analyzes LAMB convergence for general minibatch sizes by combining its update procedure with smoothness-based bounds. It treats moment parameters and gradient-related terms before summing and rearranging the inequalities.

  • The proof begins from the LAMB update and specifies moment parameters, scaling function φ, and initialization for m_0 and v_0.
  • The analysis uses L-smoothness and the Lipschitz continuity of the gradient to establish the main inequality.
  • The proof bounds term T1 separately for β2 = 0 and β2 > 0, then takes expectation and substitutes the resulting bound.
  • The update computes a bias-corrected second moment and a ratio r_t from the moment estimates.
  • Summing from t = 1 to T with a telescoping sum and rearranging yields the LAMB convergence bound.

C COMPARISON OF CONVERGENCE RATES OF LARS AND SGD

The section compares LARS and SGD convergence behavior and reports experiments on LAMB variants and design choices. LARS can improve over SGD under a stated gradient-versus-curvature or stochasticity condition, while several LAMB modifications preserve accuracy.

  • COMPARISON OF CONVERGENCE RATES OF LARS AND SGD: LARS gains over SGD when the gradient is denser than curvature or stochasticity; otherwise, SGD has the better convergence-rate upper bound.
  • N-LAMB: NESTEROV MOMENTUM FOR LAMB: N-LAMB replaces LAMB’s first-moment regular momentum with Nesterov momentum, while NN-LAMB applies Nesterov momentum to both moments.
  • N-LAMB: NESTEROV MOMENTUM FOR LAMB: N-LAMB and NN-LAMB achieve accuracy comparable to LAMB and perform much better than the momentum solver in ResNet-50 training with batch size 32K.
  • LAMB WITH LEARNING RATE CORRECTION: Removing adam-correction from LAMB produces no observed drop in test or validation accuracy for BERT and ImageNet training because it has the same effect as learning-rate warmup.
  • LAMB WITH LEARNING RATE CORRECTION: Different norms produce less than 0.1 percent difference in ResNet-50 validation accuracy, so L2 norm is used by default.

G REGULAR BATCH SIZES FOR SMALL DATASETS: MNIST AND CIFAR-10.

LAMB performs competitively on smaller datasets, achieving strong accuracy on CIFAR-10 and MNIST under carefully tuned comparisons.

  • 94.08% test accuracy: LAMB outperforms other adaptive optimizers and momentum SGD on CIFAR-10 with DavidNet.All solvers ran 24 epochs and finished in under one minute on one cloud TPU.
  • The CIFAR-10 comparison carefully tunes the adaptive optimizers' learning rates over a broad search space.The reported tuning space includes values from 0.0001 through 50.

H IMPLEMENTATION DETAILS AND ADDITIONAL RESULTS

The paper describes LAMB's default hyperparameters, emphasizes learning-rate tuning, and evaluates optimizers using accuracy or F1 rather than validation loss alone.

  • LAMB uses default λ=0.01, β1=0.9, β2=0.999, and ϵ=1e-6, without tuning β1 or β2.The paper notes that tuning β1 and β2 may produce higher accuracy.
  • Learning rate is identified as the most important hyperparameter affecting learning efficiency and final accuracy.The paper therefore carefully tunes learning rates for Adam, AdamW, AdaGrad, and momentum SGD.
  • Large-batch validation loss is unreliable, so the experiments use test or validation accuracy and dev-set F1 scores instead.A lower validation loss does not necessarily correspond to higher validation accuracy.

H.0.1 BERT

For BERT, LAMB maintains performance and smooth convergence at very large batch sizes while delivering substantial scaling efficiency.

  • 91.345 F1 score: LAMB exceeds the 90.5 target while AdamW stops scaling at batch size 16K.The table reports the best F1 score observed in the experiments.
  • LAMB produces nearly identical BERT loss curves across different batch sizes, indicating that it scales well with batch size.The comparison is described in Figure 6.
  • 64K batch size: LAMB makes mixed-batch BERT pre-training converge smoothly.The result is shown in Figure 7.
  • 76.8% scaling efficiency: ordinary batch scaling yields 49.1 times speedup with 64 times the computational resources.Mixed-batch training reaches 101.8% scaling efficiency, with 65.2 times speedup using 64 times the resources.

H.1 BASELINE TUNING DETAILS FOR IMAGENET TRAINING WITH RESNET-50

The ImageNet experiments emphasize extensive baseline tuning and conclude that existing adaptive solvers are difficult to tune or perform poorly, while LAMB's training remains smooth.

  • Baseline tuning: Five-epoch warmup and learning-rate drops at epochs 30, 60, and 80 are included when tuning adaptive ImageNet baselines.The target accuracy is around 76.3%.
  • Additional results: LAMB's ImageNet training loss curves remain smooth even at extremely large batch sizes.The figures are used to show convergence behavior rather than validation performance.
  • Baseline tuning: Existing adaptive solvers do not perform well on ImageNet training, or are at least hard to tune, according to the comprehensive tuning results.The baselines include AdaGrad, Adam, and AdamW variants with and without the learning-rate scheme from Goyal et al.
Loading 1904.00962v5…