[Paper Review] Revisiting Small Batch Training for Deep Neural Networks

Dominic Masters, Carlo Luschi|arXiv (Cornell University)|Apr 20, 2018

Advanced Neural Network Applications23 references353 citations

TL;DR

The paper shows small mini-batch sizes (2–32) yield better generalization and stable training across CIFAR-10/100 and ImageNet, challenging the push toward very large batches. It analyzes learning rate scaling, batch normalization effects, and warm-up strategies.

ABSTRACT

Modern deep neural network training is typically based on mini-batch stochastic gradient optimization. While the use of large mini-batches increases the available computational parallelism, small batch training has been shown to provide improved generalization performance and allows a significantly smaller memory footprint, which might also be exploited to improve machine throughput. In this paper, we review common assumptions on learning rate scaling and training duration, as a basis for an experimental comparison of test performance for different mini-batch sizes. We adopt a learning rate that corresponds to a constant average weight update per gradient calculation (i.e., per unit cost of computation), and point out that this results in a variance of the weight updates that increases linearly with the mini-batch size $m$. The collected experimental results for the CIFAR-10, CIFAR-100 and ImageNet datasets show that increasing the mini-batch size progressively reduces the range of learning rates that provide stable convergence and acceptable test performance. On the other hand, small mini-batch sizes provide more up-to-date gradient calculations, which yields more stable and reliable training. The best performance has been consistently obtained for mini-batch sizes between $m = 2$ and $m = 32$, which contrasts with recent work advocating the use of mini-batch sizes in the thousands.

Motivation & Objective

• Evaluate how mini-batch size affects generalization and convergence in deep nets.
• Investigate learning rate scaling under a constant weight-update-per-gradient-cost formulation.
• Assess Batch Normalization’s interaction with different batch sizes and its impact on training stability.
• Examine the role of gradual warm-up in large-batch versus small-batch regimes.
• Provide guidance for distributed training where BN and SGD batch sizes may differ.

Proposed method

• Formulate SGD updates with either the average or sum of mini-batch gradients to compare learning rate scaling.
• Empirically train CNNs on CIFAR-10, CIFAR-100, and ImageNet with varying batch sizes and base learning rates across multiple architectures.
• Evaluate performance with and without Batch Normalization and with data augmentation.
• Test gradual warm-up strategies to mitigate large-batch training instability.
• Analyze the effect of using different batch sizes for BN versus SGD updates.

Experimental results

Research questions

• RQ1 How does mini-batch size influence generalization performance and stable convergence under a constant per-gradient update cost?
• RQ2 What is the impact of Batch Normalization on small versus large batch sizes on training dynamics and final accuracy?
• RQ3 Does a gradual warm-up strategy mitigate the degradation observed with large batch sizes in diverse datasets?
• RQ4 What are optimal batch size regimes (2–32) across CIFAR-10/100 and ImageNet for best accuracy and stability?
• RQ5 How does separating BN batch size from SGD batch size affect distributed training performance?

Key findings

• Best generalization consistently occurs for mini-batch sizes between 2 and 32 across tested architectures and datasets.
• Increasing batch size reduces the range of learning rates that yield stable convergence and acceptable test performance.
• Batch Normalization improves convergence and test accuracy and allows effective training with moderate batch sizes; very small BN batches can be problematic for fully connected layers.
• Gradual warm-up helps maintain stable training in larger batch settings but does not fully restore small-batch performance advantages.
• ImageNet results show best validation accuracy for batch sizes between 16 and 64, with larger batches more sensitive to learning rate choices.