# AdAdaGrad: Adaptive Batch Size Schemes for Adaptive Gradient Methods

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

## Key facts

- Cited by 8 papers in the map (data through 14 Oct 2026).
- Builds on 58 of its 76 references that are in the map.
- First posted on arXiv 17 Feb 2024.

## Builds on (most cited first, 30 of 58)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Scaling Laws for Neural Language Models](https://mlsanity.com/paper/2001.08361/) (2020)
- [Training Compute-Optimal Large Language Models](https://mlsanity.com/paper/2203.15556/) (2022)
- [Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism](https://mlsanity.com/paper/1909.08053/) (2019)
- [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)
- [ZeRO: Memory Optimizations Toward Training Trillion Parameter Models](https://mlsanity.com/paper/1910.02054/) (2019)
- [Scaling Language Models: Methods, Analysis & Insights from Training Gopher](https://mlsanity.com/paper/2112.11446/) (2021)
- [Optimization Methods for Large-Scale Machine Learning](https://mlsanity.com/paper/1606.04838/) (2016)
- [Stochastic First- and Zeroth-order Methods for Nonconvex Stochastic Programming](https://mlsanity.com/paper/1309.5549/) (2013)
- [Large Batch Optimization for Deep Learning: Training BERT in 76 minutes](https://mlsanity.com/paper/1904.00962/) (2019)
- [PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel](https://mlsanity.com/paper/2304.11277/) (2023)
- [Using DeepSpeed and Megatron to Train Megatron-Turing NLG 530B, A Large-Scale Generative Language Model](https://mlsanity.com/paper/2201.11990/) (2022)
- [Why gradient clipping accelerates training: A theoretical justification for adaptivity](https://mlsanity.com/paper/1905.11881/) (2019)
- [Don't Decay the Learning Rate, Increase the Batch Size](https://mlsanity.com/paper/1711.00489/) (2017)
- [Lower Bounds for Non-Convex Stochastic Optimization](https://mlsanity.com/paper/1912.02365/) (2019)
- [An Empirical Model of Large-Batch Training](https://mlsanity.com/paper/1812.06162/) (2018)
- [Train longer, generalize better: closing the generalization gap in large batch training of neural networks](https://mlsanity.com/paper/1705.08741/) (2017)
- [The Marginal Value of Adaptive Gradient Methods in Machine Learning](https://mlsanity.com/paper/1705.08292/) (2017)
- [Measuring the Effects of Data Parallelism on Neural Network Training](https://mlsanity.com/paper/1811.03600/) (2018)
- [AdaGrad stepsizes: Sharp convergence over nonconvex landscapes](https://mlsanity.com/paper/1806.01811/) (2018)
- [A Simple Convergence Proof of Adam and Adagrad](https://mlsanity.com/paper/2003.02395/) (2020)
- [Better Theory for SGD in the Nonconvex World](https://mlsanity.com/paper/2002.03329/) (2020)
- [On the Convergence of Stochastic Gradient Descent with Adaptive Stepsizes](https://mlsanity.com/paper/1805.08114/) (2018)
- [Stochastic Variance Reduction for Nonconvex Optimization](https://mlsanity.com/paper/1603.06160/) (2016)
- [Fast and Faster Convergence of SGD for Over-Parameterized Models and an Accelerated Perceptron](https://mlsanity.com/paper/1810.07288/) (2018)
- [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)

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

- [Enhancing Efficiency of Safe Reinforcement Learning via Sample Manipulation](https://mlsanity.com/paper/2405.20860/) (2024)
- [Fast Catch-Up, Late Switching: Optimal Batch Size Scheduling via Functional Scaling Laws](https://mlsanity.com/paper/2602.14208/) (2026)
- [Communication-Efficient Adaptive Batch Size Strategies for Distributed Local Gradient Methods](https://mlsanity.com/paper/2406.13936/) (2024)
- [Adaptive Batch Size Schedules for Distributed Training of Language Models with Data and Model Parallelism](https://mlsanity.com/paper/2412.21124/) (2024)
- [Adaptive Batch Sizes Using Non-Euclidean Gradient Noise Scales for Stochastic Sign and Spectral Descent](https://mlsanity.com/paper/2602.03001/) (2026)
- [Efficient Training of Deep Networks using Guided Spectral Data Selection: A Step Toward Learning What You Need](https://mlsanity.com/paper/2507.04269/) (2025)
- [AdLoCo: adaptive batching significantly improves communications efficiency and convergence for Large Language Models](https://mlsanity.com/paper/2508.18182/) (2025)
- [How Does Local Landscape Geometry Evolve in Language Model Pre-Training?](https://mlsanity.com/paper/2609.39767/) (2026)

## Abstract (arXiv)

The choice of batch size in minibatch stochastic gradient optimization is critical for both optimization and generalization performance in large-scale model training. Although large-batch training is arguably the dominant paradigm in large-scale deep learning because of hardware advances, model generalization often deteriorates relative to small-batch training, leading to the so-called "generalization gap." To mitigate this issue, we investigate adaptive batch size strategies derived from adaptive sampling methods, which were originally developed for stochastic gradient descent. Given the strong interplay between learning rates and batch sizes, together with the prevalence of adaptive gradient methods in deep learning, we emphasize the need for adaptive batch size strategies in these settings. We introduce AdAdaGrad and its scalar variant AdAdaGradNorm, which progressively increase batch sizes during training while performing updates with AdaGrad and AdaGradNorm, respectively. We prove that AdAdaGradNorm converges with high probability at a rate of mathscrO(1/K) to a first-order stationary point of a smooth nonconvex function within K iterations. AdAdaGrad also exhibits similar convergence properties when combined with a novel coordinate-wise variant of our adaptive batch size strategy. We corroborate our theoretical claims with image-classification experiments that highlight the merits of the proposed schemes in terms of both training efficiency and model generalization. Our work highlights the potential of adaptive batch size strategies for adaptive gradient optimizers in large-scale model training.