DIVEBATCH: Accelerating Model Training Through Gradient-Diversity Aware Batch Size Adaptation
This addresses the computational expense of training large-scale deep neural networks, offering a practical acceleration method with some trade-offs.
The paper tackles the challenge of accelerating machine learning model training by proposing DiveBatch, an adaptive batch size SGD algorithm that dynamically adjusts batch size based on gradient diversity. Results show DiveBatch converges 1.06-5.0x faster than standard SGD and AdaBatch on datasets like CIFAR-10, CIFAR-100, and Tiny-ImageNet, with a slight performance trade-off.
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.