Section 4.3 Stochastic Compositional Momentum Methods

In this section, we present a method that matches the sample complexity of the large mini-batch approach without using large mini-batches under the smoothness conditions of \(f\) and \(F\). The idea is to design a gradient estimator such that its error can be reduced gradually. It turns out this technique, related to the momentum methods for standard stochastic optimization, is more widely applicable to other problems discussed later in this chapter. Furthermore, we introduce advanced methods to further improve the complexity to \(O(1/\epsilon^3)\) under stronger conditions.

It is worth noting that the results in this section apply to the standard stochastic optimization problem under the smoothness assumption of \(g(\mathbf{w})\) by setting \(f_i(g) =g\) and \(L_1=0\) in the complexity results and removing the \(\mathbf{u}\) update in the algorithm.

4.3.1 Moving-Average Gradient Estimator

The first method is to use the following moving-average gradient estimator: \[\begin{align}\label{eqn:mag} \mathbf{v}_{t} & = (1-\beta_t)\mathbf{v}_{t-1} + \beta_t \nabla g(\mathbf{w}_t; \zeta_t')\nabla f(\mathbf{u}_{t}; \xi_t), \end{align}\] where \(0 \leq \beta_t < 1\). With \(\mathbf{v}_{t}\), the model parameter is updated by: \[\begin{align} \mathbf{w}_{t+1} & = \mathbf{w}_t - \eta_t \mathbf{v}_{t}. \end{align}\]

We present the full steps in Algorithm 10 and refer to it as SCMA.

To understand this method, we can view \(\mathbf{v}_{t}\) as a better estimator of the gradient, with its estimation error gradually decreasing over iterations — a property we will prove shortly. This yields an enhanced stability of momentum-based metholds observed in practice.

First, we prove a generic lemma that establishes the error recursion of \(\mathbf{v}_{t}\).

Proof

Due to the update formula \(\mathbf{v}_{t} = (1-\beta_t)\mathbf{v}_{t-1} + \beta_t \mathbf{z}_t\), we have \[\begin{align*} & \mathbb{E}_{t}\left[\left\|\mathbf{v}_{t} - \nabla F(\mathbf{w}_t)\right\|_2^2\right] \\ & = \mathbb{E}_{t}\left[\left\|(1-\beta_t)\mathbf{v}_{t-1} + \beta_t \mathbf{z}_t - \nabla F(\mathbf{w}_t)\right\|_2^2\right]\\ & = \mathbb{E}_{t}\bigg[\|\underbrace{(1-\beta_t)\mathbf{v}_{t-1}- \nabla F(\mathbf{w}_t) + \beta_t \mathcal{M}_t}\limits_{\mathbf{a}_t}+\underbrace{\beta_t (\mathbf{z}_t - \mathcal{M}_t)}\limits_{\mathbf{b}_t}\|_2^2\bigg]. \end{align*}\]

Note that \(\mathbb{E}_{t}[\mathbf{a}_t^{\top}\mathbf{b}_t] = 0\). Besides, we have \(\mathbb{E}_t[\|\mathbf{b}_t\|_2^2]\leq \beta_t^2 \sigma^2\). Due to Young’s inequality, we have \(\|a+b\|_2^2 \leq (1+\alpha)\|a\|_2^2 + (1+1/\alpha)\|b\|_2^2\) for any \(\alpha>0\). Hence, \[\begin{align*} &\|\mathbf{a}_t\|_2^2=\|(1-\beta_t)(\mathbf{v}_{t-1} - \nabla F(\mathbf{w}_{t-1})) + (1-\beta_t)(\nabla F(\mathbf{w}_{t-1}) - \nabla F(\mathbf{w}_t))\\ &\quad \quad\quad \quad + \beta_t (\mathcal{M}_t - \nabla F(\mathbf{w}_t))\|_2^2\\ &\leq (1-\beta_t)^2(1+\beta_t)\|(\mathbf{v}_{t-1} - \nabla F(\mathbf{w}_{t-1}))\|_2^2\\ & + (1+\frac{1}{\beta_t})\|(1-\beta_t)(\nabla F(\mathbf{w}_{t-1}) - \nabla F(\mathbf{w}_t)) + \beta_t (\mathcal{M}_t - \nabla F(\mathbf{w}_t))\|_2^2\\ & \leq (1-\beta_t)\left\|\mathbf{v}_{t-1} - \nabla F(\mathbf{w}_{t-1})\right\|_2^2 + \frac{2(1+\beta_t)(1-\beta_t)^2}{\beta_t}\left\|\nabla F(\mathbf{w}_{t-1}) - \nabla F(\mathbf{w}_t)\right\|_2^2 \\ & + \frac{2(1+\beta_t)\beta_t^2}{\beta_t}\left\|\mathcal{M}_t - \nabla F(\mathbf{w}_t)\right\|_2^2\\ & \leq (1-\beta_t)\left\|\mathbf{v}_{t-1} - \nabla F(\mathbf{w}_{t-1})\right\|_2^2 + \frac{2L_F^2}{\beta_t}\left\|\mathbf{w}_{t-1}-\mathbf{w}_t\right\|_2^2 + 4\beta_t\left\|\mathcal{M}_{t} - \nabla F(\mathbf{w}_t)\right\|_2^2. \end{align*}\]

■

With the above lemma, we are able to establish the error recursion of \(\mathbf{v}_{t}\) of SCMA.

Proof

We denote by \(\mathbb{E}_{t}[\cdot]=\mathbb{E}_{\xi_t, \zeta'_t}[\cdot]\). Let \(\mathbf{z}_t = \nabla g(\mathbf{w}_t;\zeta'_t)\nabla f(\mathbf{u}_{t};\xi_t)\) and \(\mathcal{M}_t = \mathbb{E}_t[\mathbf{z}_t] = \nabla g(\mathbf{w}_t)\nabla f(\mathbf{u}_{t})\). Lemma 4.4 proves that \[\begin{align} \mathbb{E}_t[\|\mathbf{z}_t -\mathcal{M}_t\|_2^2]&\leq G_2^2\sigma_1^2 + G_1^2\sigma_2^2,\label{eqn:sgf} \end{align}\] and \[\begin{align*} \|\mathcal{M}_t - \nabla F(\mathbf{w}_t)\|_2^2& = \|\nabla g(\mathbf{w}_t)\nabla f(\mathbf{u}_{t}) -\nabla g(\mathbf{w}_t)\nabla f(g(\mathbf{w}_{t})) \|_2^2\\ & \leq G_2^2L_1^2\|\mathbf{u}_t - g(\mathbf{w}_t)\|_2^2. \end{align*}\]

■

To enjoy the above recursion of the gradient estimator’s error, we state the following lemma, which is a variant of the standard descent lemma of gradient descent.

Proof

Due to the smoothness of \(F\), we have \[\begin{align*} & F(\mathbf{w}_{t+1}) \leq F(\mathbf{w}_t) + \mathbf{\nabla} F(\mathbf{w}_t)^{\top}(\mathbf{w}_{t+1} - \mathbf{w}_t) + \frac{L_F}{2}\left\|\mathbf{w}_{t+1}-\mathbf{w}_t\right\|_2^2\\ & = F(\mathbf{w}_t) + (\nabla F(\mathbf{w}_t) - \mathbf{v}_{t})^{\top}(\mathbf{w}_{t+1} - \mathbf{w}_t) + \mathbf{v}_{t}^{\top}(\mathbf{w}_{t+1} - \mathbf{w}_t)+ \frac{L_F}{2}\left\|\mathbf{w}_{t+1}-\mathbf{w}_t\right\|_2^2\\ & = F(\mathbf{w}_t) - \eta_t (\nabla F(\mathbf{w}_t) - \mathbf{v}_{t})^{\top}\mathbf{v}_{t} - \left(\frac{1}{\eta_t} - \frac{L_F}{2}\right)\left\|\mathbf{w}_{t+1} - \mathbf{w}_t\right\|_2^2\\ & = F(\mathbf{w}_t) + \eta_t \left\|(\nabla F(\mathbf{w}_t) - \mathbf{v}_{t})\right\|_2^2 - \eta_t (\nabla F(\mathbf{w}_t) - \mathbf{v}_{t})^{\top}\nabla F(\mathbf{w}_t) \\ &- \left(\frac{1}{\eta_t} - \frac{L_F}{2}\right)\left\|\mathbf{w}_{t+1} - \mathbf{w}_t\right\|_2^2. \end{align*}\]

■

To prove the final convergence of SCMA, we present a useful lemma.

Proof

The proof is constructive. The idea is to construct a telescoping series of \(A_{t}+a_t\Delta_{t}+b_t \delta_{t}\) with some appropriate sequences of \(a_t, b_t\). First, we have \[\begin{align*} &A_{t+1} + a_{t+1}\Delta_{t+1} + b_{t+1} \delta_{t+1}\leq A_t + \eta_t \Delta_t - \eta_t B_t - \eta_t \Gamma_t\\ & +a_{t+1}(1-\beta_{t+1})\Delta_{t} + a_{t+1}C_1\beta_{t+1}\delta_{t+1} + a_{t+1}\frac{C_2\eta_{t}^2}{\beta_{t+1}}\Gamma_t + a_{t+1}\beta_{t+1}^2\sigma^2\\ &+b_{t+1}(1-\gamma_{t+1}) \delta_{t} + b_{t+1}\frac{C_3\eta_{t}^2}{\gamma_{t+1}} \Gamma_{t} + b_{t+1}\gamma_{t+1}^2{\sigma'}^2. \end{align*}\] Let \(a_{t+1} = \eta_t/\beta_{t+1}\leq \eta_{t-1}/\beta_{t}\) and \(b_{t+1} = C_1\eta_t(1+\gamma_{t+1})/\gamma_{t+1}\), we have \[\begin{align*} &A_{t+1} + \frac{\eta_t}{\beta_{t+1}}\Delta_{t+1} +(C_1\eta_t\frac{1+\gamma_{t+1}}{\gamma_{t+1}}- C_1\eta_t) \delta_{t+1}\leq A_t - \eta_t B_t - \eta_t \Gamma_t \\ & +\left(\eta_t+\frac{\eta_t}{\beta_{t+1}}(1-\beta_{t+1})\right)\Delta_{t} + \frac{C_2\eta_{t}^3}{\beta_{t+1}^2}\Gamma_t + \eta_t\beta_{t+1}\sigma^2\\ &+C_1\eta_t\frac{1+\gamma_{t+1}}{\gamma_{t+1}}(1-\gamma_{t+1})\delta_{t} + \frac{C_3C_1\eta_{t}^3(1+\gamma_{t+1})}{\gamma_{t+1}^2} \Gamma_t +C_1\eta_t(1+\gamma_{t+1})\gamma_{t+1}{\sigma'}^2. \end{align*}\] Thus, \[\begin{align*} &A_{t+1} + \frac{\eta_t}{\beta_{t+1}}\Delta_{t+1} + \frac{C_1\eta_t}{\gamma_{t+1}} \delta_{t+1}\leq A_t +\frac{\eta_t}{\beta_{t+1}}\Delta_{t}+\frac{C_1\eta_t}{\gamma_{t+1}}\delta_{t}\\ &- \eta_t B_t - \left(\eta_t - \frac{C_2\eta_{t}^3}{\beta_{t+1}^2} - \frac{C_3C_1\eta_{t}^3(1+\gamma_{t+1})}{\gamma_{t+1}^2}\right) \Gamma_t\\ &+ \eta_t\beta_{t+1}\sigma^2 + C_1\eta_t(1+\gamma_{t+1})\gamma_{t+1}{\sigma'}^2. \end{align*}\] Since \(\eta_t/\beta_{t+1}\leq \eta_{t-1}/\beta_{t}\) and \(\eta_t/\gamma_{t+1}\leq \eta_{t-1}/\gamma_{t}\) and \(\gamma_{t+1}\leq 1\), we have \[\begin{align*} &A_{t+1} + \frac{\eta_t}{\beta_{t+1}}\Delta_{t+1} + \frac{C_1\eta_t}{\gamma_{t+1}} \delta_{t+1}\leq A_t +\frac{\eta_{t-1}}{\beta_{t}}\Delta_{t}+\frac{C_1\eta_{t-1}}{\gamma_{t}}\delta_{t}\\ &- \eta_t B_t - \left(\eta_t - \frac{C_2\eta_{t}^3}{\beta_{t+1}^2} - \frac{2C_3C_1\eta_{t}^3}{\gamma_{t+1}^2}\right) \Gamma_t\\ &+ \eta_t\beta_{t+1}\sigma^2 + 2C_1\eta_t\gamma_{t+1}{\sigma'}^2. \end{align*}\] Since \(C_2\eta_{t}^3/\beta_{t+1}^2\leq \eta_t/4\) because \(\eta_t\leq \beta_{t+1}/\sqrt{4C_2}\) and \(2C_3C_1\eta_{t}^3/\gamma_{t+1}^2\leq \eta_t/4\) because \(\eta_t\leq \gamma_{t+1}/\sqrt{8C_1C_3}\), we have \[\begin{align*} &A_{t+1} + \frac{\eta_t}{\beta_{t+1}}\Delta_{t+1} + \frac{C_1\eta_t}{\gamma_{t+1}} \delta_{t+1}\leq A_t +\frac{\eta_{t-1}}{\beta_{t}}\Delta_{t}+\frac{C_1\eta_{t-1}}{\gamma_{t}}\delta_{t}\\ &- \eta_t B_t - \frac{1}{2}\eta_t \Gamma_t+ \eta_t\beta_{t+1}\sigma^2 + 2C_1\eta_t\gamma_{t+1}{\sigma'}^2. \end{align*}\] Define \(\Upsilon_{t+1}=A_{t+1} + \frac{\eta_t}{\beta_{t+1}}\Delta_{t+1} + \frac{C_1\eta_t}{\gamma_{t+1}} \delta_{t+1}\), we have \[\begin{align*} &\eta_t B_t + \frac{1}{2}\eta_t \Gamma_t\leq \Upsilon_{t} - \Upsilon_{t+1}+ \eta_t\beta_{t+1}\sigma^2 + 2C_1\eta_t\gamma_{t+1}{\sigma'}^2. \end{align*}\] Hence \[\begin{align*} &\sum_{t=0}^{T-1}(\eta_t B_t + \frac{1}{2}\eta_t \Gamma_t)\leq \Upsilon_{0} - A_* + \sum_{t=0}^{T-1}\left(\eta_t\beta_{t+1}\sigma^2 + 2C_1\eta_t\gamma_{t+1}{\sigma'}^2\right). \end{align*}\]

■

Finally, let us prove the convergence of SCMA.

Proof

The three inequalities in Lemma 4.8, Lemma 4.9 and Lemma 4.1 that we have proved so far are \[\begin{align*} (*)\;&F(\mathbf{w}_{t+1}) \leq F(\mathbf{w}_t) + \frac{\eta_t}{2}\left\|\nabla F(\mathbf{w}_t) - \mathbf{v}_{t}\right\|_2^2- \frac{\eta_t}{2}\left\|\nabla F(\mathbf{w}_t)\right\|_2^2- \frac{\eta_t}{4} \left\|\mathbf{v}_{t}\right\|_2^2, t\geq 0\\ (\sharp)\;&\mathbb E\left[\left\|\mathbf{v}_{t} - \nabla F(\mathbf{w}_t)\right\|_2^2\right] \leq \mathbb E[(1-\beta_t)\left\|\mathbf{v}_{t-1} - \nabla F(\mathbf{w}_{t-1})\right\|_2^2]\\ &\quad+ \mathbb E\left[4G_2^2L_1^2\beta_t\left\|\mathbf{u}_{t} - g(\mathbf{w}_t)\right\|_2^2 + \frac{2L_F^2\eta_{t-1}^2}{\beta_t}\left\|\mathbf{v}_{t-1}\right\|_2^2 + \beta_t^2\sigma^2\right],\\ (\diamond)\;&\mathbb E\left[\left\|\mathbf{u}_{t} - g(\mathbf{w}_t)\right\|_2^2\right] \leq \mathbb E[(1-\gamma_t) \left\|\mathbf{u}_{t-1} - g(\mathbf{w}_{t-1})\right\|_2^2]\\ & \quad + \mathbb E\left[\frac{G_2^2\eta_{t-1}^2}{\gamma_t} \left\|\mathbf{v}_{t-1}\right\|_2^2 + \gamma_t^2\sigma_0^2\right]. \end{align*}\]

Define \(A_t= 2(F(\mathbf{w}_t) - F_*)\) and \(B_t = \|\nabla F(\mathbf{w}_t)\|_2^2\), \(\Gamma_t = \left\|\mathbf{v}_{t}\right\|_2^2/2\), \(\Delta_t= \left\|\nabla F(\mathbf{w}_t) - \mathbf{v}_{t}\right\|_2^2\), \(\delta_{t}=\left\|\mathbf{u}_{t} - g(\mathbf{w}_t)\right\|_2^2\), and \(\Upsilon_{t}=A_{t} + \frac{\eta_{t-1}}{\beta_{t}}\Delta_{t} + \frac{C_1\eta_{t-1}}{\gamma_{t}} \delta_{t}.\)

Then the three inequalities satisfy those in Lemma 4.10 with \(C_1=4G_2^2L_1^2, C_2=4L_F^2, C_3=2G_2^2, \sigma^2 = G_1^2\sigma_2^2+G_2^2\sigma_1^2, {\sigma'}^2=\sigma_0^2\). Then \(\eta_t, \beta_t, \gamma_t\) satisfy \[ \frac{\eta_t}{\beta_{t+1}}\leq \frac{\eta_{t-1}}{\beta_{t}}, \frac{\eta_t}{\gamma_{t+1}}\leq \frac{\eta_{t-1}}{\gamma_{t}}, \eta_t\leq \min(\frac{\beta_{t+1}}{\sqrt{4C_2}}, \frac{\gamma_{t+1}}{\sqrt{8C_1C_3}}). \] Then we have \[\begin{align*} &\mathbb E\left[\sum_{t=0}^{T-1}\frac{1}{\sum_{t=0}^{T-1}\eta_t}(\eta_t \|\nabla F(\mathbf{w}_t)\|_2^2 + \frac{\eta_t}{4}\left\|\mathbf{v}_{t}\right\|_2^2)\right]\\ & \leq \frac{C_\Upsilon}{\sum_{t=1}\eta_t} + \frac{\sum_{t=0}^{T-1}\left(\eta_t\beta_{t+1}\sigma^2 + 2C_1\eta_t\gamma_{t+1}\sigma_0^2\right)}{\sum_{t=0}^{T-1}\eta_t}. \end{align*}\]

■

4.3.2 STORM Estimators

We can further reduce the error of the gradient estimator by using advanced variance reduction techniques under stronger assumptions. We make the following assumptions.

Due to Jensen’s inequality, Assumption 4.6(i) implies the Lipschitz continuity assumption of \(\nabla f\) in Assumption 4.1(i). Similarly, Assumption 4.7(i) implies that in Assumption 4.2(i), respectively. Hence, Assumption 4.6(i) and Assumption 4.7(i) are stronger, which are referred to as mean-square smoothness condition of \(f\) and \(g\).

Applying the STORM estimator to estimating the sequence of \(\{g(\mathbf{w}_t)\}_{t\geq 1}\), we have the following sequence:

\[\begin{align} \mathbf{u}_{t} = (1-\gamma_t)\mathbf{u}_{t-1}+\gamma_t g(\mathbf{w}_t;\zeta_t)+(1-\gamma_t)(g(\mathbf{w}_t;\zeta_t)-g(\mathbf{w}_{t-1};\zeta_t)). \end{align}\]

Given \(\mathbf{u}_{t}\), we can compute a moving-average gradient estimator (1) similar to SCMA. However, this will not yield an improved rate compared with SCMA. To reduce the estimator error of the gradient, we apply another STORM estimator to estimate \(\mathcal{M}_{t}=\nabla g(\mathbf{w}_t)\nabla f(\mathbf{u}_{t})\). This is computed by the following sequence:

\[\begin{align}\label{eqn:cstorm_grad} \mathbf{v}_{t} & = (1-\beta_t)\mathbf{v}_{t-1}+\beta_t\nabla g(\mathbf{w}_t;\zeta_t')\nabla f(\mathbf{u}_{t};\xi_t)\\ &\quad +(1-\beta_t)(\nabla g(\mathbf{w}_t;\zeta_t')\nabla f(\mathbf{u}_{t};\xi_t)-\nabla g(\mathbf{w}_{t-1};\zeta_t')\nabla f(\mathbf{u}_{t-1};\xi_t)).\notag \end{align}\]

With \(\mathbf{v}_{t}\), we update the model parameters by

\[\begin{align*} \mathbf{w}_{t+1} = \mathbf{w}_t-\eta\mathbf{v}_{t}. \end{align*}\]

The full steps of this method is presented in Algorithm 11, which is referred to as SCST.

Convergence Analysis

We first prove a general result of the STORM estimator that applies to both \(\mathbf{u}_t\) and \(\mathbf{v}_t\).

Proof

\[\begin{align*} &\mathbb E_{t}\left[\|\mathbf{v}_{t}-\mathcal{M}_t\|^2\right]=\mathbb E_{t}\left[\|(1-\beta_t)\mathbf{v}_{t-1}-\mathcal{M}_t+\beta_t\mathbf{z}_t+(1-\beta_t)(\mathbf{z}_t-\tilde{\mathbf{z}}_{t-1})\|^2\right]\\ &=\mathbb E_{t}\left[\|(1-\beta_t)(\mathbf{v}_{t-1}-\mathcal{M}_{t-1})+(1-\beta_t)((\mathbf{z}_t-\tilde{\mathbf{z}}_{t-1})-(\mathcal{M}_t-\mathcal{M}_{t-1}))+\beta_t(\mathbf{z}_t-\mathcal{M}_t)\|_2^2\right]. \end{align*}\]

Note that

\[\begin{align*} \mathbb E_{t}&\left[\langle(1-\beta_t)(\mathbf{v}_{t-1}-\mathcal{M}_{t-1}), (1-\beta_t)((\mathbf{z}_t-\tilde{\mathbf{z}}_{t-1})-(\mathcal{M}_t-\mathcal{M}_{t-1}))+\beta_t(\mathbf{z}_t-\mathcal{M}_t)\rangle\right]=0. \end{align*}\]

Then,

\[\begin{align*} &\mathbb E_t\left[\|\mathbf{v}_{t}-\mathcal{M}_t\|^2\right]\leq (1-\beta_t)^2\|\mathbf{v}_{t}-\mathcal{M}_{t-1}\|_2^2\\ &\quad +\|(1-\beta_t)((\mathbf{z}_t-\tilde{\mathbf{z}}_{t-1})-(\mathcal{M}_t-\mathcal{M}_{t-1}))+\beta_t(\mathbf{z}_t-\mathcal{M}_t)\|_2^2\\ &\stackrel{(\diamond)}{\leq} (1-\beta_t)^2\|\mathbf{v}_{t}-\mathcal{M}_{t-1}\|^2\\ &\quad +2(1-\beta_t)^2\mathbb E_t[\|((\mathbf{z}_t-\tilde{\mathbf{z}}_{t-1})-(\mathcal{M}_t-\mathcal{M}_{t-1}))\|_2^2]+2\beta_t^2\mathbb E_t[\|\mathbf{z}_t-\mathcal{M}_t\|_2^2]\\ &\stackrel{(*)}{\leq} (1-\beta_t)^2\|\mathbf{v}_{t-1}-\mathcal{M}_{t-1}\|_2^2+2(1-\beta_t)^2\mathbb E_t[\|\mathbf{z}_t-\tilde{\mathbf{z}}_{t-1}\|_2^2]+2\beta_t^2\sigma^2, \end{align*}\]

where \((\diamond)\) uses the Young’s inequality, \((*)\) uses the fact that \(\mathbb E[\|a-\mathbb E[a]\|_2^2]\leq \mathbb E[\|a\|_2^2]\), and \(\mathbb E_t[\mathbf{z}_t-\tilde{\mathbf{z}}_{t-1}]=\mathcal{M}_t-\mathcal{M}_{t-1}\).

■

Let us first prove an error recursion of \(\mathbf{u}_{t}\) in the lemma below.

Proof

The first part follows directly from Lemma 4.11 by noting the mean-Lipschitz continuity of \(g(\mathbf{w};\zeta)\).

To prove the second part, we proceed as follows:

\[\begin{equation*} \begin{split} \mathbb E_{t}&\left[\|\mathbf{u}_{t}-\mathbf{u}_{t-1}\|_2^2 \right]\\ =&\mathbb E_{t}\left[\left\|\gamma_{t}\left(g(\mathbf{w}_{t};\zeta_{t})-\mathbf{u}_{t-1}\right)+(1-\gamma_{t})\left(g(\mathbf{w}_{t};\zeta_{t})-g(\mathbf{w}_{t-1};\zeta_{t})\right)\right\|_2^{2}\right]\\ \leq &\mathbb E_{t}\left[2\gamma_{t}^2\left\|g(\mathbf{w}_{t};\zeta_{t})-\mathbf{u}_{t-1}\right\|_2^2+2(1-\gamma_{t})^2\left\|g(\mathbf{w}_{t};\zeta_{t})-g(\mathbf{w}_{t-1};\zeta_{t})\right\|_2^{2}\right]\\ \leq &\mathbb E_{t}\left[2\gamma_{t}^2\left\|g(\mathbf{w}_{t};\zeta_{t})-\mathbf{u}_{t-1}\right\|_2^2\right]+2(1-\gamma_{t})^2G_2^2\|\mathbf{w}_{t}-\mathbf{w}_{t-1}\|_2^{2}. \end{split} \end{equation*}\]

Next, we bound the first term on the RHS as

\[\begin{align*} &\mathbb E_{t}\left[2\gamma_{t}^2\left\|g(\mathbf{w}_{t};\zeta_{t})-\mathbf{u}_{t-1}\right\|_2^2\right]\\ &=\mathbb E_{t}\left[2\gamma_{t}^2\left\|g(\mathbf{w}_{t};\zeta_{t})-g(\mathbf{w}_t)+g(\mathbf{w}_t)-\mathbf{u}_{t-1}\right\|_2^2\right]\\ &\leq 2\gamma_{t}^2\sigma_0^2+2\gamma_{t}^2\|g(\mathbf{w}_t)-\mathbf{u}_{t-1}\|_2^2\\ &\leq 2\gamma_{t}^2\sigma_0^2+2\gamma_{t}^2\|g(\mathbf{w}_t)-g(\mathbf{w}_{t-1})+g(\mathbf{w}_{t-1})-\mathbf{u}_{t-1}\|_2^2\\ &\leq 2\gamma_{t}^2\sigma_0^2+4\gamma_{t}^2\|g(\mathbf{w}_{t-1})-\mathbf{u}_{t-1}\|_2^2+4\gamma_{t}^2G_2^2\|\mathbf{w}_t-\mathbf{w}_{t-1}\|_2^2, \end{align*}\]

where the first inequality uses the fact \(\mathbb E[g(\mathbf{w}_{t};\zeta_{t})-g(\mathbf{w}_{t})]=0\). Combining the above results, we finish the proof.

■

Next, we build an error recursion of \(\|\mathbf{v}_{t}-\mathcal{M}_t\|_2^2\).

Proof

First, (6) gives \(\mathbb E_t[\|\mathbf{z}_t-\mathcal{M}_t\|_2^2]\leq \sigma^2\).

Second,

\[\begin{equation*} \begin{aligned} &\mathbb E_t\left[\|\mathbf{z}_t-\tilde{\mathbf{z}}_{t-1}\|_2^2\right]\\ &=\mathbb E_t[\|\nabla g(\mathbf{w}_t;\zeta_t')\nabla f(\mathbf{u}_{t};\xi_t)-\nabla g(\mathbf{w}_{t-1};\zeta_t')\nabla f(\mathbf{u}_{t-1};\xi_t)\|_2^2]\\ &=\mathbb E_t[\|\nabla g(\mathbf{w}_t;\zeta_t')\nabla f(\mathbf{u}_{t};\xi_t)-\nabla g(\mathbf{w}_t;\zeta_t')\nabla f(\mathbf{u}_{t-1};\xi_t)\\ &\quad\quad\quad +\nabla g(\mathbf{w}_t;\zeta_t')\nabla f(\mathbf{u}_{t-1};\xi_t)-\nabla g(\mathbf{w}_{t-1};\zeta_t')\nabla f(\mathbf{u}_{t-1};\xi_t)\|_2^2]\\ &\stackrel{(\triangle)}{\leq} 2G_2^2L_1^2\|\mathbf{u}_{t}-\mathbf{u}_{t-1}\|_2^2+2G_1^2L_2^2\|\mathbf{w}_t-\mathbf{w}_{t-1}\|_2^2, \end{aligned} \end{equation*}\]

where \((\triangle)\) uses Assumption 4.6(i) and Assumption 4.7(i).

It then follows:

\[\begin{align*} \mathbb E_t\left[\|\mathbf{v}_{t}-\mathcal{M}_t\|^2\right] &\leq (1-\beta_t)^2\|\mathbf{v}_{t-1}-\mathcal{M}_{t-1}\|_2^2\\ &\quad +4G_2^2L_1^2\|\mathbf{u}_{t}-\mathbf{u}_{t-1}\|_2^2+4G_1^2L_2^2\|\mathbf{w}_t-\mathbf{w}_{t-1}\|_2^2+2\beta_t^2\sigma^2. \end{align*}\]

By using the second inequality of Lemma 4.12, i.e.,

\[\begin{equation*} \mathbb E_{\zeta_t}\left[\|\mathbf{u}_{t}-\mathbf{u}_{t-1}\|_2^2\right]\leq 2\gamma_{t}^2\sigma_0^2+4\gamma_{t}^2\|\mathbf{u}_{t-1}-g(\mathbf{w}_{t-1})\|_2^2+6G_2^2\|\mathbf{w}_t-\mathbf{w}_{t-1}\|_2^2, \end{equation*}\]

we have

\[\begin{align*} &\mathbb E_t\left[\|\mathbf{v}_{t}-\mathcal{M}_t\|^2\right]\leq (1-\beta_t)\|\mathbf{v}_{t-1}-\mathcal{M}_{t-1}\|_2^2+16G_2^2L_1^2\gamma_{t}^2\|\mathbf{u}_{t-1}-g(\mathbf{w}_{t-1})\|_2^2\\ &\quad +(24G_2^4L_1^2+4G_1^2L_2^2)\|\mathbf{w}_t-\mathbf{w}_{t-1}\|_2^2+2\beta_t^2\sigma^2+8G_2^2L_1^2\gamma_{t}^2\sigma_0^2. \end{align*}\]

■

Similar to Lemma 4.9, we have the following descent lemma.

This lemma can be proved following that of Lemma 4.9 by bounding \(\|\mathbf{v}_{t}-\nabla F(\mathbf{w}_t)\|_2^2\leq 2\|\mathbf{v}_{t}-\mathcal{M}_t\|_2^2+2\|\mathcal{M}_t-\nabla F(\mathbf{w}_t)\|_2^2\leq 2\|\mathbf{v}_{t}-\mathcal{M}_t\|_2^2+2G_2^2L_1^2\|\mathbf{u}_{t}-g(\mathbf{w}_t)\|_2^2\).

Proof

The proof is constructive. The idea is to multiply the second inequality by \(a_{t+1}\) and the third inequality by \(b_{t+1}\) such that we can construct a telescoping series of \(A_{t}+a_t\Delta_{t}+b_t\delta_{t}\). First, we have

\[\begin{align*} &A_{t+1}+a_{t+1}\Delta_{t+1}+b_{t+1}\delta_{t+1}\leq A_t+\eta_t\Delta_t+\eta_t\delta_t-\eta_t B_t-\eta_t\Gamma_t\\ &\quad +a_{t+1}(1-\beta_{t+1})\Delta_{t}+a_{t+1}C_1\gamma^2_{t+1}\delta_{t}+a_{t+1}C_2\eta_{t}^2\Gamma_t+a_{t+1}\beta_{t+1}^2\sigma^2+a_{t+1}\gamma_{t+1}^2{\sigma'}^2\\ &\quad +b_{t+1}(1-\gamma_{t+1})\delta_{t}+b_{t+1}C_3\eta_{t}^2\Gamma_{t}+b_{t+1}\gamma_{t+1}^2{\sigma''}^2. \end{align*}\]

Let \(a_{t+1}=c/\eta_t\) and \(b_{t+1}=c'/\eta_t\), we have

\[\begin{align*} &A_{t+1}+\frac{c}{\eta_{t}}\Delta_{t+1}+\frac{c'}{\eta_t}\delta_{t+1}\leq A_t-\eta_t B_t-\eta_t \Gamma_t\\ &\quad +\left(\eta_t+\frac{c}{\eta_t}(1-\beta_{t+1})\right)\Delta_{t}+C_2c\eta_t\Gamma_t+\frac{c\beta^2_{t+1}}{\eta_t}\sigma^2+\frac{c\gamma^2_{t+1}}{\eta_t}\sigma'^2\\ &\quad +\left(\eta_t+\frac{c}{\eta_t}C_1\gamma_{t+1}^2+\frac{c'}{\eta_t}(1-\gamma_{t+1})\right)\delta_{t}+C_3c'\eta_{t}\Gamma_t+\frac{c'\gamma^2_{t+1}}{\eta_t}\sigma''^2. \end{align*}\]

With (\(\ref{eqn:lrr}\)) we have

\[\begin{align*} &A_{t+1}+\frac{c}{\eta_{t}}\Delta_{t+1}+\frac{c'}{\eta_t}\delta_{t+1}\leq A_t+\frac{c}{\eta_{t-1}}\Delta_{t}+\frac{c'}{\eta_{t-1}}\delta_{t}-\eta_t B_t-\frac{1}{2}\eta_t\Gamma_t\\ &\quad +\frac{c\beta^2_{t+1}}{\eta_t}\sigma^2+\frac{c\gamma^2_{t+1}}{\eta_t}\sigma'^2+\frac{c'\gamma^2_{t+1}}{\eta_t}\sigma''^2. \end{align*}\]

Define \(\Upsilon_{t+1}=A_{t+1}+\frac{c}{\eta_t}\Delta_{t+1}+\frac{c'}{\eta_{t}}\delta_{t+1}\), we have

\[\begin{align*} \eta_t B_t+\frac{1}{2}\eta_t\Gamma_t\leq \Upsilon_{t}-\Upsilon_{t+1}+\frac{c\beta^2_{t+1}}{\eta_t}\sigma^2+\frac{c\gamma^2_{t+1}}{\eta_t}\sigma'^2+\frac{c'\gamma^2_{t+1}}{\eta_t}\sigma''^2. \end{align*}\]

Hence

\[\begin{align*} &\sum_{t=0}^{T-1}(\eta_t B_t+\frac{1}{2}\eta_t\Gamma_t)\leq \Upsilon_{0}-A_*+\sum_{t=0}^{T-1}\left(\frac{c\beta^2_{t+1}}{\eta_t}\sigma^2+\frac{c\gamma^2_{t+1}}{\eta_t}\sigma'^2+\frac{c'\gamma^2_{t+1}}{\eta_t}\sigma''^2\right). \end{align*}\]

Next, let us consider \(\eta_t=\eta,\beta_t=\beta,\gamma_t=\gamma\). Then we have

\[\begin{align*} &\sum_{t=0}^{T-1}\frac{1}{T}(B_t+\frac{1}{2}\Gamma_t)\leq \frac{\Upsilon_{0}-A_*}{\eta T}+\left(\frac{c\beta^2}{\eta^2}\sigma^2+\frac{c\gamma^2}{\eta^2}\sigma'^2+\frac{c'\gamma^2}{\eta^2}\sigma''^2\right). \end{align*}\]

In order to ensure the RHS is less than \(\epsilon^2\), it suffices to have

\[\begin{align*} &\beta=\frac{\epsilon\eta}{2\sqrt{c}\sigma},\quad \gamma=\min\left(\frac{\epsilon\eta}{2\sqrt{c}\sigma'},\frac{\epsilon\eta}{2\sqrt{c'}\sigma''}\right),\quad T=\frac{C_\Upsilon}{4\epsilon^2\eta}. \end{align*}\]

To ensure (\(\ref{eqn:lrr}\)), it suffices to have

\[\begin{align*} \eta^2\leq c\beta,\quad C_1c\gamma\leq c'/2,\quad \eta^2\leq c'\gamma/2,\quad c=\frac{1}{4C_2},\quad c'=\frac{1}{4C_3}. \end{align*}\]

As a result, if we set

\[\begin{align*} \eta &=\min\left(\frac{1}{L},\frac{\epsilon\sqrt{c}}{2\sigma},\frac{\epsilon c'}{4\sqrt{c}\sigma'},\frac{\epsilon \sqrt{c'}}{4\sigma''},\frac{c'}{2\sqrt{cC_1}}\right)\\ &=\min\left(\frac{1}{L},\frac{\epsilon}{4\sqrt{C_2}\sigma},\frac{\epsilon \sqrt{C_2}}{8C_3\sigma'},\frac{\epsilon}{8\sqrt{C_3}\sigma''},\frac{\sqrt{C_2}}{4C_3\sqrt{C_1}}\right)\\ \beta&=\frac{\epsilon\eta\sqrt{C_2}}{\sigma},\quad \gamma=\min\left(\frac{\epsilon\eta\sqrt{C_2}}{\sigma'},\frac{\epsilon\eta\sqrt{C_3}}{\sigma''},\frac{C_2}{2C_3C_1}\right), \end{align*}\]

we have

\[\begin{equation*} \sum_{t=0}^{T-1}\frac{1}{T}(B_t+\frac{1}{2}\Gamma_t)\leq \epsilon^2. \end{equation*}\]

Plugging the values of \(\eta\) into the requirement of \(T\) yields the order of \(T\).

■

Proof

Let us recall the three inequalities in Lemma 4.14, Lemma 4.13 and Lemma 4.12:

\[\begin{align*} (*)\;&F(\mathbf{w}_{t+1})\leq F(\mathbf{w}_t)+\eta_tG_2^2L_1^2\|\mathbf{u}_{t}-g(\mathbf{w}_t)\|_2^2+\eta_t\|\mathbf{v}_{t}-\mathcal{M}_t\|^2-\frac{\eta_t}{2}\|\nabla F(\mathbf{w}_t)\|_2^2\\ &\quad -\frac{1}{4\eta_t}\|\mathbf{w}_{t+1}-\mathbf{w}_t\|_2^2,\\ (\sharp)\;&\mathbb E\left[\|\mathbf{v}_{t}-\mathcal{M}_t\|^2\right]\leq \mathbb E[(1-\beta_t)\|\mathbf{v}_{t-1}-\mathcal{M}_{t-1}\|_2^2]+16G_2^2L_1^2\gamma_{t}^2\|\mathbf{u}_{t-1}-g(\mathbf{w}_{t-1})\|_2^2] \\ &\quad +\mathbb E[(24G_2^4L_1^2+4G_1^2L_2^2)\|\mathbf{w}_t-\mathbf{w}_{t-1}\|_2^2+2\beta_t^2\sigma^2+8G_2^2L_1^2\gamma_{t}^2\sigma_0^2],\\ (\diamond)\;&\mathbb E_{\zeta_t}\left[\|\mathbf{u}_{t}-g(\mathbf{w}_t)\|^2 \right]\leq (1-\gamma_t)\|\mathbf{u}_{t-1}-g(\mathbf{w}_{t-1})\|_2^2\\ &\quad +\mathbb E[2G_2^2\|\mathbf{w}_{t}-\mathbf{w}_{t-1}\|_2^2+2\gamma_t^2\sigma_0^2]. \end{align*}\]

Define

\[\begin{align*} A_t &= 2(F(\mathbf{w}_t)-F_*),\quad B_t=\|\nabla F(\mathbf{w}_t)\|_2^2,\\ \Gamma_t&= \|\mathbf{v}_{t}\|_2^2/2,\quad \Delta_t=2\|\mathbf{v}_{t}-\mathcal M_t\|_2^2,\quad \delta_t=2L_1^2G_2^2\|\mathbf{u}_{t}-g(\mathbf{w}_t)\|_2^2. \end{align*}\]

They satisfy the three inequalities marked by \(*,\sharp,\diamond\) in Lemma 4.15 with \(C_1=O(1),C_2=O(G_2^4L_1^2+G_1^2L_2^2),C_3=O(L_1^2G_2^4),\sigma^2=O(G_2^2\sigma_1^2+G_1^2\sigma_2^2),\sigma'^2=O(L_1^2G_2^2\sigma_0^2),\sigma''^2=O(L_1^2G_2^2\sigma_0^2)\). Plugging these into Lemma 4.15, we can finish the proof.

■

← Go Back