SGEM: stochastic gradient with energy and momentum

Abstract

In this paper, we propose SGEM, stochastic gradient with energy and momentum, to solve a class of general non-convex stochastic optimization problems, based on the AEGD method introduced in AEGD (adaptive gradient descent with energy) Liu and Tian (Numerical Algebra, Control and Optimization, 2023). SGEM incorporates both energy and momentum so as to inherit their dual advantages. We show that SGEM features an unconditional energy stability property and provide a positive lower threshold for the energy variable. We further derive energy-dependent convergence rates in the general non-convex stochastic setting, as well as a regret bound in the online convex setting. Our experimental results show that SGEM converges faster than AEGD and generalizes better or at least as well as SGDM in training some deep neural networks.

Appendices

Appendix 1. Proof of Theorem 2

For the proofs of Theorems 2 and 3, we introduce notation

$$\begin{aligned} {\tilde{F}}_t:=\sqrt{f(\theta _t; \xi _t)+c}. \end{aligned}$$

(A.1)

The initial data for $r_i$ is taken as $r_{1, i}={\tilde{F}}_1$. We also denote the update rule presented in Algorithm 1 as

$$\begin{aligned} \theta _{t+1}=\theta _{t}-2\eta r_{t+1} v_{t}, \end{aligned}$$

(A.2)

where $r_{t+1}$ is viewed as a $n\times n$ diagonal matrix that is made up of $[r_{t+1,1},$ $\ldots ,r_{t+1,i},\ldots ,r_{t+1,n}]$.

Lemma 1

Under the assumptions in Theorem 2, we have for all $t\in [T]$,

• (i) $\Vert \nabla f(\theta _t)\Vert _\infty \le G_\infty$.

• (ii) $\mathbb {E}[({\tilde{F}}_t)^2]= F^2(\theta _t)=f(\theta _t)+c$.

• (iii) $\mathbb {E}[{\tilde{F}}_t]\le F(\theta _t)$. In particular, $\mathbb {E}[r_{1,i}]= \mathbb {E}[{\tilde{F}}_1]\le F(\theta _1)$ for all $i\in [n]$.

• (iv) $\sigma ^2_g=\mathbb {E}[\Vert g_t-\nabla f(\theta _t)\Vert ^2]\le G^2_\infty$ and $\sigma ^2_f=\mathbb {E}[|f(\theta _t;\xi _t)- f(\theta _t)|^2]\le B^2.$

• (v) $\mathbb {E}[|F(\theta _t)-{\tilde{F}}_t|]\le \frac{1}{2\sqrt{a}}\sigma _f$.

• (vi) $\mathbb {E}[\Vert \nabla F(\theta _t)-\frac{g_t}{2\tilde{F}_t}\Vert ^2]\le \frac{G^2_\infty }{8a^3}\sigma ^2_f+\frac{1}{2a}\sigma ^2_g.$

Proof

• (i) By assumption $\Vert g_t\Vert _\infty \le G_\infty$, we have

  $$\begin{aligned} \Vert \nabla f(\theta _t)\Vert _\infty =\Vert \mathbb {E}[g_t]\Vert _\infty \le \mathbb {E}[\Vert g_t\Vert _\infty ]\le G_\infty . \end{aligned}$$

• (ii) This follows from the unbiased sampling of

  $$\begin{aligned} f(\theta _t)=\mathbb {E}_{\xi _t}[ f(\theta _t; \xi _t)]. \end{aligned}$$

• (iii) By Jensen’s inequality, we have

  $$\begin{aligned} \mathbb {E}[\tilde{F}_t] \le \sqrt{\mathbb {E}[\tilde{F}_t^2]}=\sqrt{F(\theta _t)^2}=F(\theta _t). \end{aligned}$$

• (iv) By assumptions $\Vert g_t\Vert _\infty \le G_\infty$ and $f(\theta _t;\xi _t)+c<B$, we have

  $$\begin{aligned} \sigma ^2_g=\mathbb {E}[\Vert g_t-\nabla f(\theta _t)\Vert ^2] = \mathbb {E}[\Vert g_t\Vert ^2] - \Vert \nabla f(\theta _t)\Vert ^2\le G^2_\infty , \end{aligned}$$
  $$\begin{aligned} \sigma ^2_f=\mathbb {E}[\Vert f(\theta _t;\xi _t)-f(\theta _t)\Vert ^2] = \mathbb {E}[\Vert f(\theta _t;\xi _t)\Vert ^2] - \Vert f(\theta _t)\Vert ^2\le B^2. \end{aligned}$$

• (v) By the assumption $0<a\le f(\theta _t;\xi _t)+c=\tilde{F}_t^2$, we have

  $$\begin{aligned} \quad \mathbb {E}[|F(\theta _t)-\tilde{F}_t|] \le \mathbb {E}\Bigg [\bigg |\frac{f(\theta _t)-f(\theta _t;\xi _t)}{F(\theta _t)+\tilde{F}_t}\bigg |\Bigg ] \le \frac{1}{2\sqrt{a}}\mathbb {E}[|f(\theta _t)-f(\theta _t;\xi _t)|] \le \frac{1}{2\sqrt{a}}\sigma _f. \end{aligned}$$

• (vi) By the definition of $F(\theta )$, we have

  $$\begin{aligned} \Vert \nabla F(\theta _t)-\frac{g_t}{2\tilde{F}_t}\Vert ^2= & {} \bigg \Vert \frac{\nabla f(\theta _t)}{2F(\theta _t)}-\frac{g_t}{2\tilde{F}_t}\bigg \Vert ^2\\= & {} \frac{1}{4}\bigg \Vert \frac{\nabla f(\theta _t)(\tilde{F}_t-F(\theta _t)) }{F(\theta _t)\tilde{F}_t}+\frac{\nabla f(\theta _t)-g_t}{\tilde{F}_t}\bigg \Vert ^2\\\le & {} \frac{1}{2} \bigg \Vert \frac{\nabla f(\theta _t)(\tilde{F}_t-F(\theta _t)) }{F(\theta _t)\tilde{F}_t}\bigg \Vert ^2 + \frac{1}{2}\bigg \Vert \frac{\nabla f(\theta _t)-g_t}{\tilde{F}_t}\bigg \Vert ^2\\\le & {} \frac{G^2_\infty }{2a^{2}}|\tilde{F}_t-F(\theta _t)|^2+\frac{1}{2a}\Vert \nabla f(\theta _t)-g_t\Vert ^2, \end{aligned}$$

  where both the gradient bound and the assumption that $0<a\le f(\theta _t;\xi _t)+c=\tilde{F}^2_t$ are essentially used. Take an expectation to get

  $$\begin{aligned} \mathbb {E}[\Vert \nabla F(\theta _t)-\frac{g_t}{2\tilde{F}_t}\Vert ^2]\le \frac{G^2_\infty }{2a^{2}}\mathbb {E}[|\tilde{F}_t-F(\theta _t)|^2]+\frac{1}{2a}\mathbb {E}[\Vert \nabla f(\theta _t)-g_t\Vert ^2]. \end{aligned}$$

  Similar to the proof for (iv), we have

  $$\begin{aligned} \mathbb {E}[|\tilde{F}_t-F(\theta _t)|^2]\le \frac{1}{4a}\sigma ^2_f. \end{aligned}$$

  This together with the variance assumption for $g_t$ gives

  $$\begin{aligned} \mathbb {E}[\Vert \nabla F(\theta _t)-\frac{g_t}{2\tilde{F}_t}\Vert ^2]\le \frac{G^2_\infty }{8a^3}\sigma ^2_f+\frac{1}{2a}\sigma ^2_g. \end{aligned}$$

Lemma 2

For any $T\ge 1$, we have

• (i) $\mathbb {E}\Big [\sum _{t=1}^{T} v_t^\top r_{t+1} v_t\Big ] \le \frac{n F(\theta _1)}{2\eta }$.

• (ii) $\mathbb {E}\Big [\sum _{t=1}^{T} m_{t-1}^\top r_{t+1} m_{t-1}\Big ]\le \mathbb {E}\Big [\sum _{t=1}^{T} m_t^\top r_{t+1} m_t\Big ] \le \frac{2Bn F(\theta _1)}{\eta }$.

• (iii) $\mathbb {E}\Big [\sum _{t=1}^{T}\Vert r_{t+1}m_t\Vert ^2\Big ] \le \frac{2Bn F^2(\theta _1)}{\eta }$.

• (iv) $\mathbb {E}\Big [\sum _{t=1}^{T} g_t^\top r_{t+1} g_t\Big ] \le \frac{8Bn F(\theta _1)}{(1-\beta )^2\eta }$.

• (v) $\mathbb {E}\Big [\sum _{t=1}^{T}\Vert r_{t+1}g_t\Vert ^2\Big ] \le \frac{8Bn F^2(\theta _1)}{(1-\beta )^2\eta }$.

Proof

From Algorithm 1 line 5, we have

$$\begin{aligned} r_{t,i}-r_{t+1,i} = 2\eta r_{t+1,i}v^2_{t,i}. \end{aligned}$$

Taking summation over t from 1 to T gives

$$\begin{aligned} r_{1,i} - r_{T+1,i} = 2\eta \sum \limits _{t=1}^{T} r_{t+1,i}v^2_{t,i} \quad \Rightarrow \quad \sum \limits _{t=1}^{T} r_{t+1,i}v^2_{t,i}\le \frac{r_{1,i}}{2\eta }. \end{aligned}$$

From which we get

$$\begin{aligned} \sum _{t=1}^{T}v_t^\top r_{t+1}v_t=\sum \limits _{i=1}^{n}\sum \limits _{t=1}^{T}r_{t+1,i}v^2_{t,i}\le \frac{n\tilde{F}_1}{2\eta }. \end{aligned}$$

Taking expectation and using (iii) in Lemma 1 gives (i). Recall that $m_t=2(1-\beta ^t)\tilde{F}_tv_t$ and $\tilde{F}_t\le \sqrt{B}$, we further get

$$\begin{aligned} \sum _{t=1}^{T} m_t^\top r_{t+1} m_t \le 4B \sum \limits _{t=1}^{T} v_t^\top r_{t+1}v_t = \frac{2Bn\tilde{F}_1}{\eta }. \end{aligned}$$

Using $r_{t+1,i}\le r_{t,i}$ and $m_{0,i}=0$, we also have

$$\begin{aligned} \sum \limits _{i=1}^{n}\sum \limits _{t=1}^{T}r_{t+1,i}m^2_{t-1,i} \le \sum \limits _{i=1}^{n}\sum \limits _{t=1}^{T}r_{t,i}m^2_{t-1,i} = \sum \limits _{i=1}^{n}\sum \limits _{t=1}^{T-1}r_{t+1,i}m^2_{t,i}\le \sum \limits _{i=1}^{n}\sum \limits _{t=1}^{T}r_{t+1,i}m^2_{t,i}. \end{aligned}$$

(A.3)

Connecting the above two inequalities and taking expectation gives (ii). Using $r_{t+1,i}\le r_{1,i}$, the above inequality further implies

$$\begin{aligned} \sum _{t=1}^{T}\Vert r_{t+1}m_t\Vert ^2= & {} \sum \limits _{i=1}^n\sum \limits _{t=1}^{T} r^2_{t+1, i}m_{t, i}^2 \le \sum \limits _{i=1}^n\sum \limits _{t=1}^{T} r_{1,i} r_{t+1, i}m_{t, i}^2 \\= & {} \bigg (\sum \limits _{i=1}^n\sum \limits _{t=1}^{T} r_{t, i}m_{t, i}^2\bigg )\tilde{F}_1 \le 2Bn\tilde{F}^2_1/\eta . \end{aligned}$$

Taking expectation and using (ii) in Lemma 1 gives (iii). By $m_t=\beta m_{t-1}+(1-\beta )g_t$, we have

$$\begin{aligned} \sum \limits _{t=1}^{T} g_t^\top r_{t+1} g_t= & {} \sum \limits _{i=1}^{n}\sum \limits _{t=1}^{T}r_{t+1,i}g^2_{t,i} = \sum \limits _{i=1}^{n}\sum \limits _{t=1}^{T}r_{t+1,i}\bigg (\frac{1}{1-\beta }m_{t,i}-\frac{\beta }{1-\beta }m_{t-1,i}\bigg )^2\\\le & {} \frac{2}{(1-\beta )^2}\sum \limits _{i=1}^{n}\sum \limits _{t=1}^{T}r_{t+1,i}m^2_{t,i} + \frac{2\beta ^2}{(1-\beta )^2}\sum \limits _{i=1}^{n}\sum \limits _{t=1}^{T}r_{t+1,i}m^2_{t-1,i}\\\le & {} \frac{2(1+\beta ^2)}{(1-\beta )^2}\sum \limits _{t=1}^{T} m_t^\top r_{t+1} m_t\le \frac{8Bn\tilde{F}_1}{(1-\beta )^2\eta }. \end{aligned}$$

Here, the third inequality is by $(a+b)^2\le 2a^2+2b^2$; (A.3) and $0<\beta <1$ are used in the fourth inequality. Taking expectation and using (iii) in Lemma 1 gives (iv). Similar as the derivation for (ii), we have

$$\begin{aligned} \sum \limits _{t=1}^{T}\Vert r_{t+1}g_t\Vert ^2 \le \bigg (\sum \limits _{i=1}^n\sum _{t=1}^{T} r_{t, i}g_{t, i}^2\bigg )\tilde{F}_1 \le \frac{8Bn\tilde{F}^2_1}{(1-\beta )^2\eta }. \end{aligned}$$

Taking expectation and using (ii) in Lemma 1 gives (v).$\square$

First, note that by (iv) in Lemma 1, $\max \{\sigma _g,\sigma _f\}\le \max \{G_\infty ,B\}$.

Recall that $F(\theta )=\sqrt{f(\theta )+c}$, then for any $x, y\in \{\theta _t\}_{t=0}^T$ we have

$$\begin{aligned} \Vert \nabla F(x)-\nabla F(y)\Vert= & {} \bigg \Vert \frac{\nabla f(x)}{2F(x)}-\frac{\nabla f(y)}{2F(y)}\bigg \Vert \\= & {} \frac{1}{2}\bigg \Vert \frac{\nabla f(x)(F(y)-F(x))}{F(x)F(y)} + \frac{\nabla f(x)-\nabla f(y)}{F(y)}\bigg \Vert \\\le & {} \frac{G_\infty }{2(F(\theta ^*))^2}|F(y)-F(x)| + \frac{1}{2F(\theta ^*)}\Vert \nabla f(x)-\nabla f(y)\Vert . \end{aligned}$$

One may check that

$$\begin{aligned} |F(y)-F(x)|\le \frac{G_\infty }{2F(\theta ^*)}\Vert x-y\Vert . \end{aligned}$$

These together with the L-smoothness of f lead to

$$\begin{aligned} \Vert \nabla F(x)-\nabla F(y)\Vert \le L_F \Vert x-y\Vert , \end{aligned}$$

where

$$\begin{aligned} L_F=\frac{1}{2\sqrt{f(\theta ^*)+c}} \left( L+ \frac{G^2_\infty }{2(f(\theta ^*)+c)}\right) . \end{aligned}$$

This confirms the $L_F$-smoothness of F, which yields

$$\begin{aligned} F(\theta _{t+1}) - F(\theta _t)\le & {} \nabla F(\theta _t)^\top (\theta _{t+1}-\theta _t) +\frac{L_F}{2}\Vert \theta _{t+1}-\theta _t\Vert ^2 \\= & {} (\nabla F(\theta _t)-\frac{g_t}{2\tilde{F}_t})^\top (\theta _{t+1}-\theta _t) + (\frac{g_t}{2\tilde{F}_t}-\frac{1-\beta ^t}{1-\beta }v_t)^\top (\theta _{t+1}-\theta _t)\\{} & {} \quad +(\frac{1-\beta ^t}{1-\beta }v_t)^\top (\theta _{t+1}-\theta _t) +\frac{L_F}{2}\Vert \theta _{t+1}-\theta _t\Vert ^2. \\ \end{aligned}$$

Summation of the above over t from 1 to T and taken with the expectation gives

$$\begin{aligned} \mathbb {E}[F(\theta _{T+1})-F(\theta _{1})]\le \sum \limits _{i=1}^{4} S_i, \end{aligned}$$

(A.4)

where

$$\begin{aligned} S_1= & {} \mathbb {E}\Bigg [\sum _{t=1}^{T}\frac{1-\beta ^t}{1-\beta }v_{t}^\top (\theta _{t+1}-\theta _t)\Bigg ],\\ S_2= & {} \mathbb {E}\Bigg [\sum _{t=1}^{T}(\frac{g_t}{2\tilde{F}_t}-\frac{1-\beta ^t}{1-\beta }v_t)^\top (\theta _{t+1}-\theta _t)\Bigg ],\\ S_3= & {} \mathbb {E}\Bigg [\sum _{t=1}^{T}(\nabla F(\theta _t)-\frac{g_t}{2\tilde{F}_t})^\top (\theta _{t+1}-\theta _t)\Bigg ],\\ S_4= & {} \mathbb {E}\Bigg [\sum _{t=1}^{T}\frac{L_F}{2}\Vert \theta _{t+1}-\theta _t\Vert ^2\Bigg ]. \end{aligned}$$

Below, we bound $S_1, S_2, S_3, S_4$ separately. To bound $S_1$, we first note that

$$\begin{aligned} r_{t+1,i}-r_{t,i}=-2\eta r_{t+1,i}v^2_{t,i} = v_{t,i}(-2\eta r_{t+1,i}v_{t,i})=v_{t,i}(\theta _{t+1,i}-\theta _i) \end{aligned}$$

from which we get

$$\begin{aligned} S_1= & {} \mathbb {E}\Bigg [\sum _{t=1}^{T}\frac{1-\beta ^{t}}{1-\beta }v_{t}^{\top } (\theta _{t+1}-\theta _{t})\Bigg ]\\= & {} \mathbb {E}\Bigg [\sum _{i=1}^{n}\sum _{t=1}^{T} \frac{1-\beta ^t}{1-\beta } (r_{t+1,i}-r_{t,i})\Bigg ]\\\le & {} \mathbb {E}\Bigg [\sum _{i=1}^{n}\sum _{t=1}^{T} r_{t+1,i}-r_{t,i}\Bigg ]\quad (\text {Since}~ r_{t+1,i}\le r_{t,i})\\= & {} \sum _{i=1}^{n}\mathbb {E}[r_{T+1,i}]-n\mathbb {E}[\tilde{F}_1]. \end{aligned}$$

For $S_2$, we have

$$\begin{aligned} S_2= & {} \mathbb {E}\Bigg [\sum \limits _{t=1}^{T}(\frac{g_t}{2\tilde{F}_t}-\frac{1-\beta ^t}{1-\beta }v_t)^\top (\theta _{t+1}-\theta _t)\Bigg ]\\= & {} \mathbb {E}\Bigg [\sum \limits _{i=1}^{n}\sum \limits _{t=1}^{T}(-\frac{1}{2\tilde{F}_t}\frac{\beta }{1-\beta }m_{t-1,i})^\top (-2\eta r_{t+1,i}v_{t,i})\Bigg ]\\\le & {} \frac{\beta \eta }{(1-\beta )\sqrt{a}}\mathbb {E}\Bigg [\left|\sum \limits _{i=1}^{n}\sum \limits _{t=1}^{T}r_{t+1,i}m_{t-1,i}v_{t,i}\right|\Bigg ]\\\le & {} \frac{\beta \eta }{(1-\beta )\sqrt{a}}\mathbb {E}\Bigg [\bigg (\sum \limits _{i=1}^{n}\sum \limits _{t=1}^{T}r_{t+1,i}m^2_{t-1,i}\bigg )^{1/2}\bigg (\sum \limits _{i=1}^{n}\sum \limits _{t=1}^{T}r_{t+1,i}v^2_{t,i}\bigg )^{1/2}\Bigg ]\\\le & {} \frac{\beta \sqrt{B}nF(\theta _1)}{(1-\beta )\sqrt{a}}, \end{aligned}$$

where the fourth inequality is by the Cauchy-Schwarz inequality, and the last inequality is by Lemma 1 (i) (ii).

For $S_3$, by the Cauchy-Schwarz inequality, we have

$$\begin{aligned} S_3= & {} \mathbb {E}\Bigg [\sum \limits _{t=1}^{T}(\nabla F(\theta _t)-\frac{g_t}{2\tilde{F}_t})^\top (\theta _{t+1}-\theta _t)\Bigg ]\\\le & {} \mathbb {E}\Bigg [\sum \limits _{t=1}^{T}\Vert \nabla F(\theta _t)-\frac{g_t}{2\tilde{F}_t})\Vert \Vert \theta _{t+1}-\theta _t)\Vert \Bigg ]\\\le & {} \mathbb {E}\Bigg [\bigg (\sum \limits _{t=1}^{T}\Vert \nabla F(\theta _t)-\frac{g_t}{2\tilde{F}_t})\Vert ^2\bigg )^{1/2}\bigg (\sum \limits _{t=1}^{T}\Vert \theta _{t+1}-\theta _t\Vert ^2\bigg )^{1/2}\Bigg ]\\\le & {} \Bigg (\mathbb {E}\Bigg [\sum \limits _{t=1}^{T}\Vert \nabla F(\theta _t)-\frac{g_t}{2\tilde{F}_t})\Vert ^2\Bigg ]\Bigg )^{1/2}\Bigg (\mathbb {E}\Bigg [\sum \limits _{t=1}^{T}\Vert \theta _{t+1}-\theta _t\Vert ^2\Bigg ]\Bigg )^{1/2}\\\le & {} F(\theta _1)\sqrt{\eta n T}\sqrt{\frac{G^2_\infty }{8a^3}\sigma ^2_f+\frac{1}{2a}\sigma ^2_g}, \end{aligned}$$

where the last inequality is by (vi) in Lemma 1 and (4.2) in Theorem 1.

For $S_4$, also by (4.2) in Theorem 1, we have

$$\begin{aligned} S_4 = \frac{L_F}{2}\mathbb {E}\Bigg [\sum _{t=1}^{T}\Vert \theta _{t+1}-\theta _t\Vert ^2\Bigg ] \le \frac{L_F\eta n F^2(\theta _1)}{2} . \end{aligned}$$

With the above bounds on $S_1, S_2, S_3, S_4$, (A.4) can be rearranged as

$$\begin{aligned}{} & {} \quad F(\theta ^*) - \frac{\beta \sqrt{B}nF(\theta _1)}{(1-\beta )\sqrt{a}} -F(\theta _1)\sqrt{\eta n T}\sqrt{\frac{G^2_\infty }{4a^3}\sigma ^2_f+\frac{1}{a}\sigma ^2_g}- \frac{L_F\eta n F^2(\theta _1)}{2}\\{} & {} \le \sum \limits _{i=1}^{n}\mathbb {E}[r_{T+1,i}] - n \mathbb {E}[\tilde{F}_1] + F(\theta _1) \\{} & {} \le \Big (\min \limits _i \mathbb {E}[r_{T+1,i}]+(n-1)\mathbb {E}[\tilde{F}_1]\Big )- (n-1)\mathbb {E}[\tilde{F}_1] + \Big (F(\theta _1)-\mathbb {E}[\tilde{F}_1]\Big )\\{} & {} \le \min \limits _i\mathbb {E}[r_{T+1,i}]+ \mathbb {E}[|F(\theta _1)-\tilde{F}_1|] \\{} & {} \le \min \limits _i\mathbb {E}[r_{T+1,i}]+ \frac{1}{2\sqrt{a}}\sigma _f, \end{aligned}$$

where (iii) in Lemma 1 was used. Hence,

$$\begin{aligned} \min \limits _i\mathbb {E}[r_{T,i}]\ge \max \{F(\theta ^*)-\eta D_1-\beta D_2-\sigma D_3,0\}, \end{aligned}$$

where $\sigma =\max \{\sigma _f,\sigma _g\}$ and

$$\begin{aligned}{} & {} D_1 = \frac{L_F n F^2(\theta _1)}{2}, \quad D_2 =\frac{\sqrt{B}nF(\theta _1)}{(1-\beta )\sqrt{a}},\\{} & {} D_3 = \frac{1}{2\sqrt{a}} + F(\theta _1)\sqrt{\eta n T}\sqrt{\frac{G^2_\infty }{4a^3}+\frac{1}{a}}. \end{aligned}$$

In the case $\sigma =0$, we obtain the stated estimate in Theorem 2.

Appendix 2. Proof of Theorem 3

The upper bound on $\sigma _g$ is given by (iv) in Lemma 1. Since f is L-smooth, we have

$$\begin{aligned} f(\theta _{t+1})\le f(\theta _t)+\nabla f(\theta _t)^\top (\theta _{t+1}-\theta _t)+\frac{L}{2}\Vert \theta _{t+1}-\theta _t\Vert ^2. \end{aligned}$$

(B.1)

Denoting $\eta _t=\eta /\tilde{F}_t$, the second term in the RHS of (B.1) can be expressed as

$$\begin{aligned}{} & {} \quad \nabla f(\theta _t)^\top (\theta _{t+1}-\theta _t)\nonumber \\= & {} \nabla f(\theta _t)^\top (-2\eta r_{t+1}v_{t})\nonumber \\= & {} -\frac{1}{1-\beta ^t}\nabla f(\theta _t)^\top \eta _t r_{t+1}m_{t}\quad (\text {since}\;m_t=2(1-\beta ^t)\tilde{F}_t v_t)\nonumber \\= & {} -\frac{1}{1-\beta ^t}\nabla f(\theta _t)^\top \eta _t r_{t+1}(\beta m_{t-1}+(1-\beta )g_t) \nonumber \\ \end{aligned}$$

$$\begin{aligned}= & {} -\frac{1-\beta }{1-\beta ^t}\nabla f(\theta _t)^\top \eta _t r_{t+1}g_{t} - \frac{\beta }{1-\beta ^t}\nabla f(\theta _t)^\top \eta _t r_{t+1}m_{t-1}\nonumber \\= & {} -\frac{1-\beta }{1-\beta ^t}\nabla f(\theta _t)^\top \eta _{t-1} r_{t}g_{t} + \frac{1-\beta }{1-\beta ^t}\nabla f(\theta _t)^\top (\eta _{t-1} r_{t}-\eta _t r_{t+1}) g_{t}\nonumber \\{} & {} \quad -\frac{\beta }{1-\beta ^t}\nabla f(\theta _t)^\top \eta _t r_{t+1}m_{t-1}. \end{aligned}$$

(B.2)

We further bound the second term and third term in the RHS of (B.2), respectively. For the second term, we note that $|\frac{1-\beta }{1-\beta ^t}|\le 1$ and

$$\begin{aligned}&\quad |&\!\!\nabla f(\theta _t)^\top (\eta _{t-1}r_{t}-\eta _tr_{t+1}) g_{t}|\nonumber \\= & {} |\nabla f(\theta _t)^\top \eta _{t-1}(r_{t}-r_{t+1}) g_{t}+ \nabla f(\theta _t)^\top (\eta _{t-1}-\eta _t)r_{t+1}g_{t}|\nonumber \\= & {} |\nabla f(\theta _t)^\top \eta _{t-1}(r_{t}-r_{t+1}) g_{t}+ (\eta _{t-1}-\eta _t)g_t^\top r_{t+1}g_{t}\nonumber \\{} & {} \!\!\!\!\!\!\quad + (\eta _{t-1}-\eta _t)(\nabla f(\theta _t)-g_t)^\top r_{t+1}g_{t}|\nonumber \\\le & {} \Vert \nabla f(\theta _t)\Vert _\infty |\eta _{t-1}|\Vert r_{t}-r_{t+1}\Vert _{1,1} \Vert g_{t}\Vert _\infty + |\eta _{t-1}-\eta _t|g_t^\top r_{t+1}g_{t}\nonumber \\{} & {} \!\!\!\!\!\!\quad + |\eta _{t-1}-\eta _t||(\nabla f(\theta _t)-g_t)^\top r_{t+1}g_{t}|\nonumber \\\le & {} (\eta G^2_\infty /\sqrt{a})(\Vert r_t\Vert _{1,1}-\Vert r_{t+1}\Vert _{1,1})+(2\eta /\sqrt{a})g_t^\top r_{t+1}g_{t}\nonumber \\{} & {} \!\!\!\!\!\!\quad + (2\eta /\sqrt{a})|(\nabla f(\theta _t)-g_t)^\top r_{t+1}g_{t}|. \end{aligned}$$

(B.3)

The third inequality holds because for a positive diagonal matrix A, $x^\top Ay\le \Vert x\Vert _\infty \Vert A\Vert _{1,1}\Vert y\Vert _\infty$, where $\Vert A\Vert _{1,1}=\sum _{i}a_{ii}$. The last inequality follows from the result $r_{t+1,i}\le r_{t,i}$ for $i\in [n]$, the assumption $\Vert g_t\Vert _\infty \le G_\infty$, $\tilde{F}_t\ge \sqrt{a}$, and (i) in Lemma (1).

For the third term in the RHS of (B.2), we note that

$$\begin{aligned} -\frac{\beta }{1-\beta ^t} \nabla f(\theta _t)^\top \eta _t r_{t+1}m_{t-1} \le \frac{\beta \eta }{(1-\beta )\sqrt{a}}|\nabla f(\theta _t)^\top \eta _t r_{t+1}m_{t-1}|, \end{aligned}$$

in which

$$\begin{aligned}{} & {} \!\!\!\!\!\!\quad |\nabla f(\theta _t)^\top r_{t+1}m_{t-1}|\nonumber \\= & {} |g_{t}^\top r_{t+1}m_{t-1}+(\nabla f(\theta _t)-g_t)^\top r_{t+1}m_{t-1}|\nonumber \\\le & {} \frac{1}{2}g_t^\top r_{t+1}g_t+\frac{1}{2}m_{t-1}^\top r_{t+1}m_{t-1}+ |(\nabla f(\theta _t)-g_t)^\top r_{t+1}m_{t-1}|, \end{aligned}$$

(B.4)

where the last inequality is because for a positive diagonal matrix A, $x^\top Ay\le \frac{1}{2}x^\top Ax+\frac{1}{2}y^\top Ay$. Substituting (B.3) and (B.4) into (B.2), we get

$$\begin{aligned} \nabla f(\theta _t)^\top{} & {} \!\!\!(\theta _{t+1}-\theta _t) \le -\frac{1-\beta }{1-\beta ^t}\nabla f(\theta _t)^\top \eta _{t-1}r_{t} g_{t} + \frac{\eta G^2_\infty }{\sqrt{a}}(\Vert r_{t}\Vert _{1,1}-\Vert r_{t+1}\Vert _{1,1})\nonumber \\{} & {} +\bigg (\frac{2\eta }{\sqrt{a}}+\frac{\beta \eta }{2(1-\beta )\sqrt{a}}\bigg )g_t^\top r_{t+1}g_t + \frac{\beta \eta }{2(1-\beta )\sqrt{a}}m_{t-1}^\top r_{t+1}m_{t-1}\nonumber \\{} & {} +\frac{2\eta }{\sqrt{a}}|(\nabla f(\theta _t)-g_t)^\top r_{t+1}g_{t}|+\frac{\beta \eta }{(1-\beta )\sqrt{a}}|(\nabla f(\theta _t)-g_t)^\top r_{t+1}m_{t-1}|. \end{aligned}$$

(B.5)

With (B.5), we take an conditional expectation on (B.1) with respect to $(\theta )$ and rearrange to get

$$\begin{aligned}{} & {} \quad \frac{1-\beta }{1-\beta ^t}\nabla f(\theta _t)^\top \eta _{t-1}r_{t}\nabla f(\theta _t) =\mathbb {E}_{\xi _t}\bigg [\frac{1-\beta }{1-\beta ^t}\nabla f(\theta _t)^\top \eta _{t-1}r_{t}g_t\bigg ]\nonumber \\{} & {} \le \mathbb {E}_{\xi _t}\Bigg [f(\theta _t)-f(\theta _{t+1})+ \frac{\eta G^2_\infty }{\sqrt{a}}(\Vert r_{t}\Vert _{1,1}-\Vert r_{t+1}\Vert _{1,1})\nonumber \\{} & {} \quad \quad +\bigg (\frac{2\eta }{\sqrt{a}}+\frac{\beta \eta }{2(1-\beta )\sqrt{a}}\bigg )g_t^\top r_{t+1}g_t + \frac{\beta \eta }{2(1-\beta )\sqrt{a}}m_{t-1}^\top r_{t+1}m_{t-1}\nonumber \\{} & {} \quad \quad +\frac{2\eta }{\sqrt{a}}|(\nabla f(\theta _t)-g_t)^\top r_{t+1}g_{t}|\nonumber \\{} & {} \quad \quad +\frac{\beta \eta }{(1-\beta )\sqrt{a}}|(\nabla f(\theta _t)-g_t)^\top r_{t+1}m_{t-1}|+\frac{L}{2}\Vert \theta _{t+1}-\theta _t\Vert ^2\Bigg ], \end{aligned}$$

(B.6)

where the assumption $\mathbb {E}_{\xi _t}[g_t]=\nabla f(\theta _t)$ is used in the first equality. Since $\xi _1,...,\xi _t$ are independent random variables, we set $\mathbb {E}=\mathbb {E}_{\xi _1}\mathbb {E}_{\xi _2}...\mathbb {E}_{\xi _T}$ and take a summation on (B.6) over t from 1 to T to get

$$\begin{aligned}{} & {} \quad \mathbb {E}\Bigg [\sum \limits _{t=1}^{T}\frac{1-\beta }{1-\beta ^t}\nabla f(\theta _t)^\top \eta _{t-1}r_{t}\nabla f(\theta _t)\Bigg ]\nonumber \\{} & {} \le \mathbb {E}\Big [f(\theta _1)-f(\theta _{T+1})\Big ] + \frac{\eta G^2_\infty }{\sqrt{a}}\mathbb {E}\Big [\Vert r_{1}\Vert _{1,1}-\Vert r_{T+1}\Vert _{1,1}\Big ]\nonumber \\{} & {} \quad + \bigg (\frac{2\eta }{\sqrt{a}}+\frac{\beta \eta }{2(1-\beta )\sqrt{a}}\bigg )\mathbb {E}\Bigg [\sum \limits _{t=1}^{T}g_t^\top r_{t+1}g_t\Bigg ] + \frac{\beta \eta }{2(1-\beta )\sqrt{a}}\mathbb {E}\Bigg [\sum \limits _{t=1}^{T}m_{t-1}^\top r_{t}m_{t-1}\Bigg ]\nonumber \\{} & {} \quad +\frac{2\eta }{\sqrt{a}}\mathbb {E}\Bigg [\sum \limits _{t=1}^{T}|(\nabla f(\theta _t)-g_t)^\top r_{t+1}g_{t}|\Bigg ]\nonumber \\{} & {} \quad +\frac{\beta \eta }{(1-\beta )\sqrt{a}}\mathbb {E}\Bigg [\sum \limits _{t=1}^{T}|(\nabla f(\theta _t)-g_t)^\top r_{t+1}m_{t-1}|\Bigg ]+\frac{L}{2}\mathbb {E}\Bigg [\sum \limits _{t=1}^{T}\Vert \theta _{t+1}-\theta _{t}\Vert ^2\Bigg ]. \end{aligned}$$

(B.7)

Below, we bound each term in (B.7) separately. By the Cauchy-Schwarz inequality, we get

$$\begin{aligned}{} & {} \quad \mathbb {E}\Bigg [\sum \limits _{t=1}^{T}|(\nabla f(\theta _t)-g_t)^\top r_{t+1}m_{t-1}|\Bigg ]\nonumber \\{} & {} \le \mathbb {E}\Bigg [\sum \limits _{t=1}^{T}\Vert \nabla f(\theta _t)-g_t\Vert \Vert r_{t+1}m_{t-1}\Vert \Bigg ]\nonumber \\{} & {} \le \mathbb {E}\Bigg [\bigg (\sum \limits _{t=1}^{T}\Vert \nabla f(\theta _t)-g_t\Vert ^2\bigg )^{1/2}\bigg (\sum \limits _{t=1}^{T}\Vert r_{t+1}m_{t-1}\Vert ^2\bigg )^{1/2}\Bigg ]\nonumber \\ \end{aligned}$$

$$\begin{aligned}{} & {} \le \Bigg (\mathbb {E}\Bigg [\sum \limits _{t=1}^{T}\Vert \nabla f(\theta _t)-g_t\Vert ^2\Bigg ]\Bigg )^{1/2}\Bigg (\mathbb {E}\Bigg [\sum \limits _{t=1}^{T}\Vert r_{t+1}m_{t-1}\Vert ^2\Bigg ]\Bigg )^{1/2}\nonumber \\{} & {} \le \sqrt{2BnT/\eta }F(\theta _1)\sigma _g, \end{aligned}$$

(B.8)

where Lemma 1 (ii) and the bounded variance assumption were used. We replace $m_{t-1}$ in () by $g_t$ and use Lemma 1 (v) to get

$$\begin{aligned}{} & {} \quad \mathbb {E}\Bigg [\sum \limits _{t=1}^{T}|(\nabla f(\theta _t)-g_t)^\top r_{t+1}g_{t}|\Bigg ] \nonumber \\{} & {} \le \Bigg (\mathbb {E}\Bigg [\sum \limits _{t=1}^{T}\Vert \nabla f(\theta _t)-g_t\Vert ^2\Bigg ]\Bigg )^{1/2}\Bigg (\mathbb {E}\Bigg [\sum \limits _{t=1}^{T}\Vert r_{t+1}g_{t}\Vert ^2\Bigg ]\Bigg )^{1/2}\nonumber \\{} & {} \le \frac{2\sqrt{2BnT/\eta }F(\theta _1)\sigma _g}{1-\beta }. \end{aligned}$$

(B.9)

By (4.2), the last term in (B.7) is bounded above by

$$\begin{aligned} \frac{L}{2}\mathbb {E}\left[ \sum _{t=0}^\infty \Vert \theta _{t+1}-\theta _t\Vert ^2 \right] \le \frac{L\eta n}{2}F^2(\theta _1). \end{aligned}$$

(B.10)

Substituting Lemma 1 (i) (iii), (B.10), (B.9), () into (B.7) to get

$$\begin{aligned} \mathbb {E}\Bigg [\sum \limits _{t=1}^{T}\frac{1-\beta }{1-\beta ^t}{} & {} \!\!\!\nabla f(\theta _t)^\top \eta _{t-1} r_{t}\nabla f(\theta _t)\Bigg ] \le (f(\theta _1)-f^*)+\frac{\eta G^2_\infty }{\sqrt{a}}nF(\theta _1)\nonumber \\{} & {} +\bigg (\frac{2}{\sqrt{a}}+\frac{\beta }{2(1-\beta )\sqrt{a}}\bigg )\frac{8BnF(\theta _1)}{(1-\beta )^2}+\frac{\beta BnF(\theta _1)}{(1-\beta )\sqrt{a}}\nonumber \\{} & {} +\frac{(4+\beta )\sqrt{2B\eta }}{(1-\beta )\sqrt{a}}F(\theta _1)\sqrt{nT}\sigma _g+\frac{L\eta n}{2}F^2(\theta _1). \end{aligned}$$

(B.11)

Note that the left hand side is bounded from below by

$$\begin{aligned} (1-\beta )\frac{\eta }{\sqrt{B}} \mathbb {E}\Bigg [\min \limits _ir_{T,i}\sum \limits _{t=1}^{T}\Vert \nabla f(\theta _t)\Vert ^2\Bigg ], \end{aligned}$$

where we used $|\frac{1-\beta }{1-\beta ^t}|\ge 1-\beta$ and $\eta _t\ge \eta /\sqrt{B}$. Thus, we have

$$\begin{aligned} \mathbb {E}\Bigg [\min \limits _ir_{T,i}\sum \limits _{t=1}^{T}\Vert \nabla f(\theta _t)\Vert ^2\Bigg ] \le \frac{C_1+C_2n+C_3\sigma _g \sqrt{ nT}}{\eta }, \end{aligned}$$

(B.12)

where

$$\begin{aligned} C_1= & {} \frac{(f(\theta _1)-f^*)\sqrt{B}}{1-\beta },\\ C_2= & {} \frac{\sqrt{B}\eta G^2_\infty F(\theta _1)}{(1-\beta )\sqrt{a}} +\bigg (\frac{2}{\sqrt{a}}+\frac{\beta }{2(1-\beta )\sqrt{a}}\bigg )\frac{8B^{3/2}F(\theta _1)}{(1-\beta )^3}\\{} & {} \quad \!\!\!\!+\frac{\beta B^{3/2}F(\theta _1)}{(1-\beta )^2\sqrt{a}}+\frac{\sqrt{B}L\eta }{2(1-\beta )^2}F^2(\theta _1),\\ C_3= & {} \frac{(4+\beta )B\sqrt{2\eta }}{(1-\beta )\sqrt{a}}F(\theta _1). \end{aligned}$$

By the Hölder inequality, we have for any $\alpha \in (0, 1)$,

$$\begin{aligned} \mathbb {E}[X^\alpha ]\le \mathbb {E}[XY]^{\alpha } \mathbb {E}[Y^{-\alpha /(1-\alpha )}]^{1-\alpha }. \end{aligned}$$

Take $X=\Delta :=\sum _{t=1}^{T}\Vert \nabla f(\theta _t)\Vert ^2, Y=\min _ir_{T,i}$, we obtain

$$\begin{aligned} \mathbb {E}[\Delta ^\alpha ]\le \mathbb {E}[\min _ir_{T,i}\Delta ]^\alpha \mathbb {E}[(\min _ir_{T,i})^{-\alpha /(1-\alpha )}]^{1-\alpha }. \end{aligned}$$

Using (B.12), and lower bounding $\mathbb {E}[\Delta ^\alpha ]$ by $T^{\alpha }\mathbb {E}[\min _{1\le t\le T}\Vert \nabla f(\theta _t)\Vert ^{2\alpha }]$, we obtain

$$\begin{aligned} \mathbb {E}\Bigg [\min _{1\le t \le T}\Vert \nabla f(\theta _t)\Vert ^{2\alpha } \Bigg ] \le \left( \frac{C_1+C_2n+C_3\sigma _g \sqrt{ nT}}{\eta T } \right) ^\alpha \mathbb {E}[(\min _ir_{T,i})^{-\alpha /(1-\alpha )}]^{1-\alpha }. \end{aligned}$$

This by taking $\alpha =1-\epsilon$ yields the stated bound.

Appendix 3. Proof of Theorem 4

Using the same argument as for (iv) in Lemma 2, we have

$$\begin{aligned} \sum \limits _{i=1}^{n}\sum \limits _{t=1}^{T}r_{t+1,i}g^2_{t,i}\le \frac{8Bn\sqrt{f_1(\theta _1)+c}}{(1-\beta )^2\eta }. \end{aligned}$$

With this estimate and the convexity of $f_t$, the regret can be bounded by

$$\begin{aligned} R(T)= & {} \sum \limits _{t=1}^{T}f_t(\theta _t)-f_t(\theta ^*) \le \sum \limits _{t=1}^{T} g_t^\top (\theta _t-\theta ^*)\\\le & {} \sum \limits _{i=1}^{n}\sum \limits _{t=1}^{T}|g_{t,i}|\sqrt{r_{t+1,i}} \frac{|\theta _{t,i}-\theta ^*_i|}{\sqrt{r_{t+1,i}}}\\\le & {} \left( \sum \limits _{i=1}^{n}\sum \limits _{t=1}^{T} r_{t+1, i} g_{t, i}^2 \right) ^{1/2}\left( \sum \limits _{i=1}^{n}\sum \limits _{t=1}^{T} \frac{|\theta _{t,i}-\theta ^*_i|^2}{r_{t+1,i}}\right) ^{1/2}\\\le & {} \frac{2D_\infty \sqrt{2B}}{1-\beta }(f_1(\theta _1)+c)^{1/4}\sqrt{nT/\eta }\left( \sum \limits _{i=1}^{n} \frac{1}{r_{T+1,i}}\right) ^{1/2}, \end{aligned}$$

where the fourth inequality is by the Cauchy-Schwarz inequality, and the assumption $\Vert x-y\Vert _\infty \le D_\infty$ for all $x,y\in \mathcal {F}$ is used in the last inequality.