An Empirical Model of Large-Batch Training

Abstract

In an increasing number of domains it has been demonstrated that deep learning models can be trained using relatively large batch sizes without sacrificing data efficiency. However the limits of this massive data parallelism seem to differ from domain to domain, ranging from batches of tens of thousands in ImageNet to batches of millions in RL agents that play the game Dota 2. To our knowledge there is limited conceptual understanding of why these limits to batch size differ or how we might choose the correct batch size in a new domain. In this paper, we demonstrate that a simple and easy-to-measure statistic called the gradient noise scale predicts the largest useful batch size across many domains and applications, including a number of supervised learning datasets (MNIST, SVHN, CIFAR-10, ImageNet, Billion Word), reinforcement learning domains (Atari and Dota), and even generative model training (autoencoders on SVHN). We find that the noise scale increases as the loss decreases over a training run and depends on the model size primarily through improved model performance. Our empirically-motivated theory also describes the tradeoff between compute-efficiency and time-efficiency, and provides a rough model of the benefits of adaptive batch-size training.

References in corpus (7)

• Large Batch Training of Convolutional Networks
• No More Pesky Learning Rates
• Extremely Large Minibatch SGD: Training ResNet-50 on ImageNet in 15 Minutes
• Qualitatively characterizing neural network optimization problems
• Gradient Descent Happens in a Tiny Subspace
• ImageNet Training in Minutes
• Gradient Diversity: a Key Ingredient for Scalable Distributed Learning

Cited by in corpus (31)

• Scaling Laws for Neural Language Models
• Improved Denoising Diffusion Probabilistic Models
• Scaling Laws for Autoregressive Generative Modeling
• Measuring the Algorithmic Efficiency of Neural Networks
• rlpyt: A Research Code Base for Deep Reinforcement Learning in PyTorch
• Data Valuation using Reinforcement Learning
• On the Origin of Implicit Regularization in Stochastic Gradient Descent
• Finite Versus Infinite Neural Networks: an Empirical Study
• SEED RL: Scalable and Efficient Deep-RL with Accelerated Central Inference
• Stochastic Weight Averaging in Parallel: Large-Batch Training that Generalizes Well
• Dynamic Mini-batch SGD for Elastic Distributed Training: Learning in the Limbo of Resources
• Learning Rate Annealing Can Provably Help Generalization, Even for Convex Problems
• Scaling Scaling Laws with Board Games
• Study on the Large Batch Size Training of Neural Networks Based on the Second Order Gradient
• A Loss Curvature Perspective on Training Instability in Deep Learning
• Towards Understanding Iterative Magnitude Pruning: Why Lottery Tickets Win
• Modularization of End-to-End Learning: Case Study in Arcade Games
• Ensemble Model Patching: A Parameter-Efficient Variational Bayesian Neural Network
• A Generalizable Approach to Learning Optimizers
• Stagewise Enlargement of Batch Size for SGD-based Learning
• On Large-Cohort Training for Federated Learning
• Layered gradient accumulation and modular pipeline parallelism: fast and efficient training of large language models
• 2kenize: Tying Subword Sequences for Chinese Script Conversion
• Extrapolation for Large-batch Training in Deep Learning
• Applying Cyclical Learning Rate to Neural Machine Translation
• Understanding Learning Dynamics for Neural Machine Translation
• Adaptive Elastic Training for Sparse Deep Learning on Heterogeneous Multi-GPU Servers
• HydaLearn: Highly Dynamic Task Weighting for Multi-task Learning with Auxiliary Tasks
• Data optimization for large batch distributed training of deep neural networks
• Full waveform inversion with random shot selection using adaptive gradient descent
• LRTuner: A Learning Rate Tuner for Deep Neural Networks