PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Jia-Nan Wang, Zixun Huang, Kairui Li, Lei Wu

arXiv:2609.02728v1stat.MLcs.LGmath.OC

TL;DR

The paper studies how momentum affects stability and data efficiency in one-pass large-batch training, where larger batches reduce the number of sequential updates. It analyzes Polyak and Nesterov momentum in power-law kernel regression, deriving risk dynamics and optimizing final-step risk across admissible hyperparameters. Polyak preserves small-batch data scaling over a wider batch range, while Nesterov improves data efficiency in the large-batch regime.

  • Problem

    The paper asks whether momentum can enable larger batches without sacrificing data efficiency or improve data efficiency itself under a fixed data budget.

  • Method

    The paper uses power-law kernel regression to analyze stability, full risk dynamics, and optimized data scaling for Polyak and Nesterov momentum under one-pass constraints.

  • Results

    Polyak preserves the best small-batch data-scaling exponent over a substantially wider batch-size range, whereas Nesterov further improves the exponent in the large-batch regime.

  • Takeaways & Limitations

    Momentum provides distinct batch-size-dependent benefits: Polyak enlarges the range preserving small-batch scaling, while Nesterov improves large-batch data efficiency.

  • Takeaways & Limitations

    The analysis treats momentum, learning rate, and batch size as fixed within each training run, so a complete theory of time-varying schedules is beyond its scope.

Abstract

from arXiv · show

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 $η_{\mathrm{SGD}}^{\mathrm{crit}}\eqsim 1$, $η_{\mathrm{Polyak}}^{\mathrm{crit}}\eqsim \min\{1,B(1-ρ)\}$, and $η_{\mathrm{Nesterov}}^{\mathrm{crit}}\eqsim \min\{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.

1 Introduction

The paper asks whether momentum can preserve or improve data efficiency as batch sizes grow in one-pass training. Using power-law kernel regression, it analyzes stability, risk dynamics, and optimized batch-size scaling for Polyak and Nesterov momentum.

  • Motivation: Momentum’s role changes with batch size because larger batches improve throughput but reduce the number of optimization steps under a fixed data budget.This creates a tradeoff between hardware efficiency and data efficiency.
  • Method: Power-law kernel regression provides a tractable setting for analyzing momentum across spectral directions and learning timescales.The analysis uses capacity exponent β > 1 and source exponent s > 0 to characterize the problem.
  • Method: The analysis proceeds from identifying stable learning-rate and momentum regions to deriving risk scaling laws and optimizing final-step risk under a one-pass constraint.The admissible set contains learning rates η and momentum factors ρ at batch size B, while T = D/B updates are available.
  • Risk dynamics: In the high-momentum regime, risk moves from an early exponential transient through power-law decay and eventually reaches a noise floor.Nesterov retains the accelerated signal term while reducing the noise floor through look-ahead damping.
  • Data scaling: The optimized batch-size phase diagram has three regimes: equal exponents at small batches, Nesterov over Polyak over SGD at intermediate batches, and equal momentum exponents above the largest threshold.Polyak preserves the best small-batch scaling over a wider range, while Nesterov improves data scaling in the large-batch regime; experiments validate the predictions.
  • Stability: Polyak preserves the SGD signal–noise tradeoff while enlarging the admissible effective learning-rate range with batch size.Its stability compensation scales as B, whereas Nesterov’s scales as B^β.

2 Related work

Prior theory studies momentum acceleration, stochastic behavior, scaling laws, and stability, but does not fully characterize batch-size-dependent momentum under a fixed data budget. This paper distinguishes itself by deriving sharp full-risk dynamics and explicit batch-dependent data-scaling exponents.

  • Momentum theory: Classical momentum theory emphasizes accelerated convergence, while stochastic analyses examine convergence, stability, and asymptotic behavior.The related literature includes Lyapunov, continuous-time, and stochastic-gradient perspectives.
  • Scaling-law theory: Scaling-law theory connects learning performance to data, model, and compute resources across linear, kernel, and related tractable models.The paper places its analysis within this broader theoretical scaling-law literature.
  • Stability: Stability theory characterizes admissible hyperparameter regions and connects them to implicit bias, generalization, and optimization geometry.The related work also studies stability specifically for momentum methods.
  • Positioning: Compared with concurrent finite-dimensional linear-regression work, this paper derives sharp full-risk scaling laws and explicit data-scaling exponents as functions of batch size.The concurrent work leaves matching upper bounds open, whereas this paper presents a complete batch-size phase diagram.

3 Preliminaries

The paper analyzes SGD, Polyak, and Nesterov momentum in power-law kernel regression, using spectral damping regimes to study their training dynamics. The model is governed by capacity and source exponents, with stochastic gradients decomposed into population signal and noise.

  • Power-law kernel regression: Power-law kernel regression uses a Gaussian feature map with covariance operator H and excess risk defined by the target error measured in H.The analysis works in the covariance eigenbasis because all three methods are orthogonally equivariant.
  • Power-law kernel regression: β > 1 controls feature-spectrum decay, while s > 0 controls target-energy decay across spectral directions.Smaller s corresponds to more target energy in low-curvature directions and a harder learning problem.
  • Optimization methods: Polyak momentum is the heavy-ball method, whereas Nesterov momentum evaluates the gradient at a look-ahead point extrapolated along the previous update direction.The look-ahead mechanism distinguishes Nesterov from Polyak.
  • Stochastic gradients: Mini-batch gradients decompose into the population gradient plus stochastic noise, whose covariance obeys a curvature-aligned structure.The noise contains iterate-dependent multiplicative noise and label-noise-driven additive noise.
  • Regime-dependent damping: For momentum on a quadratic, overdamped directions decay according to the momentum-rescaled rate η/(1 −ρ), while underdamped directions oscillate with amplitude envelope ρ^k/2.The transition between regimes occurs at ηλ = (1 −√ρ)^2.
  • Regime-dependent damping: Nesterov damps sharper directions more strongly than Polyak, while sufficiently flat directions retain approximately Polyak-like dynamics.This curvature dependence arises from its look-ahead update.

4 Risk stability and critical learning rates

Risk stability is defined through the largest learning rate compatible with stable population-risk dynamics. The critical rate couples learning rate, batch size, and momentum, with Nesterov gaining a stronger batch-size dependence than Polyak.

  • Risk stability: Risk stability requires stable population-risk dynamics for every admissible initialization with finite population risk.The paper measures stability directly through the population risk.
  • Critical learning rate: ηcrit is the largest learning rate compatible with risk stability for fixed additional hyperparameters.For SGD the additional hyperparameter is B; for momentum methods they are (B, ρ).
  • Joint dependence: Increasing ρ toward one reduces ηcrit, whereas increasing B enlarges it.The theorem gives the joint dependence of stability on learning rate, batch size, and momentum.
  • Joint dependence: Nesterov’s stability compensation scales as B^β(1 −ρ), compared with B(1 −ρ) for Polyak before order-one ceilings.Because β > 1, Nesterov permits substantially higher momentum in large-batch training.
  • Noise and stability: Multiplicative noise feeds current risk back through a renewal kernel, while additive label noise acts as external forcing.Stability is controlled by feedback strength, characterized by ||K||_1 < 1.

5 Scaling laws for risk dynamics across damping regimes

Within the stable region, momentum risk follows damping-dependent scaling laws that progress from transients to power-law signal decay and noise floors. Polyak preserves the SGD tradeoff at a rescaled rate, while Nesterov suppresses large-rate noise accumulation.

  • Damping regimes: Momentum dynamics progress from fully overdamped SGD-like behavior to mixed damping and eventually risk instability as ρ increases.The analysis tracks how increasing momentum changes damping and stability jointly.
  • Damping regimes: The spectrum separates into flat overdamped directions and sharp underdamped directions at a curvature cutoff.Because the PLK spectrum has arbitrarily small eigenvalues, overdamped directions always remain present.
  • Polyak scaling: Polyak’s underdamped transient P(k) satisfies 0 ≤ P(k) ≲ ρ^k, and instability occurs when 1 −ρ ≲ η/B.The transient originates from sharp underdamped modes.
  • Polyak scaling: After the transient, Polyak reduces to the SGD law with η replaced by ηρ, preserving the signal–noise tradeoff while enlarging the admissible effective rate.Polyak permits ηρ ≲ B, whereas SGD requires η ≲ 1.
  • Nesterov scaling: Nesterov’s transient satisfies 0 ≤ N(k) ≲ min{1, (ηk)^−s}ρ^k, and instability occurs when 1 −ρ ≲ η/B^β.In the fully overdamped regime, Nesterov and Polyak share the same scaling law.
  • Nesterov scaling: Nesterov’s look-ahead damping suppresses noise accumulation at large ηρ while leaving the asymptotic signal term (ηρk)^−s unchanged.This reduction in noise accumulation is identified as the key effect for asymptotic data scaling.

6 Optimal data scaling across batch-size regimes

The paper optimizes final-step risk under a one-pass data budget and stability constraints, revealing three batch-size regimes. Polyak preserves the small-batch scaling over a wider range, while Nesterov leads in the large-batch regime.

  • Setup: The optimization uses T = D/B sequential updates and constrains learning rate and momentum to the admissible region.Batch size is allowed to grow with data budget D.
  • Setup: The batch-size exponent b is defined by B = D^b+o(1), and larger r(b) indicates better data efficiency.The analysis compares r_sgd(b), r_polyak(b), and r_nesterov(b).
  • Small batch: For b ≤ b1, SGD, Polyak, and Nesterov achieve the same optimal data-scaling exponent.This is the small-batch regime.
  • Large batch: For b1 < b < b3, the ordering is r_nesterov(b) > r_polyak(b) > r_sgd(b), with Nesterov peaking at b = b2.This is the regime where the three methods separate.
  • Ultra-large batch: For b ≥ b3, Polyak and Nesterov share an exponent that exceeds SGD’s exponent.The methods become equal in data-scaling exponent at ultra-large batch sizes.
  • Implications: Polyak enlarges the range preserving best small-batch scaling, whereas Nesterov improves the achievable exponent in the large-batch regime.The distinction reflects Polyak’s enlarged admissible range and Nesterov’s stronger noise suppression.

7 Numerical validations

Numerical experiments validate the predicted stability boundaries, risk-dynamics scaling laws, and three-regime batch-size phase diagram. They recover the predicted method ordering across small, large, and ultra-large batch regimes.

  • Risk-dynamics scaling laws: The risk dynamics exhibit an underdamped transient, power-law decay, and a noise floor, matching theoretical predictions.The transition timescale has dominant (1−ρ)^-1 dependence, and Polyak and Nesterov share the same asymptotic power law after the transient.
  • Batch-size phase diagram: At fixed data budget, SGD deteriorates first as batch size increases, while Polyak preserves small-batch performance over a wider range.Nesterov further improves data efficiency in the large-batch regime.
  • Stability boundaries: Figure 2 validates critical-learning-rate predictions against numerical estimates across batch sizes and momentum factors.The middle panel confirms linear dependence on 1−ρ before the order-one ceiling.
  • Data-scaling exponents: For B ≂1, B ≂D^0.4, and B ≂D^0.8, empirical final-risk slopes recover the predicted ordering across the three regimes.The ordering is equal at small batch, Nesterov > Polyak > SGD at large batch, and Nesterov = Polyak > SGD at ultra-large batch.

8 Conclusion and discussion

The paper concludes that momentum’s benefits depend critically on batch size and extends the analysis to increasing-momentum schedules. Polyak preserves data scaling across a wider batch range, while Nesterov improves large-batch scaling but faces stochastic stability limits.

  • Conclusion: At small batch sizes, neither Polyak nor Nesterov improves SGD’s data-scaling exponent.The methods diverge in benefit as batch size grows.
  • Conclusion: Polyak preserves the optimal small-batch data-scaling exponent over a substantially wider batch-size range than SGD.This enlarges the data-efficient batch-size range.
  • Conclusion: Nesterov improves the data-scaling exponent in the large-batch regime, while both momentum methods share signal-limited scaling at ultra-large batch sizes.At ultra-large batch sizes, the number of sequential updates limits training.
  • Momentum scheduling: For the schedule 1−ρ_k ≂k^-1, fixed η and B eventually cross the stochastic stability boundary as ρ_k approaches one.The schedule can remain stable for a substantial fraction of training before instability.
  • Momentum scheduling: The authors recommend jointly scheduling momentum, learning rate, and batch size so 1−ρ_k remains above the stochastic stability scale.The stated scales are η_k/B_k for Polyak and η_k/B_k^β for Nesterov.
  • Limitations and outlook: A complete theory of time-varying η_k, B_k, and ρ_k remains beyond the present analysis.The current scaling analysis keeps these quantities fixed within each training run.

Use of AI

The manuscript reports AI assistance in drafting, language editing, algebraic calculations, and asymptotic estimates, while attributing proof strategies and mathematical arguments to the authors.

  • AI tools assisted with drafting and language editing throughout the manuscript.
  • In Sections B–F, AI tools also accelerated algebraic calculations and asymptotic estimates.
  • The authors developed the proof strategies and mathematical arguments and verified all AI-assisted calculations and estimates.

A.2 Rotational equivariance of SGD and momentum methods

The appendix develops the spectral and moment-based machinery underlying stability and risk-scaling results for SGD, Polyak, and Nesterov in power-law kernel regression.

  • Rotational equivariance: Unitary changes of basis preserve the predictions, risks, and recursions of SGD, Polyak, and Nesterov.The analysis can therefore work in the eigenbasis of H without loss of generality.
  • Stability analysis: A renewal-equation criterion reduces boundedness to the kernel feedback condition ∥K∥_1 < 1.The kernel captures perturbation propagation, while additive label noise acts as external forcing.
  • SGD stability: Under λ_i ≂i^-β with β > 1, SGD has an order-one critical learning-rate scale.
  • Polyak stability: Polyak stability requires η below a boundary coupling batch size, momentum, the maximum eigenvalue, and the trace of H.The critical equation explicitly contains B, ρ, λ_i, and Tr(H).
  • Polyak stability: In the high-momentum, large-batch regime, Polyak’s critical scale is governed by B(1−ρ).
  • Nesterov stability: Nesterov has a distinct critical-learning-rate boundary derived from its risk-stability theorem and order-one local ceiling.
  • Risk dynamics: Polyak dynamics transition between overdamped and underdamped spectral modes at ηλ_j ≂(1−ρ)^2.High-curvature modes can be underdamped while flatter modes remain overdamped.
  • Noise contributions: Under stability, multiplicative gradient noise is bounded by additive label noise, allowing risk to be treated asymptotically as deterministic bias plus label noise.

D.6 Risk dynamics in the mixed-damping regime

In the mixed-damping regime, momentum dynamics separate into overdamped polynomial tails and underdamped transients, while Nesterov’s look-ahead mechanism yields a lower noise floor than Polyak.

  • Bias decomposition: The deterministic bias decomposes into an overdamped polynomial tail and an underdamped oscillatory transient.The split occurs at the crossover ηλ_Jmax ≂ (1−ρ)^2.
  • Polyak: After the transient, the overdamped tail contributes a polynomial risk term, while the underdamped component can exhibit destructive interference.Consequently, the underdamped contribution lacks a strictly positive pointwise lower bound.
  • Polyak: The mixed-damping regime is characterized by η ≳ (1−ρ)^2 and stability requires η ≲ min{1, B(1−ρ)} for Polyak.
  • Nesterov: Nesterov’s coordinate dynamics use a look-ahead gradient and characteristic roots whose discriminant determines overdamped versus underdamped behavior.The corresponding Green function represents the coordinate-wise response to stochastic perturbations.
  • Nesterov: Nesterov’s multiplicative noise is absorbed into the additive label-noise term under η ≲ min{1, B^β(1−ρ)}, producing an asymptotically equivalent bias-plus-noise risk.Its label-noise contribution is lower than Polyak’s.

E.5 Risk dynamics in the fully overdamped regime

In the fully overdamped regime, all characteristic roots are real and the risk dynamics progress through an initial transient, a slowly decaying phase, and eventual power-law behavior under the stated stability condition.

  • Root structure: All characteristic roots are real in the fully overdamped regime, so the deterministic bias has no oscillatory component.
  • Root structure: The slow root remains close to one while the fast root decays rapidly, separating persistent learning dynamics from a short-lived component.
  • Temporal stages: For k ≲ 1/(1−ρ), leading modes remain essentially undecayed, leaving the deterministic bias approximately unchanged.
  • Temporal stages: For k ≳ 1/(1−ρ), the fast component is negligible and the slow component governs the subsequent risk evolution.
  • Scaling law: Under η ≲ min{1, B^β(1−ρ)} and η ≲ (1−ρ)^2, Nesterov’s expected excess risk follows the fully overdamped scaling law.
  • Batch scaling: Polyak preserves the optimal data-scaling exponent s/(s+1) up to b3, extending SGD’s critical batch size from D^b1 to D^b3.

F.3 Nesterov

For Nesterov, optimizing learning-rate and momentum scalings yields a three-branch data-scaling exponent: the small-batch SGD rate, an improved intermediate regime, and an ultra-large-batch branch shared with Polyak.

  • Optimization setup: The Nesterov stability condition scales as η ≲ min{1, B^β(1−ρ)}, while the transient constraint requires b+m<1.
  • Regime comparison: The fully overdamped Nesterov dynamics have the same batch-size-dependent scaling exponent as SGD.
  • Data-scaling regimes: For b1 ≤ b < b2, Nesterov attains the intermediate exponent sβ/(sβ+1), with b2 determined by the transient constraint.
  • Data-scaling regimes: For b ≥ b2, Nesterov’s exponent becomes 2s(1−b), matching Polyak in the ultra-large-batch regime.
  • Data-scaling regimes: For 0 ≤ b < b1, Nesterov attains the small-batch exponent s/(s+1), matching SGD and Polyak.
  • Validation: Numerical fits across additional (s,β) choices consistently track the predicted risk-dynamics scaling forms.

G.3 Batch-size phase diagram at a fixed data budget

At a fixed data budget, experiments evaluate optimized final risk across batch sizes and test representative small-, large-, and ultra-large-batch regimes.

  • Experimental design: For each batch size, the experiment searches momentum-method hyperparameters near theoretically prescribed scalings and reports the minimum final excess risk.
  • Additional validation: Additional experiments extend the batch-size phase diagram and data-scaling tests across different (s,β) choices.
  • Small batch: Small-batch experiments use b=0 and learning rates near D^-s/(s+1), with momentum damping held at order one.
  • Large batch: Large-batch experiments use b=b2−δ and tune Polyak and Nesterov near η ≂1 with prescribed momentum scalings.
  • Ultra-large batch: Ultra-large-batch experiments choose b3<b<1 and tune both momentum methods with η ≂1 and momentum damping approaching the transient boundary.
Loading 2609.02728v1…