这篇论文把HB和ASGD在批量大小上的效率权衡讲清楚了,特别是那个sqrtκ倍的窗口,对想用动量方法加速训练的人很有参考价值。
该论文研究了随机重球法(HB)和加速SGD(ASGD)在一致线性回归中的计算效率与串行运行时间权衡。结果表明HB在任意谱下无法超越SGD的计算效率前沿,但允许在比SGD临界批量大sqrtκ倍的窗口内减少串行运行时间。ASGD在快速衰减幂律谱下可提升小批量计算效率,但随着批量增大,牺牲效率换取更优串行时间。合成线性回归实验验证了这些定性规律。
Compute Efficiency and Serial Runtime Tradeoffs for Stochastic Momentum Methods
Stochastic momentum methods such as heavy ball (HB), Nesterov momentum, and variants of Accelerated SGD (ASGD) [Kidambi et al., 2018] are widely used in modern training, but their stochastic benefits depend on two distinct quantities: serial runtime, the number of iterations needed to reach a target accuracy, and compute efficiency (CE), the inverse total gradient-query or FLOP cost. Larger batches reduce serial runtime without hurting CE only when the contraction gap grows linearly with batch size. We study stochastic HB and ASGD for consistent linear regression with Gaussian covariates and prove finite-dimensional, discrete-time lower bounds on their batch-size tradeoffs. Our first result shows that HB does not improve the CE frontier over SGD for arbitrary spectra; rather, it preserves SGD-level CE over a larger batch-size window, allowing larger batches to reduce serial runtime until HB reaches its deterministic accelerated scale. This window can be a factor $\sqrtκ$ larger than the SGD critical batch size. For ASGD, the picture is more spectrum-dependent: for rapidly decaying power-law spectra, ASGD improves small-batch CE over HB/SGD, but as batch size grows it trades this CE advantage for improved serial runtime. Synthetic linear-regression experiments verify these qualitative regimes, including near-overlap of ASGD and HB for slowly decaying spectra and the predicted CE--serial tradeoff for rapidly decaying spectra.