Measuring the Effects of Data Parallelism on Neural Network Training

Christopher J. Shallue; Jaehoon Lee; Joseph Antognini; Jascha Sohl-Dickstein; Roy Frostig; George E. Dahl

Source-linked AI summary

Measuring the Effects of Data Parallelism on Neural Network Training

Christopher J. Shallue, Jaehoon Lee, Joseph Antognini, Jascha Sohl-Dickstein, Roy Frostig, George E. Dahl

arXiv:1811.03600v3Nov 8, 2018cs.LGstat.ML

TL;DR

The paper asks how increasing batch size changes neural-network training cost and out-of-sample error as data-parallel hardware scales. It experimentally measures this relationship across workloads with batch-size-specific tuning, finding proportional initial step reductions, workload-dependent diminishing returns, and no necessary degradation in out-of-sample performance.

Problem
The paper addresses uncertainty about batch-size effects on training time and out-of-sample error across training algorithms, models, and data sets.
Method
The study measures steps to a goal out-of-sample error across seven data sets and six neural-network families while evaluating learning-rate schedules and other metaparameters.
Results
Across workloads, larger batches initially reduce training steps proportionally, but the useful scaling range varies from batch size 2^4 to 2^13; tuning and compute budgets explain many quality disagreements, with no evidence of necessary degradation.
Takeaways & Limitations
Larger batches can provide valuable speedups, but blindly scaling to hardware limits will not yield large speedups for every workload.
Takeaways & Limitations
Metaparameter search spaces involved human judgment, and steps-to-result measurements were sensitive to learning-rate schedules.

Abstract

from arXiv · show

Recent hardware developments have dramatically increased the scale of data parallelism available for neural network training. Among the simplest ways to harness next-generation hardware is to increase the batch size in standard mini-batch neural network training algorithms. In this work, we aim to experimentally characterize the effects of increasing the batch size on training time, as measured by the number of steps necessary to reach a goal out-of-sample error. We study how this relationship varies with the training algorithm, model, and data set, and find extremely large variation between workloads. Along the way, we show that disagreements in the literature on how batch size affects model quality can largely be explained by differences in metaparameter tuning and compute budgets at different batch sizes. We find no evidence that larger batch sizes degrade out-of-sample performance. Finally, we discuss the implications of our results on efforts to train neural networks much faster in the future. Our experimental data is publicly available as a database of 71,638,836 loss measurements taken over the course of training for 168,160 individual models across 35 workloads.

1. Introduction

The paper studies how data parallelism and batch size affect the training cost and out-of-sample error of neural networks, addressing practical uncertainty created by expanding hardware capacity and conflicting literature. Across workloads, it finds a shared scaling pattern but substantial variation in the maximum useful batch size, while metaparameter tuning helps explain disagreements about model quality.

Data parallelism distributes examples across processors to compute and aggregate gradient updates, making it applicable across neural network architectures.
The paper asks how batch size relates to steps to a goal out-of-sample error, what governs that relationship, and whether large batches harm out-of-sample error.
Across six network families, three training algorithms, and seven data sets, batch size has the same characteristic relationship with steps to a goal error.
Increasing batch size initially reduces required steps proportionally, followed by diminishing returns and eventual saturation across workloads and error goals.
The maximum useful batch size varies with workload: momentum methods scale further than plain SGD, while model effects exceed data-set effects and are not simply determined by width or size.
Popular learning-rate scaling heuristics do not hold across all problems or batch sizes, and tuning and compute-budget choices can explain conflicting claims about quality degradation.

2. Setup and Background

The paper formalizes learning tasks, training objectives, stochastic gradients, and metaparameters, then studies mini-batch SGD and its variants through data-parallel gradient computation. It measures training cost primarily by steps because step count is hardware-agnostic and directly supports comparisons of training time.

A learning task combines a data distribution, loss, and model, with the goal of finding parameters that achieve low out-of-sample loss.
The regularized training objective adds λR(θ) to empirical average loss, where R is a differentiable penalty and λ > 0.
Mini-batch SGD updates parameters using stochastic gradients computed from subsets of training examples; uniformly sampled batches give unbiased gradient estimates.
SGD variants such as momentum, Nesterov momentum, RMSProp, and Adam all repeatedly compute stochastic gradients and therefore share the same batch-size notion.
Training metaparameters include architectural choices, learning rates, and regularization weights, selected by validation performance in the experiments.

3. Related Work

Prior work offers analytical bounds and empirical evidence for both benefits and costs of larger batches, but its conclusions depend on assumptions, workloads, budgets, and tuning procedures. The paper motivates a direct experimental comparison that independently evaluates batch sizes against a common error goal.

Analytical Studies: Analytical studies describe batch-size-dependent convergence bounds, but often rely on convexity and regularity assumptions violated by neural network training.
Analytical Studies: Theoretical bounds span no benefit from larger batches to a b-fold reduction in steps, leaving the practical relationship unresolved.
Analytical Studies: Previous critical-batch-size analyses are parameter-dependent or specific to linear least-squares regression, limiting their generalization to nonlinear neural-network optimization.
Empirical Studies: Earlier empirical studies did not consistently retune each batch size to minimize steps to the same goal error, making their speed comparisons incomplete.
Solution Quality: The literature disagrees about whether large batches worsen solution quality, with claims involving sharp minima, sufficient training duration, and possible degradation only at still larger sizes.
Solution Quality: Comparisons are constrained by wall time, hardware, metaparameter-tuning limits, and differing step, epoch, or convergence budgets.

4. Experiments and Results

Across diverse workloads, the experiments measure steps to result while independently tuning training metaparameters at each batch size. Increasing batch size initially reduces required steps proportionally, then yields diminishing returns and eventually no further benefit; model architecture and model size affect how large a batch can be exploited.

Experimental setup: The study spans seven data sets, six neural-network families, and realistic workloads, with steps to result defined as reaching a desired out-of-sample error.Models include fully connected networks, CNNs, ResNets, VGG-11, Transformers, and LSTMs across image and text tasks.
Experimental setup: Learning rate, momentum, and applicable decay schedules were independently tuned for each batch size.The linear-decay schedule uses η0, α, and T, with the learning rate decaying from η0 to ηT = αη0.
Batch-size scaling: For every workload, larger batches initially halve the required steps for each doubling, followed by diminishing returns and a maximal-parallelism region with no further step reduction.This characteristic relationship was observed across six neural-network families, three training algorithms, and seven data sets.
Batch-size scaling: Nearby validation-error goals produce similar steps-to-result curves, with nearly unchanged transition points between perfect scaling, diminishing returns, and maximal data parallelism.Changing the goal mainly produces vertical shifts for modest variations around a good absolute validation error.
Model dependence: Model architecture significantly affects the useful batch-size range, while changing depth or width affects it inconsistently across architectures.On ImageNet, ResNet-50 exploits larger batches than VGG-11; on MNIST, deeper fully connected models appear better able to exploit larger batches than a shallower model.

4.4 Momentum Extends Perfect Scaling to Larger Batch Sizes, but Matches Plain SGD at Small Batch Sizes

Momentum optimizers extend the batch-size range over which training steps scale ideally, while matching plain SGD at small batch sizes. The maximum useful batch size also varies across data sets, and regularization can matter more at larger batches.

Momentum and scaling: Momentum optimizers extend the perfect scaling region and reduce training steps at larger batch sizes, but match SGD when batches are small.Both Nesterov momentum and SGD with momentum show this pattern.
Data-set effects: Data set changes affect the batch-size relationship, but larger data sets do not consistently make larger batches more valuable.Fashion MNIST, ImageNet, Open Images, LM1B, and Common Crawl exhibit differing scaling behavior.
Data-set effects: Subset experiments find little or no subset-size effect on MNIST, while ImageNet subsets can deviate from perfect scaling at different batch sizes.The half-image subset deviates sooner, whereas the half-class subset remains close to the full-data curve.
Regularization: Label smoothing improved ResNet-50 validation error by as much as one percentage point at large batch sizes, with no apparent effect at small batch sizes.Without label smoothing, the goal error rate of 0.25 was unattainable above batch size 214 within the training budget.
Regularization: Label smoothing reduced overfitting at large batches, while increasing L2 weight penalties or adding Gaussian gradient noise did not initially control it effectively.On MNIST and Fashion MNIST, label smoothing generally helped across batch sizes without a consistent size-dependent trend.

4.7 The Best Learning Rate and Momentum Vary with Batch Size

The batch size changes which learning-rate and momentum settings train rapidly, and the fastest settings can lie near divergence boundaries. Smaller models also tolerate larger stable learning rates, while validation error depends strongly on the compute budget used for comparison.

Learning-rate scaling: The optimal effective learning rate generally increases with batch size, but does not consistently follow linear or square-root scaling.In some cases, it decreases at larger batch sizes.
Training-speed regions: Rapid training in epochs occupies a smaller metaparameter region at larger batch sizes, whereas rapid training by step-count occupies a larger region before eventually plateauing.The analysis examined constant learning rate and Nesterov momentum for Transformer on LM1B.
Feasibility and stability: The fastest metaparameters typically lie on the feasible-region boundary, so small changes can make training diverge.This boundary behavior may challenge optimization methods that assume smooth performance relationships.
Model size: Smaller models have larger stable learning rates, consistent with predictions that the largest stable rate is inversely proportional to layer width.This relationship is illustrated for Transformer models on LM1B with batch size 1024 and a fixed validation-error goal.
Compute-budget dependence: Validation error depends more on compute budget than batch size: step budgets favor large batches, whereas epoch budgets favor small batches.The comparison selected the best validation-error checkpoint across metaparameter trials within each budget.

5. Discussion

Increasing batch size initially reduces required training steps proportionally, but scaling limits vary dramatically across workloads and can shift with the optimizer. The experiments also indicate that apparent quality degradations often reflect metaparameter tuning and compute-budget choices rather than batch size itself.

Scaling behavior: Perfect scaling ended between batch sizes 24 and 2^13, after which workloads entered diminishing returns and eventually reached a maximum useful batch size.The maximum useful batch size was extremely problem-dependent.
Experimental interpretation: Comparisons that tune learning rate only at one batch size systematically favor that batch size, while optimal metaparameters need not scale simply with batch size.Epoch budgets favor smaller batches by allowing more optimizer steps, whereas step budgets can favor larger batches by exposing them to more examples.
Optimizer effects: Nesterov momentum extended perfect scaling for ResNet-8 on CIFAR-10 from batch size 16 with plain mini-batch SGD to batch size 256.This comparison shows that optimizer choice can substantially change scaling behavior for a fixed model and data set.
Solution quality: A larger batch size never increased the number of steps needed to reach a goal validation error in the experiments.Thus, a step budget was likely sufficient for larger batches to reach at least the same performance as smaller batches.
Solution quality: Increasing batch size did not inherently degrade solution quality, although additional regularization may become important at larger batch sizes.The authors attribute many disagreements in the literature to differences in computational budgets and metaparameter-selection procedures.
Limitations: The study’s protocol leaves uncertainty because search spaces may be uneven across batch sizes, the globally optimal metaparameters were not guaranteed, and robustness across trials was not assessed.The analysis reported optimal trials rather than distinguishing consistently successful batch sizes from those where only one trial reached the goal error.

6. Conclusions and Future Work

Larger batches provide useful speedups across workloads, but their benefits diminish before reaching current hardware limits. The authors propose extending scaling through optimizer and model-design research.

Batch-size increases provide valuable speedups across a range of workloads, but benefits diminish well within current hardware limits.
Blindly increasing batch size to the hardware limit will not produce large speedups for every workload.
Some optimization algorithms may consistently extend perfect scaling across many models and data sets.
Future work should test additional optimizers and uncover general principles for models that scale perfectly to larger batch sizes.

Appendix A. Data Set Details

This section provides details of the data sets summarized in Table 1.

The section contains data-set details summarized in Table 1.
Data-set descriptions are presented as supporting material for the study.
Table 1 summarizes the data sets whose details are provided in this section.

A.1 Data Set Descriptions and Pre-Processing

The study covers image, text, and web-scale data sets with task-specific preprocessing and train-validation-test splits. These include standard benchmarks and very large corpora.

Image data sets: MNIST and Fashion MNIST use 55,000 training images, 5,000 validation images, and 10,000 test images without data augmentation.
Image data sets: CIFAR-10 contains 32×32 color images across 10 classes, with 45,000 training images, 5,000 validation images, and standardized inputs.
Image data sets: ImageNet contains 1,000 classes, with 1,281,167 training images, 50,045 test images, and 50,000 validation images; training uses augmentation.
Image data sets: Open Images v4 contains 9 million images and uses 7,186 trainable classes with at least 100 human-annotated positives each.
Text data sets: LM1B uses English news articles, a 32,000-token sub-word vocabulary, and 30,301,028 training sentences with separate validation and test sets.
Text data sets: Common Crawl processing yields a 24,006-token vocabulary and approximately 25.8 billion training sentences after filtering.

A.2 Evaluation Metrics

Evaluation uses classification error, average precision, and average per-token cross entropy error across the study’s image and language workloads.

Image metrics: Classification error is used for MNIST, Fashion MNIST, CIFAR-10, and ImageNet.
Image metrics: Open Images is evaluated with class-agnostic average precision over all classes in each image.
Image metrics: Evaluating average precision over all classes may underestimate true AP because validation labels contain false negatives.
Language metrics: LM1B and Common Crawl use average per-token cross entropy error, defined as the logarithm of per-token perplexity.

Appendix B. Model Details

The appendix describes the architectures and output-layer choices used across the paper’s neural-network workloads, including feedforward, convolutional, recurrent, attention-based, and residual models.

Classification models use softmax outputs for mutually exclusive labels, while Open Images models use sigmoid outputs for multilabel classification.
Fully Connected networks use ReLU activations and dropout probability 0.4, with experiments varying hidden-layer depth and width.
Simple CNN uses two convolutional layers with max-pooling followed by one fully connected hidden layer.
ResNet-8 and ResNet-50 use residual convolutional architectures, with ResNet-50 replacing batch normalization by ghost batch normalization.
The appendix also specifies VGG-11, a one-hidden-layer LSTM, and Transformer variants used for the experiments.

Appendix C. Learning Rate Schedules

The appendix defines learning-rate schedules through decay functions and schedule metaparameters, then reports schedule experiments for ResNet-50 on ImageNet.

The schedule uses η0 as the initial learning rate, α as the decay multiplier, and T as the step until decay from η0 to αη0.
The authors experimented with cosine and piecewise linear learning-rate schedules, among other decay functions.
Piecewise linear schedules can express both the piecewise-constant schedule of He et al. and the warm-up schedule of Goyal et al.
For ResNet-50 on ImageNet, experiments used Nesterov momentum, batch size 1,024, and 150,000 training steps while tuning momentum and learning-rate schedule metaparameters.
The experiments found good results for several tested schedules and suggested that other schedules might also work well.

Appendix D. Additional Plots

The additional plots document evaluation consistency, metaparameter-search validation, normalized and zoomed comparisons, label-smoothing effects, and batch-size relationships for optimal hyperparameters.

Training-set steps to a result are almost the same as validation-set steps, with error goals specified in plot legends.
Search-space validation plots show tuned metaparameters by batch size, with blue dots for trials reaching the goal and yellow stars for fewest-step trials.
Figures 3, 5, and 6 are reproduced without y-axis normalization, while Figures 11a and 11b also receive zoomed versions.
Label smoothing reduces overfitting at batch size 216 for ResNet-50 on ImageNet and helps all batch sizes for Simple CNN on MNIST and Fashion MNIST.
Validation-error projections examine effective learning rate across metaparameter searches, while optimal learning rates and momentum are compared with batch size.