PaperlayerView exact revision on arXiv

Source-linked AI summary

A Large Batch Optimizer Reality Check: Traditional, Generic Optimizers Suffice Across Batch Sizes

Zachary Nado, Justin M. Gilmer, Christopher J. Shallue, Rohan Anil, George E. Dahl

arXiv:2102.06356v3cs.LGstat.ML

TL;DR

Whether LARS and LAMB outperform standard optimizers at large batch sizes remains an open question without fair comparisons. The paper compares standard methods with those optimizers and finds that Nesterov momentum and Adam can match or improve on their results. It establishes stronger baselines for the examined tasks and batch sizes.

  • Problem

    Fair comparisons have not established whether LARS and LAMB outperform standard optimizers at large batch sizes.

    Source: §1 · p. 2

  • Method

    The paper compares standard Nesterov momentum and Adam with LARS and LAMB on large-batch ResNet-50 and BERT training.

    Source: §1 · p. 2

  • Results

    Standard optimizers match or improve on LARS and LAMB's large-batch results across the examined tasks, with Adam also improving BERT downstream metrics.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Standard optimizers provide stronger baselines for future comparisons at the examined tasks and batch sizes.

    Source: §5 · p. 8, §6 · p. 9

    At relatively small step budgets, learning-rate schedules are optimizer-dependent and sensitive to small deviations from their optimized forms.

    Source: §2.2 · p. 6

Abstract

from arXiv · show

Recently the LARS and LAMB optimizers have been proposed for training neural networks faster using large batch sizes. LARS and LAMB add layer-wise normalization to the update rules of Heavy-ball momentum and Adam, respectively, and have become popular in prominent benchmarks and deep learning libraries. However, without fair comparisons to standard optimizers, it remains an open question whether LARS and LAMB have any benefit over traditional, generic algorithms. In this work we demonstrate that standard optimization algorithms such as Nesterov momentum and Adam can match or exceed the results of LARS and LAMB at large batch sizes. Our results establish new, stronger baselines for future comparisons at these batch sizes and shed light on the difficulties of comparing optimizers for neural network training more generally.

1 Introduction

The paper argues that standard optimizers can match or improve on LARS and LAMB at large batch sizes, addressing a gap in fair comparisons. It also situates this question amid challenges in scaling batch sizes and comparing optimizers.

1 Introduction

  • Larger batches can raise throughput, but preserving validation performance may require retuning hyperparameters and replacing gradient-noise regularization.At sufficiently large batch sizes, diminishing returns mean training steps no longer decrease proportionally with batch size.Source: §1 · pp. 1–2
  • LARS and LAMB were adopted for large-batch benchmarks, but prior comparisons left open whether standard optimizers could match them.LARS was used in MLPerf ResNet-50 submissions, while LAMB was proposed for BERT after LARS was judged ineffective across workloads.Source: §1 · p. 2
  • Nesterov matches LARS on ResNet-50 at batch 32,768, while Adam beats LAMB on largest-batch BERT pretraining and downstream metrics.Source: §1 · p. 2
  • Adam reaches 90.46 F1 in 7,818 steps at batch 65,536 and improves on LAMB's reported wall-time record.Source: §1 · pp. 2–3

1.1 Related work

  • Prior work found optimizer rankings difficult to compare fairly because hyperparameter-tuning protocols shape results, while MLPerf restricts algorithms and tuning.Source: §1.1 · p. 3

2 Matching LARS on ImageNet

The section reports a 50-run Nesterov match to the 32,768-batch ResNet-50 LARS result and explains the configuration choices and tuning caveats behind that comparison.

2 Matching LARS on ImageNet

  • The benchmark targets 75.9% validation accuracy on ResNet-50 v1.5 and compares median accuracy over 50 runs at a fixed step budget.The section matches the published 64-epoch LARS result using 2,512 steps at batch size 32,768.Source: §2 · pp. 3–4

2.1 Nesterov momentum at batch size 32k

  • 75.97% median accuracy for Configuration A matched LARS; Configuration B reached 75.92%, with the 75.9% target reached in 34/50 and 29/50 trials, respectively.The comparison used 2,512 update steps at batch size 32,768.Source: §2.1 · p. 4, §2 · p. 4

2.1.1 Nesterov Momentum Optimizer

  • Nesterov momentum is a heavy-ball variant whose update uses batch-averaged loss gradients, momentum, and a step-specific learning rate.The authors prefer it partly because it tolerates larger momentum values.Source: §2.1.1 · p. 4

2.1.2 Batch normalization

  • Configuration B matched LARS’s batch-normalization ε and ρ but re-tuned residual-block γ0; choosing γ0 between 0 and 1 was important.Source: §2.1.2 · p. 5

2.1.3 Regularization

  • Both pipelines omit L2 regularization on biases and batch-normalization γ/β, while LARS uses heavy-ball momentum—not LARS—for those parameters.The paper calls this detail extremely important for fastest training speed.Source: §2.1.3 · p. 5

2.1.4 Learning rate schedule

2.1.5 Comparing Nesterov momentum and LARS

  • Warmup power 2 favored Nesterov and 1 favored LARS; Nesterov had higher training accuracy despite similar validation performance, suggesting LARS normalization may act as regularization.Table 2 reports medians over 50 runs.Source: §2.1.5 · pp. 5–6, Table 2 · p. 6

2.2 Lessons learned

  • A small final pipeline delta understates matching effort: subtle tricks, schedule sensitivity, and narrow near-peak search regions complicate fair comparisons.The authors do not claim either optimizer wins without equally intensive independent tuning.Source: §2.2 · p. 6, §2.1.5 · p. 6

3 Stronger BERT pretraining speed baselines

The section challenges claims that Adam cannot match LAMB for BERT pretraining, presenting tuned Adam baselines with stronger reported SQuAD v1.1 results. The comparison uses the same two-phase pretraining setup and downstream F1 target as You et al. [2019].

3 Stronger BERT pretraining speed baselines

  • 7,818 steps was reported as a new BERT pretraining speed state of the art; batch-size-65,536 Adam nearly matched hybrid LAMB’s F1 in fewer steps.The authors also report improving on the 76-minute record claimed in You et al. [2019].Source: §3 · p. 7
  • The experiments matched You et al. [2019]’s two pretraining phases, with maximum sequence lengths of 128 then 512, and used downstream SQuAD v1.1 F1 as the target metric.The authors note that using a downstream metric can confound optimizer-efficiency measurement, although better pretraining performance correlated with higher fine-tuned F1 here.Source: §3 · pp. 6–7
  • Adam with decoupled weight decay beat reported LAMB and Adam results, with better SQuAD v1.1 development F1 at batch sizes 32,768 and 65,536-then-32,768.At batch size 32,768, Adam also achieved pretraining loss 1.277 versus LAMB’s 1.342.Source: Table 3 · p. 7, §3 · p. 7
  • The authors attribute improved Adam results partly to independent hyperparameter tuning and implementation differences: standard bias correction and a corrected learning-rate decay schedule.They say bias correction acts as an additional step-size warm-up and may improve early-training stability.Source: §3 · p. 7

4 Investigating a less stringent step budget

With a 6,000-step budget and simpler cosine schedule, several configurations exceeded the 75.9% MLPerf target at batch size 32,768. Pure LAMB performed poorly, while pure LARS and Nesterov momentum performed similarly.

4 Investigating a less stringent step budget

  • The authors used a 6,000-step budget and a simpler cosine schedule to reduce the dominance of highly tuned schedules and examine optimizer behavior more clearly.The budget was more generous than the 2,512 steps in Section 2 and was expected to allow validation accuracy above the 75.9% target.Source: §4 · p. 7
  • The paper excludes MLPerf’s warm-start partial-training task; its state-of-the-art “LARS” pipeline applies momentum to batch-normalization and bias parameters and LARS to others.The authors also note uncertainty about whether LAMB was applied to all parameters in You et al. [2019].Source: §4 · p. 7
  • The experiments varied whether layer-wise normalization applied to batch-normalization and ResNet-bias parameters, and compared Adam and LAMB with epsilon values of 10^-6 and 10^-8.They tuned initial learning rate and weight decay; other hyperparameters generally followed the LARS pipeline, except γ0 was set to 1.0.Source: §4 · p. 8
  • Multiple optimizer configurations exceeded the 75.9% MLPerf target with minimal tuning under a 6,000-step budget and cosine learning-rate decay.The authors conclude that stringent step budgets make hyperparameter tuning trickier, rather than large-batch training being fundamentally unstable.Source: §4 · p. 8
  • These experiments do not establish which optimizers are easiest to tune, but LARS and LAMB showed no apparent tuning-ease advantage in this setting.LAMB performed about the same as Adam with the same epsilon; Nesterov momentum and LARS had similar performance.Source: §4 · p. 8
  • Pure LAMB performed extremely poorly, while pure LARS and Nesterov momentum performed roughly the same after 6,000 steps at batch size 32,768.LAMB obtained reasonable results when layer-wise normalization was not applied to batch-normalization and ResNet-bias parameters; the MLPerf LARS pipeline also did not apply LARS to all parameters.Source: Table 4 · p. 8, §4 · p. 8

5 Discussion

The discussion argues that standard optimizers suffice on the tested large-batch workloads, while optimizer comparisons remain sensitive to tuning and evaluation choices. It cautions against treating theoretical bounds or validation results alone as evidence of practical optimizer superiority.

5 Discussion

  • Standard Momentum and Adam can match strong results at the large batch sizes considered, so new large-batch optimizers must be compared against these baselines.The paper also shows Adam satisfies a convergence bound similar to LAMB’s.Source: §5 · p. 8
  • The convergence bounds cited for LARS and LAMB rely on unrealistic assumptions and are loose worst-case upper bounds, not accurate reflections of practical performance.The authors leave open whether layer-wise normalization can help optimization or regularization, while reporting no evidence that LARS and LAMB scale better than Momentum and Adam on the cited tasks and batch sizes.Source: §5 · p. 9
  • Matching state-of-the-art training speed required implementation details and extensive retuning that are not often reported, and those details interact with optimizer choice.The authors say everything must be retuned for a new optimizer in this highly tuned regime.Source: §5 · p. 9
  • Nesterov momentum achieved higher ResNet-50 training accuracy than the tested LARS configuration despite roughly equal validation accuracy, motivating separate scrutiny of optimization and regularization effects.The authors argue that if a method’s primary benefit is regularization, it should be compared with regularization baselines rather than only other optimizers.Source: §5 · p. 9
  • Optimizer comparisons depend crucially on equal hyperparameter-tuning effort, especially for learning-rate schedules tightly coupled to the optimizer.The discussion recommends caution toward claims of one optimizer’s superiority when baseline schedules do not receive comparable tuning.Source: §5 · p. 9

6 Conclusion

The conclusion argues that standard optimizers can match or exceed LARS and LAMB without layer-wise normalization, while emphasizing the difficulty of fair optimizer comparisons.

6 Conclusion

  • Standard optimizers without layer-wise normalization can match or exceed the large-batch results used to justify LARS and LAMB.The authors call for comparisons that distinguish optimizer gains from better tuning or changes to the training pipeline.Source: §6 · pp. 9–10
  • The convergence bounds assume no momentum, and the LAMB Lavg bound assumes β2 = 0 rather than the typical 0.999.The authors also note that Lavg could remain large if L∞ is large, leaving empirical analysis for future work.Source: §6 · p. 9
  • Fair comparisons must account for tuning effort and separate gains from a new optimizer from gains due to tuning or pipeline changes.The authors frame optimizer quality across workloads in relation to the additional tuning required.Source: §6 · p. 10
  • The authors call for standardized competitive benchmarks with workloads not fully known to optimizer submitters, and better incentives for rigorous baselines and comparisons.They cite the MLCommons Algorithmic Efficiency Working Group as an effort in this direction.Source: §6 · p. 10

A Convergence Proofs

The appendix derives convergence bounds for LARS, LAMB, and Adam, while cautioning that such worst-case bounds do not establish how optimizers compare in practice.

A Convergence Proofs

  • These bounds are loose upper bounds on worst-case behavior, and the authors caution that comparing their forms does not reflect optimizers’ relative behavior in reality.They note that the derivation can follow similar algebraic operations while changing a minus sign to a plus for convenience.Source: A Convergence Proofs · p. 12
  • The analysis defines a stochastic optimization objective and assumes bounded minibatch-gradient variance when deriving convergence bounds.The setup specifies neural-network parameters, a smooth possibly nonconvex loss, batch and step quantities, and layerwise Lipschitz constants.Source: A Convergence Proofs · p. 12
  • For LARS, the cited convergence analysis requires b = T, β1 = 0, λ = 0, and bounded positive φ(v); separately, a LAMB convergence rate is derived for β2 = 0.Source: A Convergence Proofs · p. 12
  • The authors derive an Adam bound for β2 > 0 and report a similarly derived β2 = 0 case, which they describe as unrealistic in practice.The derivation bounds a term using minibatch variance, sums the resulting inequalities, and uses a telescoping sum to obtain the final bound.Source: A Convergence Proofs · pp. 12, 14

B.1 ResNet-50 training benchmark

This appendix section documents the ResNet-50 ImageNet experiment’s hardware, software, and image preprocessing.

B.1 ResNet-50 training benchmark

  • ResNet-50 experiments ran on Google TPUs, typically TPUv2-256 or TPUv3-128 for batch size 32,768, using Jax and Flax.The experiments used the standard train/validation split from prior literature and the MLPerf competition.Source: §B.1 · p. 14
  • ImageNet preprocessing used TensorFlow functions for distorted bounding-box sampling, JPEG decoding and cropping, resizing, and random horizontal flipping.The passage lists the sequence of preprocessing functions.Source: §B.1 · p. 14
  • The appendix provides a link to the full code for the experiment.Source: §B.1 · p. 14

B.2 BERT pre-training

This appendix section describes the BERT pretraining and fine-tuning setup and the hyperparameter search used for Adam.

B.2 BERT pre-training

  • BERT experiments matched the official codebase setup and prior train/test splits, using TPUv3-256 or TPUv3-512 at batch size 32,768 and TPUv3-1024 at 65,536.Source: §B.2 · p. 15
  • Pretraining used Wikipedia and Books corpora with sequence lengths 128 and 512, then fine-tuned on SQuaD v1.1 using dev-set F1 without additional fine-tuning.The setup followed the LAMB optimizer fine-tuning protocol.Source: §B.2 · p. 15
  • Adam hyperparameters were selected by quasi-random search using masked-language-model accuracy on 10,000 training examples, with up to 130 trials for a phase.Search varied learning rate, β1, β2, warmup polynomial power, and weight decay; pretraining phases used independent hyperparameters.Source: §B.2 · p. 15

C Nesterov ablations

The ablations test which components of Configuration A matter, and only the “BN & LR tuning” variant retains median accuracy above 75.9%.

C Nesterov ablations

  • The ablations separately vary the ResNet version, batch-normalization initialization, virtual batch size, BN and learning-rate tuning, and L2 regularization coverage.“BN & LR tuning” uses Configuration B values for pdecay, twarmup, η0, ρ, and ϵ; “L2 variables” applies L2 regularization to all variables.Source: C Nesterov ablations · p. 15
  • Figure 3 uses six finetuning runs from one pretraining checkpoint at batch sizes 32,768, mixed 65,536–32,768, and 65,536.The runs are presented to show the stability of the results across those settings.Source: Figure 3 · p. 16
  • Only “BN & LR tuning” keeps median accuracy above the 75.9% target; other ablations drop 0.1%–0.3%.Figure 4 summarizes accuracy distributions over 50 training runs.Source: Figure 4 · p. 16, C Nesterov ablations · p. 16

D.1 Nesterov momentum training speed on ResNet-50

The appendix records how Nesterov configurations and BERT Adam search spaces were tuned, alongside trial budgets and the progression of ResNet-50 searches.

D.1 Nesterov momentum training speed on ResNet-50

  • Configuration A tunes a wide set of Nesterov hyperparameters, while Configuration B reverts less impactful values to the LARS baseline or a simpler warmup value.Configuration B was included to show the minimal changes needed to reach the target accuracy.Source: §D.1 · p. 16

D.2 Adam on BERT

  • Adam tuning on BERT covered all pipeline phases, with the search space and best pretraining hyperparameters documented in Tables 6 and 7.All reported trials used ϵ = 10^-11.Source: §D.2 · p. 16, Table 6 · p. 17, Table 7 · p. 17

D.3 Less stringent step budget on ResNet-50

  • For 6,000-step cosine-schedule experiments, tuning covered η and λ; LAMB on all parameters required 67 trials to obtain 25 non-NaN runs.The search used at least 50 initial-range trials, 25 refined-range trials, and five seeds for the best point; Nesterov used L2 regularization while Adam used weight decay.Source: §D.3 · p. 17, Table 8 · p. 18

D.4 Nesterov ResNet50 search space chronology

  • The Nesterov search moved from 2,815- and 2,658-step trials to 2,512 steps, then narrowed learning-rate-schedule and L2 searches.The authors document the search chronology, with later trials focusing increasingly on the schedule shape and L2 regularization.Source: §D.4 · p. 18, Table 20 · p. 22, Table 21 · p. 22
Loading 2102.06356v3…