Research map / Optimization / Optimizers

Perfect Parallelization in Mini-Batch SGD with Classical Momentum Acceleration

NeurIPS 2026 · Optimization › Optimizers · approach: SGD momentum and step-size schedules · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Develops a general stochastic momentum framework for mini-batch SGD/randomized block methods on interpolation-regime quadratics, proving classical HBM/NAG momentum gives iteration complexity scaling as κ/m + √κ, i.e., perfect parallelization of mini-batches up to saturation.

Problem

Theoretical understanding of why classical momentum (heavy ball, Nesterov) accelerates mini-batch SGD is lacking; prior analyses need strong noise assumptions or prohibitively large mini-batches.

Approach

Model momentum SGD as a stochastic momentum process over stochastic contractions; bound the bias via Schur decomposition of block-diagonal 2x2 transition matrices (avoiding diagonalization), and bound variance via a matrix recursion on E[Sigma^T Sigma], yielding acceleration proportional to mini-batch size m up to a √κ saturation with simple beta = 1 - Θ(1/m).

Key results

• Theory: with beta = 1 - Θ(1/m) and m up to Θ(√κ), momentum SGD reaches ε accuracy in O((κ/m + √κ) log(κ/ε)) iterations versus O(κ log(1/ε)) without momentum.
• Kernel ridge regression experiments show CD+NAG total work stays roughly constant as mini-batch size grows until saturation, demonstrating perfect parallelization, while CDpp does not.
• For large mini-batches, classical momentum CD+NAG outperforms the tuned CDpp acceleration scheme in total work, though CDpp is better at m=1.

Surprising findings

• Mini-batching alone does not improve CD convergence, but adding momentum unlocks near-linear speed-up from mini-batching.
• A simple problem-independent choice beta = 1 - 1/m is essentially optimal in the perfect parallelization regime.

Limitations

• Analysis restricted to quadratics in the interpolation regime.
• Theory holds up to log factors and absolute constants; the bound includes a t^3 factor.
• Perfect parallelization holds only until saturation point near sqrt(kappa).

Experimental setup

Datasets: Abalone California Housing Covtype

Benchmarks: Abalone California Housing Covtype

Metrics: iterations to reach target error total work (number of sampled rows)

Baselines: CD (vanilla block coordinate descent) CDpp (Lyapunov-based accelerated coordinate descent) CD+HBM

Builds on: Polyak heavy ball Nesterov accelerated gradient Bollapragada-Chen-Ward accelerated mini-batch SGD Loizou-Richtarik momentum SGD analysis randomized Kaczmarz

Closest papers by meaning

• Dynamics of Stochastic Momentum Methods on Large-scale, Quadratic Models (NeurIPS 2021 · similarity 0.92)
• Accelerated Convergence of Stochastic Heavy Ball Method under Anisotropic Gradient Noise (ICLR 2024 · similarity 0.91)
• The Marginal Value of Momentum for Small Learning Rate SGD (ICLR 2024 · similarity 0.90)
• Gradient correlation is a key ingredient to accelerate SGD with momentum (ICLR 2025 · similarity 0.90)
• Dimension-adapted Momentum Outscales SGD (NeurIPS 2025 · similarity 0.89)
• SGD with memory: fundamental properties and stochastic acceleration (ICLR 2025 · similarity 0.88)