Momentum in large-batch training: Polyak enlarges the critical batch size, Nesterov improves data efficiency

Abstract: We study when and how momentum improves large-batch training in the one-pass regime, using power-law kernel regression as a tractable setting. We first characterize risk stability through the critical learning rate, defined as the largest learning rate for stable training, and obtain $$η<em>{\mathrm{SGD}}^{{\mathrm{crit}}\eqsim} 1$$, $$η</em>{\mathrm{Polyak}}^{{\mathrm{crit}}\eqsim} \min{1,B(1-ρ)}$$, and $$η_{\mathrm{Nesterov}}^{{\mathrm{crit}}\eqsim} \min{1,B<sup>β(1-ρ)}$$, where $B$ is the batch size, $ρ$ is the momentum factor, and $β&gt;1$ is the capacity exponent. Within this admissible region, we derive scaling laws for the full risk dynamics, capturing the progression from an early transient, through power-law decay, to a noise floor. We then minimize the final-step risk over the admissible learning rates and momentum factors under a fixed data budget, yielding a three-regime batch-size phase diagram that reveals how the role of momentum changes with batch size. Notably, Polyak enlarges the critical batch size, the largest batch size preserving the best small-batch data-scaling exponent, thereby enabling greater parallelism without sacrificing data efficiency. In contrast, Nesterov achieves better data efficiency in the large-batch regime because its look-ahead mechanism suppresses noise accumulation. Numerical experiments validate the predicted stability boundaries, risk dynamics, and batch-size phase diagram.