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 (24)

• Adam: A Method for Stochastic Optimization
• Deep Residual Learning for Image Recognition
• Large Scale GAN Training for High Fidelity Natural Image Synthesis
• Opening the Black Box of Deep Neural Networks via Information
• Don't Decay the Learning Rate, Increase the Batch Size
• Large Batch Training of Convolutional Networks
• Distributed Prioritized Experience Replay
• Train longer, generalize better: closing the generalization gap in large batch training of neural networks
• No More Pesky Learning Rates
• Extremely Large Minibatch SGD: Training ResNet-50 on ImageNet in 15 Minutes
• Adding Gradient Noise Improves Learning for Very Deep Networks
• Qualitatively characterizing neural network optimization problems
• Measuring the Effects of Data Parallelism on Neural Network Training
• Gradient Descent Happens in a Tiny Subspace
• AdaBatch: Adaptive Batch Sizes for Training Deep Neural Networks
• Accelerated Methods for Deep Reinforcement Learning
• Measuring the Intrinsic Dimension of Objective Landscapes
• ImageNet Training in Minutes
• Understanding Short-Horizon Bias in Stochastic Meta-Optimization
• Big Batch SGD: Automated Inference using Adaptive Batch Sizes
• Coupling Adaptive Batch Sizes with Learning Rates
• A Bayesian Perspective on Generalization and Stochastic Gradient Descent
• Gradient Diversity: a Key Ingredient for Scalable Distributed Learning
• Large Scale Language Modeling: Converging on 40GB of Text in Four Hours

Cited by in corpus (66)

• Scaling Laws for Neural Language Models
• Dota 2 with Large Scale Deep Reinforcement Learning
• Improved Denoising Diffusion Probabilistic Models
• Don't Use Large Mini-Batches, Use Local SGD
• Scaling Laws for Autoregressive Generative Modeling
• Measuring the Algorithmic Efficiency of Neural Networks
• ZeRO: Memory Optimizations Toward Training Trillion Parameter Models
• Batch Normalization Biases Residual Blocks Towards the Identity Function in Deep Networks
• Train Large, Then Compress: Rethinking Model Size for Efficient Training and Inference of Transformers
• 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
• Checkmate: Breaking the Memory Wall with Optimal Tensor Rematerialization
• Large batch size training of neural networks with adversarial training and second-order information
• SEED RL: Scalable and Efficient Deep-RL with Accelerated Central Inference
• The Full Spectrum of Deepnet Hessians at Scale: Dynamics with SGD Training and Sample Size
• A Closer Look at Deep Policy Gradients
• Mastering Complex Control in MOBA Games with Deep Reinforcement Learning
• Pollux: Co-adaptive Cluster Scheduling for Goodput-Optimized Deep Learning
• 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
• Which Algorithmic Choices Matter at Which Batch Sizes? Insights From a Noisy Quadratic Model
• Wide-minima Density Hypothesis and the Explore-Exploit Learning Rate Schedule
• Using a thousand optimization tasks to learn hyperparameter search strategies
• Convergence Analysis and Latency Minimization for Semi-Federated Learning in Massive IoT Networks
• Parallel Training of Deep Networks with Local Updates
• Learning Rate Annealing Can Provably Help Generalization, Even for Convex Problems
• Scaling Scaling Laws with Board Games
• ExT5: Towards Extreme Multi-Task Scaling for Transfer Learning
• Study on the Large Batch Size Training of Neural Networks Based on the Second Order Gradient
• Robust Learning Rate Selection for Stochastic Optimization via Splitting Diagnostic
• GraVAC: Adaptive Compression for Communication-Efficient Distributed DL Training
• Scavenger: A Cloud Service for Optimizing Cost and Performance of ML Training
• Stochastic Training is Not Necessary for Generalization
• Robust Optimization for Multilingual Translation with Imbalanced Data
• A Loss Curvature Perspective on Training Instability in Deep Learning
• Accelerating Distributed ML Training via Selective Synchronization
• Ensemble Model Patching: A Parameter-Efficient Variational Bayesian Neural Network
• Modularization of End-to-End Learning: Case Study in Arcade Games
• Towards Understanding Iterative Magnitude Pruning: Why Lottery Tickets Win
• Stagewise Enlargement of Batch Size for SGD-based Learning
• Stochastic Optimization with Laggard Data Pipelines
• RLScheduler: An Automated HPC Batch Job Scheduler Using Reinforcement Learning
• A Generalizable Approach to Learning Optimizers
• Distributed Training and Optimization Of Neural Networks
• Layered gradient accumulation and modular pipeline parallelism: fast and efficient training of large language models
• On Large-Cohort Training for Federated Learning
• Inefficiency of K-FAC for Large Batch Size Training
• Adaptive Elastic Training for Sparse Deep Learning on Heterogeneous Multi-GPU Servers
• Flexible Communication for Optimal Distributed Learning over Unpredictable Networks
• Neural Machine Translation: A Review and Survey
• Understanding Learning Dynamics for Neural Machine Translation
• Applying Cyclical Learning Rate to Neural Machine Translation
• 2kenize: Tying Subword Sequences for Chinese Script Conversion
• Extrapolation for Large-batch Training in Deep Learning
• Improving compute efficacy frontiers with SliceOut
• Scaling Distributed Training of Flood-Filling Networks on HPC Infrastructure for Brain Mapping
• HydaLearn: Highly Dynamic Task Weighting for Multi-task Learning with Auxiliary Tasks
• Resource-Adaptive Successive Doubling for Hyperparameter Optimization with Large Datasets on High-Performance Computing Systems
• What can linear interpolation of neural network loss landscapes tell us?
• Batch size-invariance for policy optimization
• Data optimization for large batch distributed training of deep neural networks
• Full waveform inversion with random shot selection using adaptive gradient descent
• Using a one dimensional parabolic model of the full-batch loss to estimate learning rates during training
• LRTuner: A Learning Rate Tuner for Deep Neural Networks