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 stochastic stability through the \emph{critical learning rate}, defined as the largest learning rate for stable training, and obtain ηSGDcriteqsim 1, ηPolyakcriteqsim min{1,B(1−β)}, and ηNesterovcriteqsim min{1,Bq(1−β)}, where B is the batch size, β∈[0,1) is the momentum factor, and q>1 is the capacity exponent governing spectral decay. Within the stable 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 optimize the final-step risk under a fixed data budget, obtaining a three-regime phase diagram that reveals how the benefits of momentum change with batch size. Notably, Polyak enlarges the critical batch size, the largest batch size preserving the optimal small-batch data-scaling exponent, thereby allowing fewer sequential updates without sacrificing data efficiency. In contrast, Nesterov improves the optimal data-scaling exponent in the large-batch regime through look-ahead-induced suppression of noise amplification. Experiments on PLK regression quantitatively support the theoretical scaling predictions, while neural-network experiments provide qualitative evidence for the predicted batch-size dependence and stability behavior.
Related
- Continuized Nesterov Momentum Achieves the O(ε−7/4) Complexity in Smooth Nonconvex Optimization
- Muon with Nesterov Momentum: Heavy-Tailed Noise and (Randomized) Inexact Polar Decomposition
- When Are Nesterov Trajectories Eventually Locally Minimizing?
- Improved Gradient Descent Lower Bounds Beyond Nesterov
- On the Nonoccurrence of Gradient Restart in Nesterov Accelerated Gradient Method
BibTeX
@misc{wang2026momentum,
title = {{Momentum in large-batch training: Polyak enlarges the critical batch size, Nesterov improves data efficiency}},
author = {Jia-Nan Wang and Zixun Huang and Kairui Li and Zilin Wang and Lei Wu},
howpublished = {arXiv},
year = {2026},
month = oct,
doi = {10.48550/arxiv.2609.02728},
eprint = {2609.02728},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2609.02728},
note = {CC BY 4.0},
}
RIS
TY - UNPB
TI - Momentum in large-batch training: Polyak enlarges the critical batch size, Nesterov improves data efficiency
AU - Jia-Nan Wang
AU - Zixun Huang
AU - Kairui Li
AU - Zilin Wang
AU - Lei Wu
PY - 2026
DA - 2026/10/27
PB - arXiv
DO - 10.48550/arxiv.2609.02728
UR - https://arxiv.org/abs/2609.02728
AB - 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 stochastic stability through the \emph{critical learning rate}, defined as the largest learning rate for stable training, and obtain η_SGD^criteqsim 1, η_Polyak^criteqsim min{1,B(1−β)}, and η_Nesterov^criteqsim min{1,B^q(1−β)}, where B is the batch size, β∈[0,1) is the momentum factor, and q>1 is the capacity exponent governing spectral decay. Within the stable 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 optimize the final-step risk under a fixed data budget, obtaining a three-regime phase diagram that reveals how the benefits of momentum change with batch size. Notably, Polyak enlarges the critical batch size, the largest batch size preserving the optimal small-batch data-scaling exponent, thereby allowing fewer sequential updates without sacrificing data efficiency. In contrast, Nesterov improves the optimal data-scaling exponent in the large-batch regime through look-ahead-induced suppression of noise amplification. Experiments on PLK regression quantitatively support the theoretical scaling predictions, while neural-network experiments provide qualitative evidence for the predicted batch-size dependence and stability behavior.
ER -