Source-linked AI summary
On the insufficiency of existing momentum schemes for Stochastic Optimization
Rahul Kidambi, Praneeth Netrapalli, Prateek Jain, Sham M. Kakade
TL;DR
The paper asks whether HB and NAG genuinely improve stochastic optimization beyond SGD, given that their classical guarantees assume exact gradients. It proves limitations for HB, finds similar empirical behavior for NAG, and proposes ASGD, which improves performance across the studied problem instances.
Problem
It is unknown whether momentum methods provide provable improvement over SGD in the stochastic first-order-oracle model with small mini-batches.
Method
The paper proves lower-bound results for HB, evaluates NAG empirically, and develops an intuitive ASGD variant for stochastic optimization.
Results
ASGD significantly improves over SGD, HB, and NAG on the studied problem instances, including 5−7% better test error than HB and NAG during initial CIFAR-10 training.
Takeaways & Limitations
The results suggest that HB and NAG’s practical gains are primarily attributable to mini-batching, whereas ASGD can improve over SGD across batch-size regimes.
Takeaways & Limitations
The theoretical HB result applies to specified problem instances, while NAG’s comparable behavior is observed empirically rather than proved in this paper.
Abstract
from arXiv · showhide
Momentum based stochastic gradient methods such as heavy ball (HB) and Nesterov's accelerated gradient descent (NAG) method are widely used in practice for training deep networks and other supervised learning models, as they often provide significant improvements over stochastic gradient descent (SGD). Rigorously speaking, "fast gradient" methods have provable improvements over gradient descent only for the deterministic case, where the gradients are exact. In the stochastic case, the popular explanations for their wide applicability is that when these fast gradient methods are applied in the stochastic case, they partially mimic their exact gradient counterparts, resulting in some practical gain. This work provides a counterpoint to this belief by proving that there exist simple problem instances where these methods cannot outperform SGD despite the best setting of its parameters. These negative problem instances are, in an informal sense, generic; they do not look like carefully constructed pathological instances. These results suggest (along with empirical evidence) that HB or NAG's practical performance gains are a by-product of mini-batching. Furthermore, this work provides a viable (and provable) alternative, which, on the same set of problem instances, significantly improves over HB, NAG, and SGD's performance. This algorithm, referred to as Accelerated Stochastic Gradient Descent (ASGD), is a simple to implement stochastic algorithm, based on a relatively less popular variant of Nesterov's Acceleration. Extensive empirical results in this paper show that ASGD has performance gains over HB, NAG, and SGD.
1 Introduction
The paper questions whether momentum methods genuinely improve stochastic optimization, showing that HB and NAG can fail to outperform SGD while ASGD provides significant gains. Its results attribute practical momentum benefits primarily to mini-batching and report improved deep-learning performance for ASGD.
- ASGD is designed to improve over SGD across small and large mini-batch sizes, unlike HB and NAG, which often struggle with small constant batches.
- On CIFAR-10 residual-network training, ASGD achieved 5−7% better test error than HB and NAG during initial optimization with batch size 128.
- HB is not optimal in the stochastic first-order-oracle model: on some linear regression problems, it performs no better than SGD while ASGD improves over both.
- Experiments suggest HB’s suboptimality is widespread across linear regression problems, with similar empirical behavior observed for NAG.
- Mini-batching reduces stochastic-gradient variance and moves the stochastic oracle toward the exact-gradient setting, explaining momentum methods’ practical superiority.
- The paper presents an easier-to-tune ASGD variant that converges faster than SGD, HB, and NAG while achieving favorable or comparable asymptotic accuracy on several deep-learning problems.
2 Notation
The notation defines the objective, exact and stochastic gradients, linear-regression losses, and the momentum updates used by HB and NAG.
- The objective is f(w) = 1/n P_i f_i(w), where w contains the model parameters and ∇f(w) is its exact gradient.
- The stochastic gradient b∇f_t(w_t) samples one component gradient uniformly, and for linear regression f_i(w) = 0.5 · (b_i − ⟨w, a_i⟩)^2.
- HB combines the SGD update with the momentum term w_t − w_{t−1}, while α denotes the momentum parameter.
- NAG is represented as a stochastic version of the commonly used accelerated-gradient method of Nesterov (1983).
3 Suboptimality of Heavy Ball Method
The paper constructs linear regression instances where Heavy Ball cannot improve over SGD beyond a constant, although ASGD achieves a substantial improvement. The section derives contrasting convergence rates and identifies when ASGD’s gains matter.
- Problem instance: The constructed instance samples feature directions σ1·z·e1 and σ2·z·e2 with equal probability and sets b = ⟨w∗, a⟩.This defines the two-dimensional linear regression distribution used for the comparison.
- Convergence-rate comparison: ASGD improves upon SGD by a factor of √κ on the constructed linear regression problem.Here, the computational condition number is κ = cσ2 while the statistical condition number is ˜κ = c.
- Heavy Ball lower bound: For any Heavy Ball stepsize and momentum, its performance does not improve over SGD by more than a constant.The proposition establishes this asymptotically for sufficiently large t, with a bound whose constant depends on κ, δ, and α.
- ASGD procedure: The section’s ASGD procedure combines short-step and long-step updates through a running average and returns the final iterate.Its inputs include a short step, a long-step parameter κ, and a statistical advantage parameter ξ.
- Convergence-rate comparison: ASGD reaches an ϵ-approximation in O(√κ log κ log 1/ϵ) iterations, whereas the displayed SGD rate is slower on this instance.The improvement is meaningful when κ > O(c); otherwise, all algorithms have nearly the same rates up to constant factors.
- Scope beyond Heavy Ball: The section notes that NAG empirically obtains nearly the same rate as HB and SGD on the same problem, while a matching lower bound remains conjectural.This observation is empirical rather than a proved NAG lower bound.
4 Algorithm
ASGD is presented as an intuitive stochastic method with three tunable inputs: a short step, a long-step parameter, and a statistical advantage parameter. The short step matches the step size used by SGD, HB, and NAG, while the long-step parameter estimates curvature variation.
- ASGD inputs: ASGD takes a short step δ, a long-step parameter κ, and a statistical advantage parameter ξ as inputs.The method is presented as an easier-to-tune version of accelerated stochastic gradient descent.
- Parameter interpretation: The short step δ is the same step size used by SGD, HB, and NAG and scales inversely with function smoothness for convex problems.The long-step parameter κ estimates the ratio of the largest and smallest curvatures.
5 Experiments
Experiments across linear regression and deep networks test whether momentum methods' advantages persist with small minibatches and whether ASGD improves convergence. Results indicate that HB and NAG gains depend strongly on minibatching, while ASGD often converges faster, especially at small batch sizes.
- 5.1 Linear Regression: Experiments vary κ across two-dimensional discrete and Gaussian linear-regression problems, using 100 runs per κ and t = 5κ iterations.Convergence is assessed by whether the error in the second half of training exceeds the starting error.
- 5.1 Linear Regression: ASGD scales with √κ, whereas SGD, HB, and NAG show nearly linear dependence on condition number κ in linear-regression experiments.The comparison uses discrete and Gaussian distributions, with parameters for HB and NAG selected by grid search.
- 5.2 Deep Autoencoders for MNIST: For MNIST autoencoders, HB and NAG decay loss faster than SGD with minibatch size 8, but this advantage disappears at minibatch size 1.With batch size 1, SGD, HB, and NAG decrease loss at similar rates, while ASGD decays faster than all three.
- 5.3 Deep Residual Networks for CIFAR-10: For CIFAR-10 residual networks, NAG converges faster than SGD and HB at batch size 128, but this rate advantage disappears at batch size 8.Final errors for SGD, HB, and NAG are close for both batch sizes, while the batch-size comparisons use different training durations.
- 5.3 Deep Residual Networks for CIFAR-10: ASGD converges faster than NAG across the tested CIFAR-10 settings and surpasses HB after fully optimizing ASGD's learning rate and decay schedule.Using HB's schedule initially obscures ASGD's advantage at batch size 8 because that schedule is suboptimal for ASGD.
6 Related Work
Related work establishes that momentum methods are effective with exact gradients but their stochastic-oracle behavior remains theoretically unclear. Prior results show that acceleration with noisy gradients is delicate, while several finite-sum methods achieve faster rates under different settings.
- Accelerating SGD is difficult because fast gradient methods can be unstable under statistical, numerical, and adversarial errors.
- Jain et al. (2017) developed a provably accelerated SGD method for linear regression that achieves minimax rates, inspired by a Nesterov variant.
- HB and NAG are effective with exact first-order oracles, but their theoretical performance in the stochastic first-order oracle model is not well understood.
- Finite-sum methods including SDCA, SAG, SVRG, and SAGA improve over SGD’s sub-linear rates for strongly convex finite sums by obtaining linear convergence.
- Momentum methods are standard in deep-network training, but the paper attributes their practical improvement over SGD to mini-batching rather than stochastic acceleration.
7 Conclusions and Future Directions
The paper argues that HB’s stochastic performance gains come from mini-batching, not intrinsic acceleration, and presents ASGD as an alternative with faster convergence on several problem classes. Its analysis establishes parameter-independent lower bounds through update-operator eigenvalues.
- HB’s stochastic performance gains are attributed to mini-batching, whereas ASGD is designed to improve over SGD across small and large mini-batch sizes.
- The paper concludes that ASGD significantly improves convergence over SGD, HB, and NAG on the studied problems and motivates extensions to NLP and automatic momentum tuning.
- For every stepsize and momentum choice, the analysis shows that one of two update operators has an eigenvalue preventing HB from improving sufficiently over SGD.
- The analysis decouples across the two coordinates, allowing the stochastic dynamics to be studied separately in each direction.
A.1 Proof
The proof analyzes the characteristic polynomial of HB’s expected covariance-update operator and bounds its eigenvalues in separate momentum regimes. These bounds show that HB cannot attain accelerated convergence on the constructed stochastic problems.
- The proof analyzes the characteristic polynomial of the update operator and evaluates it at selected points to bound its roots.
- When the characteristic polynomial is nonpositive at one and nonnegative at infinity, continuity yields a root at least one.
- Increasing momentum requires reducing the stochastic heavy-ball stepsize for convergence, unlike the deterministic c = 1 case described in the remark.
- For sufficiently large momentum, the update operator has an eigenvalue with magnitude close to one, independently of the learning rate.
- For momentum bounded away from one, the proof combines separate eigenvalue bounds to show that some eigenvalue remains at least 1−500/κ in magnitude.
- The eigendirection argument applies to any starting iterate distinct from the optimum, strengthening the result beyond initialization-specific lower bounds.
B Equivalence of Algorithm 3 and ASGD
The appendix verifies that the intuitive Algorithm 3 implementation is algebraically equivalent to the ASGD updates. Substitution identifies its running average and convex-combination iterate updates.
- The appendix starts from ASGD’s two-iterate updates and specifies the corresponding step-size parameters.
- Substituting the parameter choices shows that the auxiliary iterate is exactly the ASGD running average.
- The main iterate is updated as a convex combination of the current running average and a short gradient step.
C More details on experiments
This section introduces additional details about the paper’s experimental setup.
- The section is devoted to presenting more details on the experimental setup.
- It serves as a setup-details section rather than reporting a specific experimental result.
- No particular dataset, algorithm comparison, or quantitative outcome is specified here.
C.1 Linear Regression
The linear-regression experiments analyze algorithmic error decay through covariance-operator dynamics and compare optimized rates across condition numbers. ASGD exhibits square-root dependence on the condition number, unlike the linear dependence observed for SGD, HB, and NAG.
- C.1 Linear Regression: Expected errors for SGD, HB, NAG, and ASGD can be computed by tracking covariance matrices that evolve as linear systems.
- C.1 Linear Regression: The largest-magnitude eigenvalue of each algorithm’s covariance operator determines its error-decay rate, with smaller values indicating faster decay.
- C.1 Linear Regression: The experiments vary κ across powers of two, estimate fourth-moment tensors from 1000 samples, and grid-search learning-rate and momentum parameters for NAG and HB.
- C.1 Linear Regression: ASGD has √κ dependence on the condition number, while SGD, HB, and NAG have linear dependence after parameter grid search.Figure 7 reports this pattern for both discrete and Gaussian linear-regression problems.
C.3 Deep Residual Networks for CIFAR-10
The CIFAR-10 experiments compare SGD, HB, NAG, and ASGD across minibatch sizes using grid-searched parameters and matched schedules. NAG’s early advantage depends on batch size, while ASGD shows trends consistent with its validation-error comparisons against momentum methods.
- C.3 Deep Residual Networks for CIFAR-10: The experiments use separate parameter grids for batch sizes 8 and 128, including learning rates, momentum or long-step settings, advantage parameters, and decay factors.
- C.3 Deep Residual Networks for CIFAR-10: Across algorithm comparisons, ASGD keeps the comparator’s starting learning rate, decay factor, and schedule fixed while searching its long-step and advantage parameters.
- C.3 Deep Residual Networks for CIFAR-10: NAG substantially improves over SGD and HB early for batch size 128 but not batch size 8, then decreases training function value rapidly for both sizes later.
- C.3 Deep Residual Networks for CIFAR-10: ASGD training-error comparisons with HB and NAG show trends similar to the corresponding validation-error plots.Figures 9 and 10 compare ASGD with HB and NAG, respectively.