PaperlayerView exact revision on arXiv

Source-linked AI summary

An Empirical Model of Large-Batch Training

Sam McCandlish, Jared Kaplan, Dario Amodei, OpenAI Dota Team

arXiv:1812.06162v1cs.LGstat.ML

TL;DR

The paper addresses the limited ability to predict useful batch sizes across training domains. It measures gradient noise scale and develops an empirically tested theory connecting it to batch-size efficiency. Across diverse tasks, the noise scale approximately predicts the critical batch size and the associated compute–time tradeoff, while supporting adaptive batch-size tuning.

  • Problem

    Useful batch-size limits differ substantially across domains, but there is limited conceptual guidance for predicting them in a new dataset or model.

  • Method

    The paper measures gradient noise scale and develops a theoretical framework relating it to data parallelism, critical batch size, compute–time tradeoffs, and dynamic batch-size adjustment.

  • Results

    Across diverse tasks, simplified gradient noise scale approximately predicts the critical batch size over six orders of magnitude and the tradeoff curve between compute and optimization steps.

  • Takeaways & Limitations

    A single run can estimate the largest useful batch size and reduce trial-and-error hyperparameter searches, while larger batches show predictable diminishing returns.

  • Takeaways & Limitations

    The framework is strictly local and may underestimate the benefits of noise when noise helps avoid bad minima or cope with ill-conditioned landscapes.

Abstract

from arXiv · show

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.

1 Introduction

Deep-learning training is increasingly parallelized with large batches, but the useful batch-size limit varies widely across tasks and is difficult to predict. The paper introduces the gradient noise scale as an empirical predictor and tests a theory of batch-size, time, and compute tradeoffs across diverse applications.

  • Motivation: Parallelism can make training practically feasible, but insufficient parallelization leaves models dependent on excessive serial time.Data parallelism splits batches across devices, requiring both fast communication and algorithmically effective large batches.
  • Motivation: Large-batch limits vary across datasets and models, yet researchers commonly discover effective batch sizes through costly experimentation and tuning.Warmup periods or unusual learning-rate schedules may be required, allowing useful batch sizes to remain undiscovered.
  • Contribution: The gradient noise scale, an easily measured gradient signal-to-noise statistic, approximately predicts the largest efficient batch size across a wide range of tasks.The paper combines elementary theory with extensive empirical testing.
  • Contribution: Training is predicted to parallelize nearly linearly up to a batch size equal to the noise scale, after which further parallelism yields rapidly diminishing benefits.The framework also predicts increasing noise scale during training, larger values for complex tasks, and weak dependence on model size itself.
  • Empirical evaluation: Experiments across 8 supervised, reinforcement-learning, and generative tasks found order-of-magnitude agreement between noise scale and largest usable batch size.The study also reports predicted degradation of parallelism gains, increasing noise scale during training, and benefits from dynamic batch-size tuning.

2 Theory and Predictions for the Gradient Noise Scale

The paper models how batch size controls the tradeoff between optimization speed and computation, with the gradient noise scale setting the approximate boundary of useful data parallelism. It derives predictions for this boundary and its variation during training, while identifying assumptions that limit the theory’s scope.

  • Intuitive Picture: Small batches produce noisy gradient estimates, so aggregating updates through larger batches can improve progress and enable near-linear training speedups.With very large batches, gradients nearly match the true gradient, so further increases mainly add computation without improving the update.
  • Intuitive Picture: The transition between efficient and wasteful batching occurs when gradient signal and variance are balanced, defining the gradient noise scale.This scale marks the bend in the serial-time versus total-compute tradeoff curve.
  • Theory: For a batch B much smaller than Bnoise, increasing B linearly increases progress, whereas for B much larger than Bnoise, further increases have almost no effect.The switch between these regimes occurs at B ≈ Bnoise.
  • Efficiency Tradeoff: The model predicts Bcrit ≈ Bnoise, but the critical batch size depends on the performance level because the noise scale can vary significantly during training.At B = Bcrit, training requires twice the data passes of an optimally data-efficient run and twice the optimization steps of an optimally time-efficient run.
  • Assumptions and Caveats: The framework is local and assumption-dependent: noise may sometimes aid global optimization, poor conditioning can distort per-step progress, and the model does not address generalization.Despite these caveats, the paper reports that the noise scale remains a good empirical predictor of critical batch size and that most runs fit Equation 2.11 well.
  • Predictions: The theory predicts that the speed-efficiency tradeoff follows Equation 2.11 and that Bcrit can be estimated from the simplified noise scale Bsimple at order-of-magnitude accuracy.Bsimple is preferred empirically because it requires less computational expense than the exact noise scale.

3 Experiments

Experiments across supervised learning, reinforcement learning, and generative modeling test whether the simple gradient noise scale predicts useful batch sizes and related efficiency tradeoffs. The results broadly support the predictions, while revealing task-specific deviations.

  • Optimal learning rates: The optimal learning rate initially increases with batch size, then levels off; for Adam and RMSProp, the initial scaling follows ϵ(B) ∝ B^α with α between 0.5 and 1.0.The leveling-off scale is generally somewhat smaller than the typical noise scale.
  • Critical batch size: The critical batch size typically increases by an order of magnitude or more during training, marking the transition from per-example similarity to per-step similarity.It is determined by fitting the observed Pareto fronts.
  • Pareto frontiers: Equation 2.11 fits most time-efficiency versus compute-efficiency Pareto frontiers, with small batches using compute efficiently and large batches requiring fewer optimizer steps.The turning point increases later in training.
  • Results: Across many tasks, the simple noise scale approximately predicts the batch size where returns from increasing scale diminish significantly.The comparison spans tasks whose critical batch sizes vary by six orders of magnitude.
  • Supervised Learning: ImageNet noise scales range from 2,000 to 100,000 during main training and reach hundreds of thousands or millions during fine-tuning.The reported critical batch size is slightly lower than some literature values: 15k versus 64k.
  • Generative Modeling: For generative models, the VAE has larger Bcrit and Bsimple than the autoencoder, but Bsimple is significantly smaller than Bcrit for both.Both generative models have much lower Bsimple than SVHN image classifiers.
  • Model Size Dependence: The noise scale is roughly independent of LSTM size at fixed loss, while larger models reach larger scales because they achieve lower loss.More complex image datasets also have larger noise scales in a way not directly determined by dataset size.
  • Reinforcement Learning: Across Atari games, tradeoff curves generally agree with Equation 2.11, although some small-batch deviations and noisy estimates appear.The authors note that Pong has a lower noise scale consistent with being easier to learn than other Atari games.

4 Related Work

The paper situates its contribution within prior work on large-batch training, gradient noise, adaptive schedules, optimization landscapes, and computational scaling. It combines these threads by linking the gradient noise scale to batch-size selection and efficiency.

  • Large-Batch Training: Prior studies established that large-batch training can work on specific supervised-learning datasets and reinforcement-learning tasks.ImageNet and RL received particular empirical attention, including layer-wise adaptive learning rates.
  • Gradient Noise: Earlier work connected gradient noise to optimization speed, adaptive learning rates, sampling distributions, temperature, and critical batch size.The paper distinguishes its observations from work predicting a dependence on dataset size.
  • Loss Landscapes: Related research studied neural-network loss landscapes, including their limits as guides to overall optimal training and the possible usefulness of early-training noise.Some work linked noise to minima that generalize better.
  • Positioning: The paper’s broader motivation concerns scaling learning with computational and data resources, complementing prior work focused on scaling model size.Its key contributions connect gradient noise to optimization speed and measure critical batch sizes across tasks.

5 Discussion

The discussion presents the noise scale as a practical estimator of useful batch size and a framework for understanding time–compute tradeoffs across tasks. It also identifies important variation and scope limits in the approximation.

  • Main Findings: The simplified gradient noise scale approximately predicts the actual diminishing-return point Bcrit across diverse problems spanning six orders of magnitude.The associated compute–optimization-step tradeoff is roughly hyperbolic, and optimal learning-rate scaling is also approximately predicted.
  • Practical Use: A single, possibly partial, training run can estimate the largest useful batch size and reduce extensive trial-and-error hyperparameter searches.The method also predicts a common form for diminishing returns from larger batches across tasks.
  • Task Complexity: The largest useful batch size is likely correlated with task complexity because the noise scale measures data diversity as seen by the model.This is presented as one aspect of task complexity rather than a complete definition.
  • Time–Compute Tradeoff: Training at Bcrit is a compromise using twice the minimum optimizer steps and twice the minimum data, while larger batches provide diminishing returns.The critical batch size is identified from the time–compute relation at a target loss.
  • Implications: The critical batch size provides a natural compromise between time and compute because it stays within twice both minima.For batches above Bcrit, increasing data parallelism yields diminishing returns.
  • Reinforcement Learning: Dota’s noise scale is roughly a thousand times Atari’s, yet the total number of optimization steps is not much larger.The authors suggest that additional compute for more powerful models may therefore be substantially parallelizable.
  • Limitations: Bsimple/Bcrit varies by about an order of magnitude between tasks, with no explanation yet for differences across autoencoders, VAEs, Dota 1v1, LSTMs, image classification, and Atari.The authors attribute possible discrepancies to omitted Hessian conditioning or global loss-landscape effects.
  • Limitations: The theory focuses on a chosen batch size with optimized learning rate or temperature and analyzes a specific loss-landscape point, leaving warm-up and early–late training relationships unresolved.These issues are identified as relevant to adaptive batch-size schedules.

A Methods

The methods estimate the noise scale from gradient norms at different batch sizes and use Pareto fronts to measure critical batch sizes across training goals. Measurements rely on averaging and smoothing procedures, with learning rates tuned separately for batch-size comparisons.

  • Noise-scale measurement: The noise scale is estimated by comparing gradient norms for small and large batches, yielding estimates of gradient magnitude and gradient covariance trace.The method obtains unbiased estimates of |G|^2 and tr(Σ) from two batch sizes.
  • Estimator caveat: The ratio S/|G|^2 is biased for Bnoise, so the procedure reduces variance by averaging |G|^2 over many batches instead of applying a correction.The bias correction is possible but omitted to minimize complexity.
  • Noise-scale measurement: Data-parallel training provides the two gradient norms with minimal effort by comparing pre- and post-averaging gradients across devices.The local batch is Bsmall, while the global batch after averaging is Bbig.
  • Noise-scale measurement: Exponentially weighted moving averages stabilize per-step estimates before their ratio is used as the noise-scale estimate.The decay parameters are tuned so that the estimates remain stable.
  • Pareto-front construction: Pareto fronts are constructed by grid-searching batch sizes and learning rates, then selecting the fastest run for each target loss, perplexity, or game-score.Different learning rates can be optimal for different goals at the same batch size.
  • Pareto-front construction: Critical-batch-size estimates are fit from these Pareto-front measurements and reported with standard errors, including averages over optimization steps throughout training.Target attainment is determined after smoothing the loss with an exponentially weighted moving average.

A.4 Details of Learning Tasks

The experiments cover supervised learning, reinforcement learning, language modeling, and generative image modeling using several architectures and optimizers. Task-specific implementations vary in datasets, rollout structure, sequence handling, and stochasticity.

  • Experimental scope: The study measures simple noise scale during one well-optimized run for models trained with SGD, momentum, or Adam.The reported tasks span classification, reinforcement learning, language modeling, and generative modeling.
  • Supervised learning: Image-classification experiments include MNIST, SVHN, CIFAR10, and ImageNet with convolutional networks and ResNets.SVHN uses both SGD and Adam, while CIFAR10 and ImageNet use momentum.
  • Reinforcement learning: Reinforcement-learning experiments include Atari and Dota tasks, using A2C for Atari and asynchronous PPO for Dota.Atari batch size varies through the number of parallel environments, and Dota performance is measured with TrueSkill.
  • Language modeling: Language modeling uses 2048-, 1024-, and 512-unit LSTMs on the One Billion Word Benchmark with Adam and byte-pair encoding.The main model uses a 40,000-word vocabulary, 512-dimensional embeddings, gradient clipping, and 20-token sequences.
  • Generative modeling: Generative image modeling uses a Variational Autoencoder and a simple autoencoder with the same InfoGAN-based architecture on SVHN.The simple autoencoder provides a comparison without the VAE’s additional stochasticity.

B Results for All Tasks

Across classification, reinforcement learning, and generative modeling, the experiments compare compute-efficiency and time-efficiency across batch sizes. The critical batch size is compared with simple and true noise-scale estimates, and the noise scale can change substantially during training.

  • Cross-task results: The experiments plot compute–time tradeoff curves and compare critical batch size with simple and true noise-scale estimates across multiple task families.Critical batch size is defined as the point within 50% of maximum efficiency in both compute power and speed.
  • Task-specific scaling: Figure 9 summarizes scaling behavior for generative and language-modeling tasks.The figure belongs to the broader set of classification, reinforcement-learning, and generative-modeling comparisons.
  • Task-specific scaling: Figure 10 evaluates ImageNet scaling using error-based learning-rate decays at training errors 0.487, 0.312, and 0.229.The critical batch size is computed separately for each schedule span.
  • Task-specific scaling: Figures 11 and 12 show A2C scaling behavior across seven Atari tasks, while Figure 13 shows Dota 1v1 scaling trained to top-level professional performance.The Atari tasks include Beam Rider, Breakout, Pong, Space Invaders, Alien, Qbert, and Seaquest.
  • Task-specific scaling: Figure 14 presents scaling behavior for image-classification tasks.The figure is part of the study’s cross-domain comparison of batch-size efficiency.
  • Noise-scale dynamics: A sixteenfold decrease in batch size caused the simple noise scale to increase sixteenfold, then return near its original value after batch size and learning rate were reset.This behavior was observed between specified early-training intervals for SVHN and Billion Word.

C Temperature and the Noise Scale

The paper treats the ratio of learning rate to batch size as a training temperature and studies how the noise scale depends on that temperature. Empirically, the noise scale tracks inverse temperature and remains consistent across optimally tuned batch sizes at a fixed temperature and loss.

  • Temperature: The noise scale primarily depends on learning rate and batch size through their ratio, which the paper calls training temperature.For pure SGD in the small-batch regime, temperature is approximated by T ≈ ϵ/B.
  • Temperature: In equilibrium, the noise scale varies in proportion to inverse temperature.The paper motivates temperature using a toy quadratic-loss model and interprets it through the loss landscape.
  • Empirical scaling: Decaying the learning rate by a constant factor often increases the noise scale by roughly the same factor, including an ImageNet example.The paper presents Figure 15 as a direct investigation of the relation between simple noise scale and temperature.
  • Empirical scaling: Because well-tuned runs at different batch sizes share temperature, measured noise scale is consistent across those runs and depends on temperature and loss.This consistency supports using the noise scale to characterize temperature.
  • Implication: The noise scale prescribes an optimal batch size at a given training temperature rather than an optimal temperature schedule.The paper distinguishes batch-size selection from temperature scheduling.

A Toy Model for the Temperature

A locally quadratic toy model links the gradient noise scale to learning rate and batch size, predicting how this statistic changes during training. The model fits neural-network measurements even with Adam, while relying on effective-equilibrium assumptions.

  • Model: The toy model approximates the loss as locally quadratic and models per-example losses as shifted versions of the true loss.This yields a gradient covariance matrix Σ = HΣcH and average gradient G = Hθ.
  • Predictions: The noise scale is expected to increase when the learning rate decreases or the batch size increases.Scaling learning rate and batch size together is expected to leave the noise scale unchanged.
  • Empirical check: The toy analysis accurately models the dependence of noise scale on learning rate and batch size during neural-network training, including with Adam.The tests used SVHN and Billion Word training runs with specified initial learning-rate and batch-size settings.
  • Assumption: The analysis assumes that the optimizer has approached an effective equilibrium.The authors expect equilibration to take longer in low-curvature directions.
  • Training dynamics: Measured noise-scale quantities and the empirical critical batch size increase by at least an order of magnitude during training.This motivates studying dynamically scaled batch sizes, although predicted gains are relatively modest unless the critical batch size changes greatly.

D.1 Theory

The theory treats batch size as an exchange between optimization steps and processed examples, deriving an optimal schedule and a Pareto tradeoff between training cost and time. Adaptive batching helps only when the noise scale varies substantially.

  • D.1 Theory: A batch of size B requires δS = 1 + B optimizer steps and δE = BδS training examples to match one full-batch step’s progress.Here B is the noise scale.
  • D.1 Theory: The training trajectory is parameterized by the number s of full-batch optimizer steps.The minimum-step quantity Smin is abbreviated to s for the derivation.
  • D.1 Theory: Batch size choices trade processed examples for optimizer steps through an exchange rate that can vary along the training trajectory.The theory frames adaptive batching as optimally distributing examples across training.
  • D.1 Theory: An optimal batch-size schedule varies in proportion with the square root of the noise scale.This follows because transferring examples between training phases should not save optimization steps when the exchange rate is constant.
  • D.1 Theory: When the noise scale is constant, γ = 1 and adaptive batching provides no benefit; greater variation pushes γ toward 0 and improves the predicted Pareto front.The parameter γ summarizes noise-scale variation over training.
  • D.1 Theory: The exchange rate r controls the preferred balance between training time and compute, with r = Emin/Smin making both efficiencies within a factor 1 + √γ of optimal.This choice corresponds to the turning point in Figure 16.
  • D.1 Theory: Adaptive batch-size training shows a modest but visible efficiency improvement over fixed-batch Pareto fronts.The figure compares adaptive-batch data points with hyperbolic fits for fixed batch sizes.

D.2 An SVHN Case Study

The SVHN case study applies noise-scale-based adaptive batching during training and observes a modest efficiency benefit. The observed gains can exceed the theoretical prediction, particularly because adaptive batching reduces minimum optimization steps at large batch sizes.

  • D.2 An SVHN Case Study: Adaptive batch training dynamically sets the batch size using periodically measured Bsimple.Its performance is compared with fixed batch sizes across various exchange-rate values r.
  • D.2 An SVHN Case Study: Adaptive training produces a modest efficiency benefit compared with fixed batch-size training.The comparison uses the same learning-rate relation for fixed and adaptive runs, yielding an adaptive learning rate.
  • D.2 An SVHN Case Study: For SVHN, the critical batch size is fit by Bcrit(s) ≈ 10√s, where s is the number of steps in the very-large-batch limit.This fit is used to predict adaptive-training efficiency gains.
  • D.2 An SVHN Case Study: The theoretical prediction gives an efficiency gain of around 4%.The prediction is obtained by applying Equation D.5 to the fitted critical-batch-size behavior.
  • D.2 An SVHN Case Study: Observed benefits in some Figure 16 runs appear too large to be fully explained by the theoretical analysis.The authors identify this as an unresolved discrepancy.
  • D.2 An SVHN Case Study: Adaptive batching seems to reduce Smin in the large-batch regime, although the mechanism is unresolved.Possible factors include learning-rate warmup or a larger, more consistent proportion of gradient noise.
  • D.2 An SVHN Case Study: Adaptive batch sizes may maintain a constant proportion of gradient noise, which some prior work has argued could benefit generalization.This is presented as an additional possible advantage rather than an established result here.

E.1 Deterministic Training Performs Poorly

Nearly deterministic large-batch training can perform poorly because later line-search step sizes rapidly shrink, whereas fixed-learning-rate training often stabilizes at updates about twice the optimal size. Related scaling arguments connect optimizer behavior to batch-dependent gradient variance.

  • E.1 Deterministic Training Performs Poorly: At B ≫ Bnoise, training approaches a minimally stochastic procedure that line-searches the true loss along the true gradient.This motivates studying nearly deterministic or “greedient descent” updates.
  • E.1 Deterministic Training Performs Poorly: Nearly deterministic “greedient descent” performs poorly: after initially reducing loss, its step sizes rapidly shrink and make little further progress.The cited behavior is reported for this minimally stochastic training approach.
  • E.1 Deterministic Training Performs Poorly: With fixed learning rates, the optimal line-search step is often almost exactly half the actual update magnitude across several learning rates and model types.The paper reports this pattern for an LSTM and a CNN, among other models.
  • E.1 Deterministic Training Performs Poorly: The authors interpret the half-update pattern as consistent with large-Hessian-curvature directions dominating the update.This is offered as a natural interpretation of the observed behavior.
  • E.2 Motivations for Learning Rate Scaling Rules: For Adam, batch-dependent gradient variance motivates a square-root learning-rate scaling rule, while slow second-moment adaptation pushes the exponent toward 1.0.This may explain why the exponent varies between tasks.
  • E.3 Test-Loss Scaling: At the end of training, the test-loss critical batch size shows a small dip consistently across the small image-classification datasets.Early in training, train-loss and test-loss results do not differ.
Loading 1812.06162v1…