Blog

The latest from Google Research

Measuring the Limits of Data Parallel Training for Neural Networks

Tuesday, March 19, 2019

Posted by Chris Shallue, Senior Software Engineer and George Dahl, Senior Research Scientist, Google AIimage classificationmachine translationspeech recognitionCloud TPU Podsdistribute computationsmodel parallelismdata parallelismstochastic gradient descentbatch sizeMeasuring the Effects of Data Parallelism in Neural Network Trainingshare our raw data our paperUniversal Relationship Between Batch Size and Training Timestepsout-of-sample error

For all workloads we tested, we observed a universal relationship between batch size and training speed with three distinct regimes: perfect scaling (following the dashed line), diminishing returns (diverging from the dashed line), and maximal data parallelism (where the trend plateaus). The transition points between the regimes vary dramatically between different workloads.

Cloud TPU PodsResNet-8CIFAR-10ResNet-50ImageNetOptimizing Workloads

Left: A transformer neural network scales to much larger batch sizes than an LSTM neural network on the LM1B dataset. Right: The Common Crawl dataset does not benefit from larger batch sizes than the LM1B dataset, even though it is 1,000 times the size.

momentumFuture WorkAcknowledgementsThe authors of this study were Chris Shallue, Jaehoon Lee, Joe Antognini, Jascha Sohl-Dickstein, Roy Frostig and George Dahl (Chris and Jaehoon contributed equally). Many researchers have done work in this area that we have built on, so please see our paper for a full discussion of related work.