The Impact of the Mini-batch Size on the Variance of Gradients in Stochastic Gradient Descent

TL;DR

This paper analyzes how mini-batch size affects gradient variance in stochastic gradient descent, providing theoretical insights and empirical evidence that smaller batches lead to lower gradient variance and better training outcomes.

Contribution

It offers the first theoretical analysis of gradient variance dependence on mini-batch size in linear and deep linear networks, supported by empirical validation.

Findings

Gradient norm decreases with larger mini-batch size

Gradient variance is a decreasing function of batch size

Smaller batches tend to produce lower loss values

Abstract

The mini-batch stochastic gradient descent (SGD) algorithm is widely used in training machine learning models, in particular deep learning models. We study SGD dynamics under linear regression and two-layer linear networks, with an easy extension to deeper linear networks, by focusing on the variance of the gradients, which is the first study of this nature. In the linear regression case, we show that in each iteration the norm of the gradient is a decreasing function of the mini-batch size $b$ and thus the variance of the stochastic gradient estimator is a decreasing function of $b$. For deep neural networks with $L_2$ loss we show that the variance of the gradient is a polynomial in $1/b$. The results back the important intuition that smaller batch sizes yield lower loss function values which is a common believe among the researchers. The proof techniques exhibit a relationship between stochastic gradient estimators and initial weights, which is useful for further research on the dynamics of SGD. We empirically provide further insights to our results on various datasets and commonly used deep network structures.