Online Learning and Information Exponents: On The Importance of Batch size, and Time/Complexity Tradeoffs

TL;DR

This paper investigates how batch size affects training time in neural networks, identifying optimal batch sizes based on target complexity, and introduces a new protocol to surpass existing limitations, supported by theoretical and experimental validation.

Contribution

It characterizes the optimal batch size for minimizing training time based on information exponents and proposes Correlation loss SGD to improve time complexity beyond traditional limits.

Findings

Optimal batch size scales with input dimension and target complexity.

Large batch sizes beyond a threshold hinder training time improvements.

Correlation loss SGD effectively reduces auto-correlation, enhancing training efficiency.

Abstract

We study the impact of the batch size $n_b$ on the iteration time $T$ of training two-layer neural networks with one-pass stochastic gradient descent (SGD) on multi-index target functions of isotropic covariates. We characterize the optimal batch size minimizing the iteration time as a function of the hardness of the target, as characterized by the information exponents. We show that performing gradient updates with large batches $n_b \lesssim d^{\frac{\ell}{2}}$ minimizes the training time without changing the total sample complexity, where $\ell$ is the information exponent of the target to be learned \citep{arous2021online} and $d$ is the input dimension. However, larger batch sizes than $n_b \gg d^{\frac{\ell}{2}}$ are detrimental for improving the time complexity of SGD. We provably overcome this fundamental limitation via a different training protocol, \textit{Correlation loss SGD}, which suppresses the auto-correlation terms in the loss function. We show that one can track the training progress by a system of low-dimensional ordinary differential equations (ODEs). Finally, we validate our theoretical results with numerical experiments.