PaperlayerView exact revision on arXiv

Source-linked AI summary

Momentum and Stochastic Momentum for Stochastic Gradient, Newton, Proximal Point and Subspace Descent Methods

Nicolas Loizou, Peter Richtárik

arXiv:1712.09677v2math.OCcs.LGmath.NAstat.ML

TL;DR

The paper addresses limited evidence on combining stochastic approximation with heavy-ball momentum. It analyzes equivalent primal and dual stochastic methods in a structured quadratic setting, proving linear convergence and showing that stochastic momentum can improve overall complexity in some sparse regimes.

  • Problem

    The paper studies the largely unknown combined effect of stochastic approximation and momentum, including the absence of known linear convergence results for such methods.

  • Method

    The paper analyzes SGD, stochastic Newton, stochastic proximal point, and stochastic dual subspace ascent in an equivalent stochastic quadratic and linear-system framework.

  • Results

    The paper proves global linear convergence for momentum variants across several measures and shows linear convergence with stochastic momentum, with better overall complexity than deterministic momentum in some sparse settings.

  • Takeaways & Limitations

    Stochastic momentum can reduce per-iteration momentum cost while retaining linear convergence and improving overall complexity for some sparse-data regimes.

  • Takeaways & Limitations

    The analysis is centered on special quadratics; beyond that setting, algorithmic equivalence can break and each method may require separate analysis.

Abstract

from arXiv · show

In this paper we study several classes of stochastic optimization algorithms enriched with heavy ball momentum. Among the methods studied are: stochastic gradient descent, stochastic Newton, stochastic proximal point and stochastic dual subspace ascent. This is the first time momentum variants of several of these methods are studied. We choose to perform our analysis in a setting in which all of the above methods are equivalent. We prove global nonassymptotic linear convergence rates for all methods and various measures of success, including primal function values, primal iterates (in L2 sense), and dual function values. We also show that the primal iterates converge at an accelerated linear rate in the L1 sense. This is the first time a linear rate is shown for the stochastic heavy ball method (i.e., stochastic gradient descent method with momentum). Under somewhat weaker conditions, we establish a sublinear convergence rate for Cesaro averages of primal iterates. Moreover, we propose a novel concept, which we call stochastic momentum, aimed at decreasing the cost of performing the momentum step. We prove linear convergence of several stochastic methods with stochastic momentum, and show that in some sparse data regimes and for sufficiently small momentum parameters, these methods enjoy better overall complexity than methods with deterministic momentum. Finally, we perform extensive numerical testing on artificial and real datasets, including data coming from average consensus problems.

1 Introduction

The paper studies stochastic optimization with heavy-ball momentum in a setting linking stochastic, approximation, and dual quadratic problems. This shared structure makes several primal and dual methods equivalent and enables a unified momentum analysis.

  • Stochastic approximation replaces costly quantities such as gradients with cheaper stochastic approximations, while heavy-ball momentum is intended to accelerate gradient-type methods.
  • The paper addresses the limited knowledge about combining stochastic approximation and momentum, establishing linear convergence for important problem classes and parameter ranges.
  • The stochastic optimization problem uses random least-squares-type convex quadratics whose minimizers coincide with solutions of a consistent linear system.
  • Because of the problem's structure, stochastic gradient descent, stochastic Newton, and stochastic proximal point methods with fixed stepsize are identical.
  • The equivalence extends to 0 < ω < 2 and allows momentum variants of the primal methods to be studied through a unified algorithmic analysis.
  • The same framework connects best approximation and bounded concave quadratic maximization through stochastic dual subspace ascent, which moves in random subspaces spanned by sampled matrix columns.

2 Momentum Methods and Our Contributions

The paper analyzes heavy-ball momentum across several stochastic optimization methods in an equivalent quadratic setting, establishing linear and accelerated convergence results while introducing stochastic momentum.

  • Motivation: Prior stochastic heavy-ball analyses established only sublinear rates, leaving acceleration from momentum unresolved.The paper identifies the absence of known linear convergence for stochastic approximation combined with momentum as a central gap.
  • New momentum methods: The paper studies momentum variants of SGD, stochastic Newton, stochastic proximal point, and stochastic dual subspace ascent in a quadratic setting where the methods are equivalent.Momentum variants of stochastic Newton, stochastic proximal point, and stochastic dual subspace ascent were not previously analyzed to the authors’ knowledge.
  • Linear rate: The primal momentum methods achieve global, non-asymptotic linear convergence for expected squared iterate error and expected function-value suboptimality, with corresponding dual results for mSDSA.These rates hold for ranges of stepsize ω > 0 and momentum parameter β ≥ 0.
  • Accelerated linear rate: The expected primal iterates achieve an accelerated linear rate depending on the square root of the Hessian condition number.The paper describes this as the first accelerated rate for the stochastic heavy-ball method and a quadratic speedup over no-momentum methods.
  • Sublinear rate for Cesaro averages: Cesaro averages of primal iterates converge at a sublinear O(1/k) rate under weaker assumptions than those required for linear convergence.This result applies to all primal momentum methods.
  • Primal-dual correspondence: The methods are related through affine primal-dual correspondence, extending prior equivalence results beyond the no-momentum and unit-stepsize cases.The correspondence connects SGD, stochastic Newton, and stochastic proximal point with stochastic dual subspace ascent.
  • Stochastic momentum: Stochastic momentum approximates deterministic momentum coordinate-wise, reducing per-iteration cost but adding noise and iteration overhead; for sparse randomized Kaczmarz settings, overall complexity can improve.The paper proves linear convergence for several methods using stochastic momentum and reports better overall complexity for SGD with stochastic momentum in some sparse regimes.
  • Scope: The analysis is specialized to quadratic objectives; beyond this setting, algorithmic symmetry breaks and separate analyses are needed for each method.The authors present generalization beyond special quadratics as future work.

3 Technical Preliminaries

The paper uses stochastic reformulations of consistent linear systems to place several optimization and subspace methods in a common framework. Under exactness, these methods converge to the B-projection of the initial point onto the solution set.

  • Stochastic reformulation: A consistent linear system is reformulated as equivalent stochastic optimization, linear-system, fixed-point, and intersection problems using random matrices S drawn from a user-defined distribution D.The reformulations share identical solution sets and connect linear algebra with stochastic optimization.
  • Stochastic reformulation: The stochastic objective is a least-squares-type convex quadratic depending on S, A, b, and a positive definite matrix B defining the inner product and norm.Its Hessian is represented through the expectation of a random positive semidefinite matrix H, with the Moore–Penrose pseudoinverse used in H's construction.
  • Assumptions: Exactness ensures that the minimizers of the stochastic objective coincide with the solution set L of Ax = b; positive definiteness of E[H] is sufficient for exactness.Without exactness, the objective may also vanish outside L, producing a larger minimizer set.
  • Equivalent methods: SGD is equivalent in this setting to stochastic Newton and stochastic proximal point methods, while randomized Kaczmarz appears as a special case of SGD.The stochastic gradient is taken with respect to the B-induced inner product, and the stochastic Newton method uses a B-pseudoinverse because the sampled Hessian may be singular.
  • Primal–dual methods: Sketch-and-project methods replace L by a sketched system L_S and project the current iterate onto L_S; their framework expresses randomized variants of many classical algorithms.The cited framework also supports randomized methods for matrix inverse and pseudoinverse estimation and average consensus.
  • Primal–dual methods: The primal methods converge linearly to x∗ = ΠB_L(x0), the B-projection of the initial iterate onto the linear-system solution set, and SDSA moves in random subspaces generated by S.SDSA samples S and greedily chooses a dual update maximizing the dual objective, with the smallest-norm solution selected when the maximizer is nonunique.

4 Primal Methods with Momentum

The paper analyzes heavy-ball momentum variants of SGD, stochastic Newton, and stochastic proximal point methods in a shared stochastic optimization setting. Under suitable parameter conditions, these methods achieve global linear convergence, while Cesàro averages achieve a sublinear rate and primal iterates an accelerated linear rate in L1.

  • Method: mSGD, mSN, and mSPP are obtained by adding the heavy-ball term to equivalent stochastic methods.The methods use stepsize ω and momentum parameter β.
  • L2 convergence and function values: linear rate: Global linear convergence holds for mSGD/mSN/mSPP over a range of ω > 0 and β ≥ 0.The result is stated for stochastic heavy-ball methods and implies convergence of expected function values.
  • L2 convergence and function values: linear rate: The linear-rate result is the first established for a stochastic heavy-ball method, and the momentum variants of stochastic Newton and proximal point methods are also new.Prior stochastic heavy-ball results cited by the paper were sublinear.
  • L2 convergence and function values: linear rate: For β ≥ 0, the lower bound on the convergence factor increases with β, and the rate is inferior to SGD’s β = 0 rate.The paper leaves open whether mSGD can ever have a strictly better rate than SGD.
  • Cesaro average: sublinear rate without exactness assumption: Cesàro averages have an O(1/k) convergence rate under weaker conditions, including 0 ≤ β < 1 and ω + 2β < 2.This extends the stochastic heavy-ball analysis beyond the exactness assumptions used for the global linear result.
  • L1 convergence: accelerated linear rate: A suitable combination of ω and β yields an accelerated linear convergence rate for primal iterates in the L1 sense.The convergence factor in this result equals the momentum parameter β.

5 Dual Methods with Momentum

The paper introduces mSDSA as a momentum variant of stochastic dual subspace ascent and establishes its relationship to the primal momentum methods. This correspondence transfers convergence results to dual function values.

  • Method: mSDSA extends stochastic dual subspace ascent by adding a heavy-ball momentum term.SDSA is recovered when β = 0.
  • Correspondence Between Primal and Dual Methods: Under shared ω, β, and random matrices, primal iterates are affine images of mSDSA’s dual iterates.The mapping φ makes the dual recursion satisfy the same momentum recursion as mSGD.
  • Convergence of dual objective: The correspondence enables a linear convergence result for mSDSA in dual objective values.The theorem assumes exactness, 0 ≤ ω ≤ 2, β ≥ 0, and q < 1.
  • Convergence of dual objective: When ω = 1 and β = 0, the mSDSA result specializes to the previously established SDSA rate.The same rate can also be proved for the expected primal-dual gap.

6 Methods with Stochastic Momentum

The paper replaces the full heavy-ball momentum step with an unbiased stochastic approximation to reduce per-iteration cost. Although this worsens the convergence rate, it can improve total complexity in sparse-data regimes and for small momentum parameters.

  • Motivation and construction: When the momentum step dominates the stochastic-gradient cost, the paper replaces it with a cheap stochastic approximation.A uniformly sampled coordinate gives an estimator computable in O(1) time.
  • Motivation and construction: The stochastic momentum estimator is unbiased and reduces iteration cost from O(∥A_j:∥0 + n) to O(∥A_j:∥0).The saving removes the O(n) deterministic momentum cost.
  • Methods with stochastic momentum: smSGD, smSN, and smSPP apply stochastic momentum to the equivalent SGD, stochastic Newton, and stochastic proximal point methods.The paper develops the construction from the SGD perspective.
  • Convergence: L2 linear convergence is established for smSGD/smSN/smSPP when B = I and 0 < ω < 2 with admissible β.The theorem requires the stated exactness and parameter conditions.
  • Complexity comparison: Using momentum parameter βn makes stochastic momentum unbiased relative to deterministic momentum parameter β.This scaling supports the rate and complexity comparison between smSGD and mSGD.
  • Complexity comparison: The stochastic-momentum method has a worse convergence rate than deterministic momentum, despite potentially lower total complexity.The paper explicitly states that mSGD’s rate is always better than smSGD’s rate.
  • Complexity comparison: For small β, smSGD’s total complexity is approximately 1 + n/g times smaller than mSGD’s.Here n is the number of columns of A and g is the expected stochastic-gradient evaluation cost.

7 Special Cases: Randomized Kaczmarz with Momentum and Randomized Coordinate Descent with Momentum

The paper specializes its momentum framework to randomized Kaczmarz and randomized coordinate descent, connecting both to stochastic gradient descent. These variants retain the underlying randomized updates while adding heavy-ball momentum and achieve quadratic improvements in stated L1 iteration-complexity bounds.

  • Special cases: mRK and mRCD are momentum variants of randomized Kaczmarz and randomized coordinate descent within the mSGD framework.The framework also includes randomized coordinate Newton and randomized Gaussian Kaczmarz as momentum special cases.
  • Randomized Kaczmarz: For β = 0, mRK reduces to randomized Kaczmarz with relaxation, and with ω = 1 it becomes the standard randomized Kaczmarz method.The corresponding no-momentum, no-relaxation method has the previously established iteration complexity.
  • Randomized Kaczmarz: Theorem 4 yields an L1 iteration-complexity bound for mRK that is a quadratic improvement over the previous best result.The supplied passages state the contrast and improvement but do not preserve the full displayed complexity expression.
  • Randomized coordinate descent: For mRCD, the choice B = A and coordinate sampling probability pi = Aii / Trace(A) produces a randomized coordinate-descent specialization for positive definite A.Its momentum complexity result is likewise reported as a quadratic improvement over the previous best result.
  • Acceleration mechanism: The graphical interpretation decomposes mRK into a randomized hyperplane projection followed by a momentum displacement parallel to the preceding iterate difference.The momentum step is described as moving the next iterate closer to x∗ than the projection point in the illustrated example.
  • Acceleration mechanism: After iterate x3, momentum becomes more effective in the illustration, a pattern also observed experimentally after roughly the first m + 1 iterations.The remark reports similar behavior in experiments with matrices having many rows.

8 Numerical Experiments

The numerical experiments compare momentum and no-momentum methods on synthetic systems, real datasets, stochastic-momentum variants, and average-consensus problems. Across these tests, appropriately chosen momentum commonly improves convergence, with benefits depending on conditioning and momentum size.

  • Experimental design: Global L2 error and expected function values converge linearly for mSGD, but the experiments do not show a faster rate than the no-momentum method on these measures.The accelerated linear rate is reported only for the weaker L1 sense.
  • Experimental design: The experiments evaluate mSGD special cases mRK and mRCD on synthetic Gaussian systems and real matrices, then test stochastic momentum and average-consensus applications.Performance is measured using relative error, function values, iterations, and wall-clock time.
  • Synthetic data: For ill-conditioned systems with large 1/λ+ min, momentum is most beneficial; for well-conditioned systems, the no-momentum method already converges rapidly.The paper reports this conditioning dependence explicitly for the synthetic comparisons.
  • Synthetic data: With ω = 1, smaller condition numbers favor smaller β, while β = 0.3 or β = 0.4 are reported as good options in the first two rows of Figures 2 and 3.For large 1/λ+ min, the paper identifies β = 0.5 as the best choice in the reported examples.
  • Synthetic data: β = 0.5 halves mRK’s iterations to relative error 10^-10 for A ∈ R300×280, while reaching relative error around 10^-9 in 12 seconds versus over 25 seconds for RK.The example uses 1/λ+ min = 208,730.
  • Real data: mRK and mRCD with appropriately chosen 0 < β ≤ 0.5 are observed to converge faster than RK and RCD, respectively.The reported momentum range is smaller than β ≈ 0.9 commonly used with mSGD for deep-neural-network training.
  • Average consensus: Momentum improves pairwise randomized gossip across the studied networks, with β = 0.4 improving vanilla PRG for all networks and β = 0.5 favorable for cycle and line graphs.These experiments address average-consensus problems, including rank-deficient incidence matrices.

A.1 Lemmas

This appendix section introduces a recurrence lemma used to establish geometric decay and records identities linking the stochastic quadratic objective to its gradient and projection structure.

  • Recurrence lemma: Lemma 9 studies nonnegative sequences satisfying a two-step recurrence with coefficients a1 and a2.The assumptions require a2 ≥ 0, a1 + a2 < 1, and at least one coefficient positive.
  • Recurrence lemma: The lemma bounds the sequence geometrically as Fk+1 ≤ q^k(1 + δ)F0 for k ≥ 1.Here q is constructed from a1 and a2 through the auxiliary parameter δ.
  • Recurrence lemma: The proof adds δFk to the recurrence so that the resulting expression contracts by q across successive steps.Unrolling this transformed recurrence yields the stated geometric bound.
  • Quadratic identities: For x∗ satisfying Ax∗ = b, the paper identifies the B-inner product of the gradient with x − x∗ as 2f(x).This follows by taking expectations of the corresponding stochastic quadratic identity.

A.2 The Proof

The proof derives a two-step inequality for the momentum iterates, bounds its component expressions, and applies the recurrence lemma to obtain linear convergence.

  • Proof structure: The proof decomposes the main inequality into three expressions that are bounded separately before being recombined.The resulting relation is prepared for application of Lemma 9.
  • Proof structure: Taking expectation over Sk produces an inequality involving the B-norm error, function value f(xk), and a momentum cross term.The displayed bound contains the coefficient 2ω(2 − ω) multiplying f(xk).
  • Proof structure: The cross term is controlled using ⟨∇f(xk), xk−1 − xk⟩ ≤ f(xk−1) − f(xk), together with ωβ ≥ 0.This converts the momentum interaction into terms compatible with the recurrence analysis.
  • Convergence conclusion: The recurrence coefficients satisfy the lemma’s conditions under the stated assumptions, so Lemma 9 yields the geometric convergence estimate.The proof checks a2 ≥ 0 and a1 + a2 < 1, including the β = 0 case.
  • Convergence conclusion: Expected function-value convergence follows as a corollary by applying inequality (48) to the established iterate bound.The appendix explicitly derives E[f(xk)] from the preceding inequality.

B Proof of Theorem 3

The proof constructs a momentum-adjusted potential and derives a one-step expected descent inequality, which yields convergence bounds after telescoping.

  • B Proof of Theorem 3: The expected distance recursion is obtained by expanding the momentum update and conditioning on the random matrix S_t.The expansion separates the current distance, gradient inner products, and momentum cross-term.
  • B Proof of Theorem 3: Convexity bounds the momentum cross-term using the difference between consecutive function values.The resulting inequality replaces an inner product with f(x_t) − f(x_{t−1}).
  • B Proof of Theorem 3: The coefficient α = 4ω/(1−β) − 2ω^2/(1−β)^2 must be positive for the descent inequality.Under this condition, the potential satisfies θ_{t+1} + αE[f(x_t)] ≤ θ_t.
  • B Proof of Theorem 3: Summing the potential inequality over iterations provides cumulative control of expected function values, followed by Jensen’s inequality.The proof uses telescoping before applying Jensen’s inequality to obtain the final bound.
  • B Proof of Theorem 3: The auxiliary potential θ_t combines expected distance to x∗ with a weighted preceding function value.This definition converts the descent relation into a recursive potential inequality.

C Proof of Theorem 4

The proof transforms the stochastic momentum dynamics into coordinatewise second-order recurrences and analyzes their characteristic roots under spectral conditions.

  • C Proof of Theorem 4: When the characteristic roots are imaginary, Lemma 13 supplies a bound based on the initial conditions.The imaginary-root condition is expressed as 1 + 4a2 < 0.
  • C Proof of Theorem 4: The eigenvalue decomposition of W converts the transformed iterate recursion into coordinatewise recurrences.The coordinates s_k = UᵀB^1/2E[x_k − x∗] are analyzed separately.
  • C Proof of Theorem 4: Each coordinate follows a second-degree homogeneous recurrence with coefficients a1 = 1 + β − ωλ_i and a2 = −β.The proof uses 0 ≤ λ_i ≤ 1 and 0 < ω ≤ 1/λ_max to bound ωλ_i.
  • C Proof of Theorem 4: The two spectral cases are combined to obtain a bound valid for all k ≥ 0.The proof handles the zero-eigenvalue and positive-eigenvalue cases before combining them.
  • C Proof of Theorem 4: The lower bound on β is tight, whereas the upper bound is not.The proof notes that large β is less relevant because β represents the convergence rate and is interesting mainly below 1.

D Proof of Theorem 7

The proof applies a multi-stage tower-property argument to stochastic momentum, derives a recurrence for expected squared distance, and obtains function-value convergence as a corollary.

  • D Proof of Theorem 7: Stochasticity in the momentum term requires a more involved tower-property argument than the deterministic-momentum proof.The expectations are evaluated from the innermost to the outermost level.
  • D Proof of Theorem 7: The conditional squared-distance expression is decomposed into three terms and bounded separately.The first term is constant under the relevant conditional expectation; the remaining terms use vector identities and bounds.
  • D Proof of Theorem 7: Substituting the intermediate bounds and applying the nested expectations yields a recurrence for F_k = E[∥x_k − x∗∥^2_B].The recurrence is formed after the outermost expectation is taken.
  • D Proof of Theorem 7: Lemma 9 applies because the recurrence coefficients satisfy the required sign and sum conditions under the assumptions.The proof verifies a2 ≥ 0 and a1 + a2 < 1, including the β = 0 case.
  • D Proof of Theorem 7: Function-value convergence follows by applying inequality (48) to the established distance relation.The function-value result is stated as a corollary of the preceding recursion.
Loading 1712.09677v2…