Momentum in large-batch training: Polyak enlarges the critical batch size, Nesterov improves data efficiency
stat.ML, cs.LG, math.OC
Submitted: 2026-09-02
Updated: 2026-09-02
Comments: 69 pages, 8 figures
License: http://creativecommons.org/licenses/by/4.0/
The gist: We study when and how momentum improves large-batch training in the one-pass regime, using power-law kernel regression as a tractable setting.
Terminology
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 η SGD crit 1, η Polyak crit 1,B(1-ρ), and η Nesterov crit 1,B β(1-ρ), where B is the batch size, ρ is the momentum factor, and β>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.
Sources
- DeepSeek-V4: Towards Highly Efficient Million-Token Context Intelligence
- Scaling Law for Stochastic Gradient Descent in Quadratically Parameterized Linear Regression
- Logarithmic-time Schedules for Scaling Language Models with Momentum
- The Llama 3 Herd of Models
- Deep Learning Scaling is Predictable, Empirically
- Learning Curve Theory
- Scaling Laws for Neural Language Models
- Optimal Learning Rate Schedules under Functional Scaling Laws: Power Decay and Warmup-Stable-Decay
- Accelerating SGD with momentum for over-parameterized learning
- Muon is Scalable for LLM Training
- A Solvable Model of Neural Scaling Laws
- Compute Efficiency and Serial Runtime Tradeoffs for Stochastic Momentum Methods
- 4+3 Phases of Compute-Optimal Neural Scaling Laws
- Benchmarking Optimizers for Large Language Model Pretraining
- Practical Efficiency of Muon for Pretraining
- Asymmetric Scaling Laws from Sparse Features
- Convergence of Momentum-Based Optimization Algorithms with Time-Varying Parameters
- The Sharpness Disparity Principle in Transformers for Accelerating Language Model Pre-Training
- A Theoretical Analysis of Noise Geometry in Stochastic Gradient Descent
- The Optimality of (Accelerated) SGD for High-Dimensional Quadratic Optimization
Related papers
- Behavior of prediction performance metrics with rare events
- Optimal Estimation of Generic Dynamics by Path-Dependent Neural Jump ODEs
- A Posterior-Dynamics Framework for Imaging Inverse Problems with Pretrained Diffusion Priors
- Efficient quadratic entropy with distance sketches
- Embedding-Bias in Conditional Independence Testing
- From Log-Odds to Shapley Values: An Explanatory Geometry for the Weighted Naive Bayes Classifier