# DIVEBATCH: Accelerating Model Training Through Gradient-Diversity Aware Batch Size Adaptation

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 14 of its 32 references that are in the map.
- First posted on arXiv 19 Sep 2025.

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

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [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)
- [Don't Decay the Learning Rate, Increase the Batch Size](https://mlsanity.com/paper/1711.00489/) (2017)
- [An Empirical Model of Large-Batch Training](https://mlsanity.com/paper/1812.06162/) (2018)
- [SGD: General Analysis and Improved Rates](https://mlsanity.com/paper/1901.09401/) (2019)
- [Handbook of Convergence Theorems for (Stochastic) Gradient Methods](https://mlsanity.com/paper/2301.11235/) (2023)
- [On the Convergence of Stochastic Gradient Descent with Adaptive Stepsizes](https://mlsanity.com/paper/1805.08114/) (2018)
- [Improving Generalization Performance by Switching from Adam to SGD](https://mlsanity.com/paper/1712.07628/) (2017)
- [Revisiting Small Batch Training for Deep Neural Networks](https://mlsanity.com/paper/1804.07612/) (2018)
- [Coupling Adaptive Batch Sizes with Learning Rates](https://mlsanity.com/paper/1612.05086/) (2016)
- [AdaBatch: Adaptive Batch Sizes for Training Deep Neural Networks](https://mlsanity.com/paper/1712.02029/) (2017)
- [Active Mini-Batch Sampling using Repulsive Point Processes](https://mlsanity.com/paper/1804.02772/) (2018)
- [Adaptive Learning Rates for Faster Stochastic Gradient Methods](https://mlsanity.com/paper/2208.05287/) (2022)

## Abstract (arXiv)

The goal of this paper is to accelerate the training of machine learning models, a critical challenge since the training of large-scale deep neural models can be computationally expensive. Stochastic gradient descent (SGD) and its variants are widely used to train deep neural networks. In contrast to traditional approaches that focus on tuning the learning rate, we propose a novel adaptive batch size SGD algorithm, DiveBatch, that dynamically adjusts the batch size. Adapting the batch size is challenging: using large batch sizes is more efficient due to parallel computation, but small-batch training often converges in fewer epochs and generalizes better. To address this challenge, we introduce a data-driven adaptation based on gradient diversity, enabling DiveBatch to maintain the generalization performance of small-batch training while improving convergence speed and computational efficiency. Gradient diversity has a strong theoretical justification: it emerges from the convergence analysis of SGD. Evaluations of DiveBatch on synthetic and CiFar-10, CiFar-100, and Tiny-ImageNet demonstrate that DiveBatch converges significantly faster than standard SGD and AdaBatch (1.06 -- 5.0x), with a slight trade-off in performance.