Linearly convergent stochastic heavy ball method for minimizing generalization error

Introduction

In this paper we study the stochastic optimization problem:

In , the authors consider solving (1) via stochastic gradient descent (SGD)

where $\omega>0$ is a fixed stepsize and ${\bf S}_{k}$ is sampled afresh in each iteration from ${\cal D}$. It is shown that, SGD converges to an $x_{*}$ which satisfies

where $x_{0}$ is the starting point. It was observed that, surprisingly, SGD is in this setting equivalent to the stochastic (pseudo)-Newton method, and the stochastic proximal point method, and that it converges at a linear rate despite the following obstacles: $f$ is not necessarily strongly convex, (1) is not a finite-sum problem, and a fixed stepsize $\omega$ is used.

In this paper we take an alternative route, and develop a stochastic variant of the heavy ball method for solving the stochastic optimization problem (1). Applied to (1), the classical heavy ball method of Polyak , with constant stepsize $\omega>0$ and constant momentum parameter $\beta\geq 0$, takes the form

This method introduces the momentum term $\beta(x_{k}-x_{k-1})$ into the gradient descent method to achieve acceleration.

Our stochastic variant of the heavy ball method, which we henceforth simply refer to by the name stochastic heavy ball method (SHB), replaces the (costly) computation of the gradient by an unbiased estimator of the gradient (“stochastic gradient”) which is hopefully much cheaper to compute:

2 Related Work

Stochastic variants of heavy ball method have been employed widely in practice, especially in the area of deep learning . Despite the popularity of the method both in convex and non-convex optimization its convergence properties are not very well understood. Recent papers that provide complexity analysis of SHB (in different setting than ours) include and . In the authors analyzed SHB for general Lipshitz continuous convex objective functions (with bounded variance) and proved the sublinear rate $O(1/\sqrt{k})$. In , a complexity analysis is provided for the case of quadratic strongly convex smooth coercive functions. A sublinear convergence rate $O(1/k^{\beta})$, where $\beta\in(0,1)$, was proved. In contrast to our results, where we assume fixed stepsize $\omega$, both papers analyze SHB with diminishing stepsizes. For our problem, variance reduction methods like SVRG , S2GD , mS2GD , SAG and SAGA are not necessary. To the best of our knowledge, our work provides the first linear convergence rate for SHB in any setting.

Convergence Results

In this section we state our convergence results for SHB.

satisfy $a_{1}+a_{2}<1$. Let $x_{*}$ be the projection of $x_{0}$ onto $\{x\;:\;\mathbf{A}x=b\}$. Then

where $q=\frac{a_{1}+\sqrt{a_{1}^{2}+4a_{2}}}{2}$ and $\delta=q-a_{1}$. Moreover, $a_{1}+a_{2}\leq q<1$.

In the above theorem we obtain global linear rate. To the best of our knowledge, this is the first time that linear rate is established for a stochastic variant of the heavy ball method in any setting. All existing results are sublinear.

If we choose $\omega\in(0,2)$, then the condition $a_{1}+a_{2}<1$ is satisfied for all

If $\beta=0$, SHB reduces to the “basic method” in (SGD with constant stepsize). In this special case, $q=1-\omega(2-\omega)\lambda_{\min}^{+}$, which is the rate established in . Hence, our result is more general.

Let $q(\beta)$ be the rate as a function of $\beta$. Note that since $\beta\geq 0$, we have

Clearly, the lower bound on $q$ is an increasing function of $\beta$. Also, for any $\beta$ the rate is always inferior to that of SGD ($\beta=0$). It is an open problem whether one can prove a strictly better rate for SHB than for SGD.

2 Cesaro average: sublinear rate without exactness assumption

In this section we present convergence results for function values computed at the Cesaro average of all past iterates. Again, our results are global in nature. To the best of our knowledge, an analysis of the Cesaro average for the SHB with $O(1/k)$ rate was not established before for any class of functions. Moreover, the result holds without the exactness assumption.

Choose $x_{0}=x_{1}$ and let $\{x_{k}\}_{k=0}^{\infty}$ be the random iterates produced by SHB, where the momentum parameter $0\leq\beta<1$ and relaxation parameter (stepsize) $\omega>0$ satisfy $\omega+2\beta<2$. Let $x_{*}$ be any vector satisfying $f(x_{*})=0$. If we let $\hat{x}_{k}=\frac{1}{k}\sum_{t=1}^{k}x_{t}$, then

3 L​1𝐿1L1 convergence: accelerated linear rate

In this section we show that by a proper combination of the stepsize parameter $\omega$ and the momentum parameter $\beta$ the proposed algorithm enjoys accelerated linear convergence rate with respect to the expected iterates.

If we choose $\omega=1$ and $\beta=\left(1-\sqrt{0.99\lambda_{\min}^{+}}\right)^{2}$, then

and the iteration complexity becomes ${\cal O}\left(\sqrt{1/\lambda_{\min}^{+}}\log(1/\epsilon)\right)$.

If we choose $\omega=1/\lambda_{\max}$ and $\beta=\left(1-\sqrt{0.99\lambda_{\min}^{+}/\lambda_{\max}}\right)^{2}$, then

and the iteration complexity becomes ${\cal O}\left(\sqrt{\lambda_{\max}/\lambda_{\min}^{+}}\log(1/\epsilon)\right)$

Note that the convergence factor is precisely equal to the value of the momentum parameter.

Experiments

This is a randomized Kaczmarz method (RK) with momentum. Note that for $\beta=0$ and $\omega=1$ this reduces to the celebrated Randomized Kaczmarz method (RK) of Strohmer and Vershynin . In Figure 1, RK with momentum is tested for several values of the momentum parameters $\beta$ and fixed stepsize $\omega=1$. For the evaluation we use both the relative error measure $\|x_{k}-x_{*}\|^{2}/\|x_{0}-x_{*}\|^{2}$ and the function suboptimality $f(x_{k})-f(x_{*})$. The starting point is chosen as $x_{0}=0$. For the horizontal axis we use either the number of iterations or the wall-clock time measured using the tic-toc Julia function. It is clear that in this setting the addition of momentum parameter is beneficial and leads to faster convergence.

References