mlsanityData to 7 Oct 2026

Batch size invariant Adam

arXiv 2402.18824 ↗First posted 29 Feb 2024cs.LGcited by 4

Not on the map. Why?

Abstract (arXiv)

We propose a batch size invariant version of Adam, for use in large-scale, distributed settings, in which the mini-batch is divided into micro-batches which are distributed among worker nodes. For the v term, standard Adam first computes the average over micro-batch gradients, then squares, while in the batch size invariant Adam proposed here, we first square the micro-batch gradients, then average. Previous work (e.g. Malladi et al. 2022) used an alternative approach that involved a square-root scaling of the learning rate, but this approach requires strong assumptions to work; in particular that the gradient variance dominates the square of the expected gradient. In contrast, the approach proposed here gives batch size invariance without this assumption. We confirm that in practice our scheme gives batch size invariance in a much larger range of scenarios than the previous approach.

Key facts. Cited by 4 papers in the map (data through 7 Oct 2026). Builds on 24 of its 27 references that are in the map. First posted on arXiv 29 Feb 2024.

Builds on

24 of 27 references.

References · 24

  1. 2015Deep Residual Learning for Image Recognitioncited by 51,809
  2. 2020An Image is Worth 16x16 Words: Transformers for Image Recognition at Scalecited by 29,915
  3. 2023Llama 2: Open Foundation and Fine-Tuned Chat Modelscited by 12,865
  4. 2022OPT: Open Pre-trained Transformer Language Modelscited by 3,464
  5. 2017Accurate, Large Minibatch SGD: Training ImageNet in 1 Hourcited by 1,479
  6. 2016On Large-Batch Training for Deep Learning: Generalization Gap and Sharp Minimacited by 1,337
  7. 2018AutoAugment: Learning Augmentation Policies from Datacited by 728
  8. 2023Scaling Vision Transformers to 22 Billion Parameterscited by 690
  9. 2017Don't Decay the Learning Rate, Increase the Batch Sizecited by 313
  10. 2017Stochastic Gradient Descent as Approximate Bayesian Inferencecited by 298
  11. 2018An Empirical Model of Large-Batch Trainingcited by 268
  12. 2014One weird trick for parallelizing convolutional neural networkscited by 264
12 more
  1. 2018Measuring the Effects of Data Parallelism on Neural Network Trainingcited by 190
  2. 2023Small-scale proxies for large-scale Transformer training instabilitiescited by 182
  3. 2017The Power of Interpolation: Understanding the Effectiveness of SGD in Modern Over-parametrized Learningcited by 150
  4. 2017A Bayesian Perspective on Generalization and Stochastic Gradient Descentcited by 134
  5. 2019Which Algorithmic Choices Matter at Which Batch Sizes? Insights From a Noisy Quadratic Modelcited by 116
  6. 2022On the SDEs and Scaling Rules for Adaptive Gradient Algorithmscited by 100
  7. 2019BackPACK: Packing more into backpropcited by 77
  8. 2020Learning Rates as a Function of Batch Size: A Random Matrix Theory Approach to Neural Network Trainingcited by 50
  9. 2021Batch size-invariance for policy optimizationcited by 33
  10. 2018On the Computational Inefficiency of Large Batch Sizes for Stochastic Gradient Descentcited by 25
  11. 2023MoMo: Momentum Models for Adaptive Learning Ratescited by 24
  12. 2023Adam Accumulation to Reduce Memory Footprints of both Activations and Gradients for Large-scale DNN Trainingcited by 5

Built on by

Cited by year

Citing papers per year: 2024: 2; 2025: 2.

4 papers.

  1. 2024How to set AdamW's weight decay as you scale model and dataset sizecited by 42
  2. 2025Small Batch Size Training for Language Models: When Vanilla SGD Works, and Why Gradient Accumulation Is Wastefulcited by 36
  3. 2024FAdam: Adam is a natural gradient optimizer using diagonal empirical Fisher informationcited by 17
  4. 2025Per-example gradients: a new frontier for understanding and improving optimizerscited by 2