# Batch size invariant Adam

arXiv: https://arxiv.org/abs/2402.18824 · HTML: https://mlsanity.com/paper/2402.18824/

## Key facts

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

## Builds on (most cited first, 24 of 24)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Llama 2: Open Foundation and Fine-Tuned Chat Models](https://mlsanity.com/paper/2307.09288/) (2023)
- [OPT: Open Pre-trained Transformer Language Models](https://mlsanity.com/paper/2205.01068/) (2022)
- [Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour](https://mlsanity.com/paper/1706.02677/) (2017)
- [On Large-Batch Training for Deep Learning: Generalization Gap and Sharp Minima](https://mlsanity.com/paper/1609.04836/) (2016)
- [AutoAugment: Learning Augmentation Policies from Data](https://mlsanity.com/paper/1805.09501/) (2018)
- [Scaling Vision Transformers to 22 Billion Parameters](https://mlsanity.com/paper/2302.05442/) (2023)
- [Don't Decay the Learning Rate, Increase the Batch Size](https://mlsanity.com/paper/1711.00489/) (2017)
- [Stochastic Gradient Descent as Approximate Bayesian Inference](https://mlsanity.com/paper/1704.04289/) (2017)
- [An Empirical Model of Large-Batch Training](https://mlsanity.com/paper/1812.06162/) (2018)
- [One weird trick for parallelizing convolutional neural networks](https://mlsanity.com/paper/1404.5997/) (2014)
- [Measuring the Effects of Data Parallelism on Neural Network Training](https://mlsanity.com/paper/1811.03600/) (2018)
- [Small-scale proxies for large-scale Transformer training instabilities](https://mlsanity.com/paper/2309.14322/) (2023)
- [The Power of Interpolation: Understanding the Effectiveness of SGD in Modern Over-parametrized Learning](https://mlsanity.com/paper/1712.06559/) (2017)
- [A Bayesian Perspective on Generalization and Stochastic Gradient Descent](https://mlsanity.com/paper/1710.06451/) (2017)
- [Which Algorithmic Choices Matter at Which Batch Sizes? Insights From a Noisy Quadratic Model](https://mlsanity.com/paper/1907.04164/) (2019)
- [On the SDEs and Scaling Rules for Adaptive Gradient Algorithms](https://mlsanity.com/paper/2205.10287/) (2022)
- [BackPACK: Packing more into backprop](https://mlsanity.com/paper/1912.10985/) (2019)
- [Learning Rates as a Function of Batch Size: A Random Matrix Theory Approach to Neural Network Training](https://mlsanity.com/paper/2006.09092/) (2020)
- [Batch size-invariance for policy optimization](https://mlsanity.com/paper/2110.00641/) (2021)
- [On the Computational Inefficiency of Large Batch Sizes for Stochastic Gradient Descent](https://mlsanity.com/paper/1811.12941/) (2018)
- [MoMo: Momentum Models for Adaptive Learning Rates](https://mlsanity.com/paper/2305.07583/) (2023)
- [Adam Accumulation to Reduce Memory Footprints of both Activations and Gradients for Large-scale DNN Training](https://mlsanity.com/paper/2305.19982/) (2023)

## Built on by (most cited first, 5 of 5)

- [How to set AdamW's weight decay as you scale model and dataset size](https://mlsanity.com/paper/2405.13698/) (2024)
- [Small Batch Size Training for Language Models: When Vanilla SGD Works, and Why Gradient Accumulation Is Wasteful](https://mlsanity.com/paper/2507.07101/) (2025)
- [FAdam: Adam is a natural gradient optimizer using diagonal empirical Fisher information](https://mlsanity.com/paper/2405.12807/) (2024)
- [Per-example gradients: a new frontier for understanding and improving optimizers](https://mlsanity.com/paper/2510.00236/) (2025)
- [The Best Optimizer Depends on Batch Size](https://mlsanity.com/paper/2610.08975/) (2026)

## 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.