Momentum and stochastic momentum for stochastic gradient, Newton, proximal point and subspace descent methods

Abstract

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: convex quadratic problems. We prove global non-asymptotic linear convergence rates for all methods and various measures of success, including primal function values, primal iterates, and dual function values. We also show that the primal iterates converge at an accelerated linear rate in a somewhat weaker 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 Cesàro 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.

Appendices

Appendix 1: Technical Lemmas

Lemma 9

Fix $F_1=F_0\ge 0$ and let $\{F_k\}_{k\ge 0}$ be a sequence of nonnegative real numbers satisfying the relation

$$\begin{aligned} F_{k+1}\le a_1F_k +a_2 F_{k-1}, \quad \forall k\ge 1, \end{aligned}$$

(42)

where $a_2 \ge 0$, $a_1 + a_2 <1$ and at least one of the coefficients $a_1,a_2$ is positive. Then the sequence satisfies the relation $F_{k+1}\le q^{k} (1+ \delta ) F_0$ for all $k\ge 1,$ where $q=\frac{a_1+\sqrt{a_1^2+4a_2}}{2}$ and $\delta =q-a_1\ge 0$. Moreover,

$$\begin{aligned} q \ge a_1 + a_2, \end{aligned}$$

(43)

with equality if and only if $a_2=0$ (in which case $q=a_1$ and $\delta =0$).

Proof

Choose $\delta = \frac{-a_1+\sqrt{a_1^2+4a_2}}{2}$. We claim $\delta \ge 0$ and $a_2 \le (a_1+\delta )\delta$. Indeed, non-negativity of $\delta$ follows from $a_2\ge 0$, while the second relation follows from the fact that $\delta$ satisfies

$$\begin{aligned} (a_1+\delta )\delta - a_2 = 0. \end{aligned}$$

(44)

In view of these two relations, adding $\delta F_k$ to both sides of (42), we get

$$\begin{aligned} F_{k+1} + \delta F_k \le (a_1+\delta )F_k + a_2 F_{k-1} \le (a_1+\delta )(F_k + \delta F_{k-1}) = q(F_k+\delta F_{k-1}). \end{aligned}$$

(45)

Let us now argue that $0<q<1$. Non-negativity of q follows from non-negativity of $a_2$. Clearly, as long as $a_2>0$, q is positive. If $a_2=0$, then $a_1>0$ by assumption, which implies that q is positive. The inequality $q<1$ follows directly from the assumption $a_1+a_2<1$. By unrolling the recurrence (45), we obtain $F_{k+1} \le F_{k+1} + \delta F_k \le q^k (F_1+ \delta F_0) = q^{k}(1+\delta ) F_{0}.$

Finally, let us establish (45). Noting that $a_1 = q-\delta$, and since in view of (44) we have $a_2=q\delta$, we conclude that $a_1+a_2 = q + \delta (q-1) \le q$, where the inequality follows from $q <1$. $\square$

The following identities were established in [78]. For completeness, we include different (and somewhat simpler) proofs here.

Lemma 10

([78]) For all $x \in \mathbb {R}^n$ we have

$$\begin{aligned} f_{\mathbf {S}}(x) = \frac{1}{2}\Vert \nabla f_{\mathbf {S}}(x)\Vert ^2_{\mathbf {B}}. \end{aligned}$$

(46)

Moreover, if $x_*\in \mathcal{L}$ (i.e., if $x_*$ satisfies $\mathbf{A}x_* =b$), then for all $x\in \mathbb {R}^n$ we have

$$\begin{aligned} f_{\mathbf {S}}(x) = \frac{1}{2}\langle \nabla f_{\mathbf {S}}(x),x-x_* \rangle _{\mathbf {B}}, \end{aligned}$$

(47)

and

$$\begin{aligned} f(x) = \frac{1}{2}\langle \nabla f(x),x-x_* \rangle _{\mathbf {B}}. \end{aligned}$$

(48)

Proof

In view of (10), and since $\mathbf{Z}\mathbf{B}^{-1} \mathbf{Z}= \mathbf{Z}$ (see [78]), we have

$$\begin{aligned} \Vert \nabla f_{\mathbf {S}}(x)\Vert ^2_{\mathbf {B}}&\overset{(10)}{=} \Vert \mathbf {B}^{-1} \mathbf{Z}(x-x_*)\Vert ^2_{\mathbf {B}} \\&= (x-x_*)^\top \mathbf{Z}\mathbf {B}^{-1} \mathbf {Z}(x-x_*) = (x-x_*)^\top \mathbf{Z}(x-x_*) \\&\overset{(7)}{=} (x-x_*)^\top \mathbf{A}^\top \mathbf{H}\mathbf{A}(x-x_*) = (\mathbf{A}x-b)^\top \mathbf{H}(\mathbf{A}x- b ) \quad \overset{(6)}{=}\quad 2f_{\mathbf {S}}(x). \end{aligned}$$

Moreover,

$$\begin{aligned} \langle \nabla f_{\mathbf {S}}(x),x-x_* \rangle _{\mathbf {B}}&\overset{(10)}{=} \langle \mathbf {B}^{-1} \mathbf {Z}(x-x_*),x-x_* \rangle _{\mathbf {B}}\\&= (x-x_*)^\top \mathbf {Z}\mathbf {B}^{-1} \mathbf {B}(x-x_*) \quad = \quad 2f_{\mathbf {S}}(x). \end{aligned}$$

By taking expectations in the last identity with respect to the random matrix $\mathbf {S}$, we get $\langle \nabla f(x),x-x_* \rangle _{\mathbf {B}}=2f(x).$ $\square$

Lemma 11

([78]) For all $x \in \mathbb {R}^n$ and $x_* \in \mathcal{L}$

$$\begin{aligned} \lambda _{\min }^+ f(x) \le \frac{1}{2} \Vert \nabla f(x) \Vert ^2_{\mathbf {B}} \le \lambda _{\max } f(x) \end{aligned}$$

(49)

and

$$\begin{aligned} f(x) \le \frac{\lambda _{\max }}{2} \Vert x-x_*\Vert ^2_{\mathbf {B}}. \end{aligned}$$

(50)

Moreover, if exactness is satisfied, and we let $x_* =\varPi ^{\mathbf {B}}_{\mathcal{L}}(x)$, we have

$$\begin{aligned} \frac{\lambda _{\min }^+}{2} \Vert x-x_*\Vert ^2_{\mathbf {B}} \le f(x) . \end{aligned}$$

(51)

Finally, let us present a simple lemma of an identity that we use in our main proofs. This preliminary result is known to hold for the case of Euclidean norms ($\mathbf {B}=\mathbf {I}$). We provide the proof for the more general $\mathbf {B}-$norm for completeness.

Lemma 12

Let a, b, c be arbitrary vectors in $\mathbb {R}^n$ and let $\mathbf {B}$ be a positive definite matrix. Then the following identity holds: $2 \langle a-c,c-b \rangle _{\mathbf {B}}=\Vert a-b\Vert ^2_{\mathbf {B}}-\Vert c-b\Vert ^2_{\mathbf {B}}-\Vert a-c\Vert ^2_{\mathbf {B}},$

Proof

$$\begin{aligned} LHS= & {} 2 \langle a-c,c-b \rangle _{\mathbf {B}} = 2(a-c)^\top \mathbf {B}(c-b)\\= & {} 2a^\top \mathbf {B}c-2a^\top \mathbf {B}b-2c^\top \mathbf {B}c+2c^\top \mathbf {B}b \end{aligned}$$

and

$$\begin{aligned} RHS= & {} \Vert a-b\Vert ^2_{\mathbf {B}}-\Vert c-b\Vert ^2_{\mathbf {B}}-\Vert a-c\Vert ^2_{\mathbf {B}}\\= & {} (a-b)^\top \mathbf {B}(a-b)- (c-b)^\top \mathbf {B}(c-b)-(a-c)^\top \mathbf {B}(a-c)\\= & {} a^\top \mathbf {B}a-a^\top \mathbf {B}b-b^\top \mathbf {B}a+b^\top \mathbf {B}b-c^\top \mathbf {B}c+c^\top \mathbf {B}b+b^\top \mathbf {B}c-b^\top \mathbf {B}b\\&- a^\top \mathbf {B}a+a^\top \mathbf {B}c+c^\top \mathbf {B}a-c^\top \mathbf {B}c \\= & {} 2a^\top \mathbf {B}c-2a^\top \mathbf {B}b-2c^\top \mathbf {B}c+2c^\top \mathbf {B}b \end{aligned}$$

LHS (left-hand side) = RHS (right-hand side) and this completes the proof. $\square$

Appendix 2: Proof of Theorem 1

First, we decompose

(52)

We will now analyze the three expressions separately. The first expression can be written as

(53)

We will now bound the second expression. First, we have

(54)

Using the identity from Lemma 12 for the vectors $x_k, x_*$ and $x_{k-1}$ we obtain:

$$\begin{aligned} 2 \langle x_k-x_*, x_*-x_{k-1} \rangle _{\mathbf {B}}= \Vert x_k-x_{k-1}\Vert ^2_{\mathbf {B}}- \Vert x_{k-1}-x_*\Vert ^2_{\mathbf {B}}-\Vert x_k-x_*\Vert ^2_{\mathbf {B}}. \end{aligned}$$

Substituting this into (54) gives

(55)

The third expression can be bounded as

(56)

By substituting the bounds (53), (55), (56) into (52) we obtain

$$\begin{aligned}&\Vert x_{k+1}-x_*\Vert ^2_{\mathbf {B}}\\&\quad \le \Vert x_k-x_*\Vert ^2_{\mathbf {B}}-2\omega (2-\omega )f_{\mathbf {S}_k}(x_k)\\&\qquad + \beta \Vert x_k-x_*\Vert ^2_{\mathbf {B}}+\beta \Vert x_{k}-x_{k-1}\Vert ^2_{\mathbf {B}}-\beta \Vert x_{k-1}-x_*\Vert ^2_{\mathbf {B}} \\&\qquad + 2\omega \beta \langle \nabla f_{\mathbf {S}_k}(x_k),x_{k-1}- x_k \rangle _{\mathbf {B}} + 2\beta ^2\Vert x_{k}-x_*\Vert ^2_{\mathbf {B}}+2\beta ^2\Vert x_{k-1}-x_*\Vert ^2_{\mathbf {B}}\\&\quad \le (1+3\beta + 2\beta ^2)\Vert x_k-x_*\Vert ^2_{\mathbf {B}}+ (\beta + 2\beta ^2 )\Vert x_{k-1}-x_*\Vert ^2_{\mathbf {B}}-2\omega (2-\omega )f_{\mathbf {S}_k}(x_k)\\&\qquad + 2\omega \beta \langle \nabla f_{\mathbf {S}_k}(x_k),x_{k-1}- x_k \rangle _{\mathbf {B}}. \end{aligned}$$

Now by first taking expectation with respect to $\mathbf{S}_k$, we obtain:

$$\begin{aligned} \mathbb {E}_{\mathbf{S}_k}[\Vert x_{k+1}-x_*\Vert ^2_{\mathbf {B}}]\le & {} (1+3\beta +2\beta ^2)\Vert x_k-x_*\Vert ^2_{\mathbf {B}}+ (\beta +2\beta ^2)\Vert x_{k-1}-x_*\Vert ^2_{\mathbf {B}} \\&-2\omega (2-\omega )f(x_k) + 2\omega \beta \langle \nabla f(x_k),x_{k-1}- x_k \rangle _{\mathbf {B}}\\\le & {} (1+3\beta +2\beta ^2)\Vert x_k-x_*\Vert ^2_{\mathbf {B}}+ (\beta +2\beta ^2)\Vert x_{k-1}-x_*\Vert ^2_{\mathbf {B}} \\&-2\omega (2-\omega )f(x_k) + 2\omega \beta (f(x_{k-1})-f(x_k))\\= & {} (1+3\beta +2\beta ^2)\Vert x_k-x_*\Vert ^2_{\mathbf {B}}+ (\beta +2\beta ^2)\Vert x_{k-1}-x_*\Vert ^2_{\mathbf {B}} \\&- (2\omega (2-\omega ) +2\omega \beta )f(x_k) + 2\omega \beta f(x_{k-1}). \end{aligned}$$

where in the second step we used the inequality $\langle \nabla f(x_k),x_{k-1}- x_k \rangle \le f(x_{k-1})-f(x_k)$ and the fact that $\omega \beta \ge 0$, which follows from the assumptions. We now apply inequalities (50) and (51), obtaining

$$\begin{aligned} \mathbb {E}_{\mathbf{S}_k}[\Vert x_{k+1}-x_*\Vert ^2_{\mathbf {B}}]\le & {} \underbrace{(1+3\beta +2\beta ^2 - (\omega (2-\omega ) +\omega \beta )\lambda _{\min }^+)}_{a_1}\Vert x_k-x_*\Vert ^2_{\mathbf {B}} \\&\quad + \underbrace{(\beta +2\beta ^2 + \omega \beta \lambda _{\max })}_{a_2}\Vert x_{k-1}-x_*\Vert ^2_{\mathbf {B}}. \end{aligned}$$

By taking expectation again, and letting $F_k{:}{=}\mathbb {E}[\Vert x_{k}-x_*\Vert ^2_{\mathbf {B}}]$, we get the relation

$$\begin{aligned} F_{k+1} \le a_1 F_k + a_2 F_{k-1} . \end{aligned}$$

(57)

It suffices to apply Lemma 9 to the relation (57). The conditions of the lemma are satisfied. Indeed, $a_2\ge 0$, and if $a_2=0$, then $\beta =0$ and hence $a_1=1-\omega (2-\omega )\lambda _{\min }^+>0$. The condition $a_1+a_2<1$ holds by assumption.

The convergence result in function values, $\mathbb {E}[f(x_k)]$, follows as a corollary by applying inequality (50) to (23).

Appendix 3: Proof of Theorem 3

Let $p_t=\frac{\beta }{1-\beta }(x_t-x_{t-1})$ and $d_t = \Vert x_t + p_t -x_*\Vert _{\mathbf{B}}^2$. In view of (22), we can write

$$\begin{aligned} x_{t+1}+p_{t+1}&= x_{t+1}+\frac{\beta }{1-\beta }(x_{t+1}-x_{t}) \nonumber \\&\overset{(22)}{=} x_{t}-\omega \nabla f_{\mathbf {S}_t}(x_t)+\beta (x_t-x_{t-1})\nonumber \\&\qquad +\frac{\beta }{1-\beta }\left( -\omega \nabla f_{\mathbf {S}_t}(x_t)+\beta (x_t-x_{t-1})\right) \nonumber \\&= x_{t}-[\omega +\frac{\beta }{1-\beta }\omega ] \nabla f_{\mathbf {S}_t}(x_t)+[\beta +\frac{\beta ^2}{1-\beta }](x_t-x_{t-1})\nonumber \\&= x_{t}-\frac{\omega }{1-\beta }\nabla f_{\mathbf {S}_t}(x_t)+\frac{\beta }{1-\beta }(x_t-x_{t-1})\nonumber \\&= x_t+p_t-\frac{\omega }{1-\beta } \nabla f_{\mathbf {S}_t}(x_t) \end{aligned}$$

(58)

and therefore

$$\begin{aligned} d_{t+1}&\overset{(58)}{=} \left\| x_t+p_t-\frac{\omega }{1-\beta } \nabla f_{\mathbf {S}_t}(x_t) -x_* \right\| ^2_{\mathbf {B}} \\&= d_t -2 \frac{\omega }{1-\beta } \langle x_t+p_t-x_*, \nabla f_{\mathbf {S}_t}(x_t) \rangle _{\mathbf {B}} + \frac{\omega ^2}{(1-\beta )^2} \Vert \nabla f_{\mathbf {S}_t}(x_t)\Vert ^2_{\mathbf {B}}\\&= d_t -\frac{2\omega }{1-\beta } \langle x_t-x_*, \nabla f_{\mathbf {S}_t}(x_t) \rangle _{\mathbf {B}} - \frac{2 \omega \beta }{(1-\beta )^2} \langle x_t-x_{t-1}, \nabla f_{\mathbf {S}_t}(x_t) \rangle _{\mathbf {B}}\\&\quad + \frac{\omega ^2}{(1-\beta )^2} \Vert \nabla f_{\mathbf {S}_t}(x_t)\Vert ^2_{\mathbf {B}}. \end{aligned}$$

Taking expectation with respect to the random matrix $\mathbf {S}_t$ we obtain:

$$\begin{aligned} \mathbb {E}_{\mathbf{S}_t}[d_{t+1}]&= d_t -\frac{2\omega }{1-\beta } \langle x_t-x_*, \nabla f(x_t) \rangle _{\mathbf {B}} - \frac{2 \omega \beta }{(1-\beta )^2} \langle x_t-x_{t-1}, \nabla f(x_t) \rangle _{\mathbf {B}} \\&\quad + \frac{\omega ^2}{(1-\beta )^2} 2 f(x_t) \\&\overset{(48)}{=} d_t -\frac{4\omega }{1-\beta } f(x_t) - \frac{2 \omega \beta }{(1-\beta )^2} \langle x_t-x_{t-1}, \nabla f(x_t) \rangle _{\mathbf {B}} + \frac{\omega ^2}{(1-\beta )^2} 2 f(x_t)\\&\le d_t -\frac{4\omega }{1-\beta } f(x_t) - \frac{2 \omega \beta }{(1-\beta )^2} [f(x_t)-f(x_{t-1})] + \frac{\omega ^2}{(1-\beta )^2} 2 f(x_t)\\&= d_t + \left[ -\frac{4\omega }{1-\beta } - \frac{2 \omega \beta }{(1-\beta )^2} +\frac{2 \omega ^2}{(1-\beta )^2}\right] f(x_t) + \frac{2 \omega \beta }{(1-\beta )^2} f(x_{t-1}), \end{aligned}$$

where the inequality follows from convexity of f. After rearranging the terms we get

$$\begin{aligned} \mathbb {E}_{\mathbf{S}_t}[d_{t+1}] + \frac{2 \omega \beta }{(1-\beta )^2} f(x_t) + \alpha f(x_t) \le d_t + \frac{2 \omega \beta }{(1-\beta )^2} f(x_{t-1}), \end{aligned}$$

where $\alpha = \frac{4\omega }{1-\beta } -\frac{2 \omega ^2}{(1-\beta )^2} > 0$. Taking expectations again and using the tower property, we get

$$\begin{aligned} \theta _{t+1} + \alpha \mathbb {E}[f(x_t)] \le \theta _t, \qquad t=1,2,\dots , \end{aligned}$$

(59)

where $\theta _t = \mathbb {E}[d_t] + \frac{2 \omega \beta }{(1-\beta )^2}\mathbb {E}[ f(x_{t-1})]$. By summing up (59) for $t=1,\dots , k$ we get

$$\begin{aligned} \sum _{t=1}^k \mathbb {E}[f(x_t)] \le \frac{\theta _1-\theta _{k-1}}{\alpha } \le \frac{\theta _1}{\alpha }. \end{aligned}$$

(60)

Finally, using Jensen’s inequality, we get

$$\begin{aligned} \mathbb {E}[f(\hat{x}_k)] = \mathbb {E}\left[ f\left( \frac{1}{k}\sum _{t=1}^k x_t\right) \right] \le \mathbb {E}\left[ \frac{1}{k}\sum _{t=1}^k f(x_t)\right] = \frac{1}{k}\sum _{t=1}^k \mathbb {E}[f(x_t)] \overset{(60)}{\le } \frac{\theta _1}{\alpha k}. \end{aligned}$$

It remains to note that $\theta _1 = \Vert x_0-x_*\Vert _{\mathbf{B}}^2 + \frac{2\omega \beta }{(1-\beta )^2 }f(x_0).$

Appendix 4: Proof of Theorem 4

In the proof of Theorem 4 the following two lemmas are used.

Lemma 13

([78]) Assume exactness. Let $x\in \mathbb {R}^n$ and $x_* = \varPi _\mathcal {L}^\mathbf {B}(x)$. If $\lambda _i=0$, then $u_i^\top \mathbf {B}^{1/2} (x-x_*)=0$.

Lemma 14

([18, 21]) Consider the second degree linear homogeneous recurrence relation:

$$\begin{aligned} r_{k+1}= a_1r_k+a_2 r_{k-1} \end{aligned}$$

(61)

with initial conditions $r_0,r_1 \in \mathbb {R}$. Assume that the constant coefficients $a_1$ and $a_2$ satisfy the inequality $a_1^2 +4a_2<0$ (the roots of the characteristic equation $t^2-a_1t-a_2=0$ are imaginary). Then there are complex constants $C_0$ and $C_1$ (depending on the initial conditions $r_0$ and $r_1$) such that:

$$\begin{aligned} r_k=2 M^k (C_0 \cos ( \theta k) + C_1 \sin (\theta k)) \end{aligned}$$

where $M= \bigg (\sqrt{\frac{a_1^2}{4}+\frac{(-a_1^2-4a_2)}{4}} \bigg )=\sqrt{-a_2}$ and $\theta$ is such that $a_1=2 M \cos (\theta )$ and $\sqrt{-a_1^2-4a_2}=2 M \sin (\theta )$.

We can now turn to the proof of Theorem 4. Plugging in the expression for the stochastic gradient, mSGD can be written in the form

$$\begin{aligned} x_{k+1}= & {} x_k -\omega \nabla f_{\mathbf {S}_k}(x_k) + \beta (x_k - x_{k-1}) \nonumber \\&\overset{(10)}{=}&x_k- \omega {\mathbf {B}}^{-1} \mathbf {Z}_k(x_k-x_*) + \beta (x_k - x_{k-1}). \end{aligned}$$

(62)

Subtracting $x_*$ from both sides of (62), we get

$$\begin{aligned} x_{k+1}-x_*= & {} (\mathbf {I}- \omega {\mathbf {B}}^{-1} \mathbf {Z}_k)(x_k-x_*) + \beta (x_k -x_* +x_* - x_{k-1})\\= & {} \left( (1+\beta )\mathbf {I}- \omega {\mathbf {B}}^{-1} \mathbf {Z}_k\right) (x_k-x_*) - \beta (x_{k-1}-x_*). \end{aligned}$$

Multiplying the last identity from the left by $\mathbf {B}^{1/2}$, we get

$$\begin{aligned} \mathbf {B}^{1/2} (x_{k+1}-x_*)= & {} \left( (1+\beta )\mathbf {I}- \omega \mathbf {B}^{-1/2} \mathbf {Z}_k \mathbf {B}^{-1/2}\right) \mathbf {B}^{1/2}(x_{k} -x_*) \\&- \beta \mathbf {B}^{1/2}(x_{k-1}-x_*). \end{aligned}$$

Taking expectations, conditioned on $x_k$ (that is, the expectation is with respect to $\mathbf {S}_k$):

$$\begin{aligned} \mathbf {B}^{1/2} \mathbb {E}[x_{k+1} -x_* \;|\; x_k]= & {} \left( (1+\beta )\mathbf {I}- \omega \mathbf {B}^{-1/2} \mathbb {E}[\mathbf {Z}] \mathbf {B}^{-1/2}\right) \mathbf {B}^{1/2}(x_{k} -x_*) \\&- \beta \mathbf {B}^{1/2}(x_{k-1}-x_*) . \end{aligned}$$

Taking expectations again, and using the tower property, we get

$$\begin{aligned} \mathbf {B}^{1/2} \mathbb {E}[x_{k+1} -x_*]= & {} \mathbf {B}^{1/2}\mathbb {E}\left[ \mathbb {E}[x_{k+1} -x_* \;|\; x_k]\right] \\= & {} \left( (1+\beta )\mathbf {I}- \omega \mathbf {B}^{-1/2} \mathbb {E}[\mathbf {Z}] \mathbf {B}^{-1/2} \right) \mathbf {B}^{1/2} \mathbb {E}[x_{k} -x_*] \\&- \beta \mathbf {B}^{1/2} \mathbb {E}[x_{k-1}-x_*]. \end{aligned}$$

Plugging the eigenvalue decomposition ${\mathbf {U}}\varvec{\varLambda } {{\mathbf {U}}}^\top$ of the matrix $\mathbf {W}=\mathbf {B}^{-1/2} \mathbb {E}[\mathbf {Z}] \mathbf {B}^{-1/2}$ into the above, and multiplying both sides from the left by ${{\mathbf {U}}}^\top$, we obtain

$$\begin{aligned} {{\mathbf {U}}}^\top \mathbf {B}^{1/2} \mathbb {E}[x_{k+1} -x_*]&= {{\mathbf {U}}}^\top \left( (1+\beta )\mathbf {I}- \omega {\mathbf {U}}\varvec{\varLambda } {{\mathbf {U}}}^\top \right) \mathbf {B}^{1/2} \mathbb {E}[x_{k} -x_*]\nonumber \\&\quad \, - \beta {{\mathbf {U}}}^\top \mathbf {B}^{1/2} \mathbb {E}[x_{k-1}-x_*]. \end{aligned}$$

(63)

Let us define $s_k{:}{=}{{\mathbf {U}}}^\top \mathbf {B}^{1/2} \mathbb {E}[x_{k} -x_*] \in \mathbb {R}^n$. Then relation (63) takes the form of the recursion

$$\begin{aligned} s_{k+1}= [(1+\beta )\mathbf{I}- \omega \varvec{\varLambda } ] s_k - \beta s_{k-1}, \end{aligned}$$

which can be written in a coordinate-by-coordinate form as follows:

$$\begin{aligned} s_{k+1}^i= [(1+\beta ) - \omega \lambda _i ] s_k^i - \beta s_{k-1}^i \quad \text {for all} \quad i= 1,2,3,\ldots ,n, \end{aligned}$$

(64)

where $s_k^i$ indicates the ith coordinate of $s_k$.

We will now fix i and analyze recursion (64) using Lemma 14. Note that (64) is a second degree linear homogeneous recurrence relation of the form (61) with $a_1=1+\beta - \omega \lambda _i$ and $a_2=- \beta$. Recall that $0\le \lambda _i \le 1$ for all i. Since we assume that $0< \omega \le 1/\lambda _{\max }$, we know that $0\le \omega \lambda _i \le 1$ for all i. We now consider two cases:

• 1. $\lambda _i =0$.

  In this case, (64) takes the form:

  $$\begin{aligned} s_{k+1}^i=(1+\beta )s_k^i-\beta s_{k-1}^i. \end{aligned}$$

  (65)

  Applying Proposition 2, we know that $x_*=\varPi _\mathcal {L}^\mathbf {B}(x_0)=\varPi _\mathcal {L}^\mathbf {B}(x_1)$. Using Lemma 13 twice, once for $x=x_0$ and then for $x=x_1$, we observe that $s_0^i=u_i^\top \mathbf {B}^{1/2} (x_0-x_*)=0$ and $s_1^i=u_i^\top \mathbf {B}^{1/2} (x_1-x_*)=0$. Finally, in view of (65) we conclude that

  $$\begin{aligned} s_k^i=0 \quad \text {for all} \quad k\ge 0 . \end{aligned}$$

  (66)

• 2. $\lambda _i >0$.

  Since $0<\omega \lambda _i \le 1$ and $\beta \ge 0$, we have $1+\beta - \omega \lambda _i \ge 0$ and hence

  $$\begin{aligned} a_1^2+4a_2=(1+\beta -\omega \lambda _i)^2-4\beta \le (1+\beta -\omega \lambda _{\min }^+)^2-4\beta < 0, \end{aligned}$$

  where the last inequality can be shown to hold¹⁷ for $\left( 1-\sqrt{\omega \lambda _{\min }^+}\right) ^2< \beta < 1$. Applying Lemma 14 the following bound can be deduced

  $$\begin{aligned} s_k^i= & {} 2(-a_2)^{k/2} (C_0 \cos (\theta k) +C_1 \sin (\theta k)) \; \le \; 2 \beta ^{k/2} P_i, \end{aligned}$$

  (67)

  where $P_i$ is a constant depending on the initial conditions (we can simply choose $P_i = |C_0| + |C_1|$).

Now putting the two cases together, for all $k\ge 0$ we have

$$\begin{aligned} \Vert \mathbb {E}[x_{k} -x_*]\Vert _{\mathbf {B}}^2&= \mathbb {E}[x_{k} -x_*]^\top \mathbf {B}\mathbb {E}[x_{k} -x_*] \; = \; \mathbb {E}[x_{k} -x_*] \mathbf {B}^{1/2} \mathbf {U}{\mathbf {U}}^\top \mathbf {B}^{1/2} \mathbb {E}[x_{k} -x_*] \\&= \Vert {\mathbf {U}}^\top \mathbf {B}^{1/2} \mathbb {E}[x_{k} -x_*] \Vert _2^2 = \Vert s_k\Vert ^2 = \sum _{i=1}^{n} (s_k^i)^2 \\&= \sum _{i: \lambda _i=0} (s_k^i)^2 + \sum _{i: \lambda _i>0} (s_k^i)^2 \; \overset{(66)}{=}\; \sum _{i: \lambda _i>0} (s_k^i)^2\\&\overset{(67)}{\le } \sum _{i: \lambda _i >0} 4 \beta ^k P_i^2 \\&= \beta ^k C, \end{aligned}$$

where $C=4\sum _{i: \lambda _i >0} P_i^2$.

Appendix 5: Proof of Theorem 7

The proof follows a similar pattern to that of Theorem 1. However, stochasticity in the momentum term introduces an additional layer of complexity, which we shall tackle by utilizing a more involved version of the tower property.

For simplicity, let $i=i_k$ and $r_{k}^i {:}{=}e_i^\top (x_k-x_{k-1})e_i$. First, we decompose

$$\begin{aligned} \Vert x_{k+1}-x_*\Vert ^2= & {} \Vert x_k-\omega \nabla f_{\mathbf {S}_k}(x_k)+\gamma r_k^i -x_*\Vert ^2 \nonumber \\= & {} \Vert x_k-\omega \nabla f_{\mathbf {S}_k}(x_k)-x_*\Vert ^2+2\langle x_k-\omega \nabla f_{\mathbf {S}_k}(x_k)\nonumber \\&-x_*, \gamma r_k^i \rangle + \gamma ^2\Vert r_k^i\Vert ^2. \end{aligned}$$

(68)

We shall use the tower property in the form

$$\begin{aligned} \mathbb {E}[\mathbb {E}[\mathbb {E}[ X \;|\; x_k, \mathbf{S}_k] \;|\; x_k]] = \mathbb {E}[X], \end{aligned}$$

(69)

where X is some random variable. We shall perform the three expectations in order, from the innermost to the outermost. Applying the inner expectation to the identity (68), we get

(70)

We will now analyze the three expressions separately. The first expression is constant under the expectation, and hence we can write

(71)

We will now bound the second expression. Using the identity

$$\begin{aligned} \mathbb {E}[r_k^i \;|\; x_k, \mathbf{S}_k] = \mathbb {E}_i [r_k^i] = \sum _{i=1}^n \frac{1}{n}r_k^i = \frac{1}{n}(x_k-x_{k-1}), \end{aligned}$$

(72)

we can write

(73)

Using the fact that for arbitrary vectors $a,b,c \in \mathbb {R}^n$ we have the identity $2 \langle a-c,c-b \rangle =\Vert a-b\Vert ^2-\Vert c-b\Vert ^2-\Vert a-c\Vert ^2,$ we obtain

$$\begin{aligned} 2 \langle x_k-x_*, x_*-x_{k-1} \rangle = \Vert x_k-x_{k-1}\Vert ^2- \Vert x_{k-1}-x_*\Vert ^2-\Vert x_k-x_*\Vert ^2. \end{aligned}$$

Substituting this into (73) gives

(74)

The third expression can be bound as

(75)

By substituting the bounds (71), (74), (75) into (70) we obtain

$$\begin{aligned} \mathbb {E}[\Vert x_{k+1}-x_*\Vert ^2 \;|\; x_k, \mathbf{S}_k ]&\le \Vert x_k-x_*\Vert ^2-2\omega (2-\omega ) f_{\mathbf {S}_k}(x_k)\nonumber \\&\quad + \tfrac{\gamma }{n} \Vert x_k-x_*\Vert ^2+ \tfrac{\gamma }{n} \Vert x_{k}-x_{k-1}\Vert ^2 -\tfrac{\gamma }{n} \Vert x_{k-1}-x_*\Vert ^2 \nonumber \\&\quad + 2\omega \tfrac{\gamma }{n} \langle \nabla f_{\mathbf {S}_k}(x_k), x_{k-1}- x_k \rangle + 2\tfrac{\gamma ^2}{n} \Vert x_{k}-x_*\Vert ^2 \nonumber \\&\quad + 2\tfrac{\gamma ^2}{n}\Vert x_{k-1}-x_*\Vert ^2 \nonumber \\&\overset{(56)}{\le } \left( 1+3\tfrac{\gamma }{n} + 2\tfrac{\gamma ^2}{n}\right) \Vert x_k-x_*\Vert ^2+ \left( \tfrac{\gamma }{n} + 2\tfrac{\gamma ^2}{n} \right) \Vert x_{k-1}-x_*\Vert ^2 \nonumber \\&\quad - 2\omega (2-\omega )f_{\mathbf {S}_k}(x_k) + 2\omega \tfrac{\gamma }{n} \langle \nabla f_{\mathbf {S}_k}(x_k),x_{k-1}- x_k \rangle . \end{aligned}$$

(76)

We now take the middle expectation (see (69)) and apply it to inequality (76):

$$\begin{aligned}&\mathbb {E}[\mathbb {E}[\Vert x_{k+1}-x_*\Vert ^2 \;|\; x_k, \mathbf{S}_k ] \;|\; x_k] \\&\quad \le \left( 1+3\tfrac{\gamma }{n} + 2\tfrac{\gamma ^2}{n}\right) \Vert x_k-x_*\Vert ^2+ \left( \tfrac{\gamma }{n} + 2\tfrac{\gamma ^2}{n} \right) \Vert x_{k-1}-x_*\Vert ^2 \\&\qquad -2\omega (2-\omega )f(x_k) + 2\omega \tfrac{\gamma }{n} \langle \nabla f(x_k),x_{k-1}- x_k \rangle \\&\quad \le \left( 1+3\tfrac{\gamma }{n} + 2\tfrac{\gamma ^2}{n}\right) \Vert x_k-x_*\Vert ^2+ \left( \tfrac{\gamma }{n} + 2\tfrac{\gamma ^2}{n} \right) \Vert x_{k-1}-x_*\Vert ^2 \\&\qquad -2\omega (2-\omega )f(x_k) + 2\omega \tfrac{\gamma }{n}(f(x_{k-1})-f(x_k))\\&\quad = \left( 1+3\tfrac{\gamma }{n} + 2\tfrac{\gamma ^2}{n}\right) \Vert x_k-x_*\Vert ^2+ \left( \tfrac{\gamma }{n} + 2\tfrac{\gamma ^2}{n} \right) \Vert x_{k-1}-x_*\Vert ^2 \\&\qquad - \left( 2\omega (2-\omega ) +2\omega \tfrac{\gamma }{n}\right) f(x_k) + 2\omega \tfrac{\gamma }{n} f(x_{k-1}). \end{aligned}$$

where in the second step we used the inequality $\langle \nabla f(x_k),x_{k-1}- x_k \rangle \le f(x_{k-1})-f(x_k)$ and the fact that $\omega \gamma \ge 0$, which follows from the assumptions. We now apply inequalities (50) and (51), obtaining

$$\begin{aligned}&\mathbb {E}[\mathbb {E}[\Vert x_{k+1}-x_*\Vert ^2 \;|\; x_k, \mathbf{S}_k ] \;|\; x_k]\\&\quad \le \underbrace{\left( 1+3\tfrac{\gamma }{n}+2\tfrac{\gamma ^2}{n} - \left( \omega (2-\omega ) +\omega \tfrac{\gamma }{n}\right) \lambda _{\min }^+ \right) }_{a_1}\Vert x_k-x_*\Vert ^2 \\&\qquad + \underbrace{\tfrac{1}{n}\left( \gamma +2\gamma ^2 + \omega \gamma \lambda _{\max }\right) }_{a_2}\Vert x_{k-1}-x_*\Vert ^2. \end{aligned}$$

By taking expectation again (outermost expectation in the tower rule (69)), and letting $F_k{:}{=}\mathbb {E}[\Vert x_{k}-x_*\Vert ^2_{\mathbf {B}}]$, we get the relation

$$\begin{aligned} F_{k+1} \le a_1 F_k + a_2 F_{k-1} . \end{aligned}$$

(77)

It suffices to apply Lemma 9 to the relation (57). The conditions of the lemma are satisfied. Indeed, $a_2\ge 0$, and if $a_2=0$, then $\gamma =0$ and hence $a_1=1-\omega (1-\omega )\lambda _{\min }^+>0$. The condition $a_1+a_2<1$ holds by assumption.

The convergence result in function values follows as a corollary by applying inequality (50) to (32).

Appendix 6: Notation glossary

For the frequently used notation, see Table 8.

Table 8 Frequently used notation