SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\begin{array}{l} \mathbb {E} \left[ \left\| v _ {i} ^ {t} - \nabla f _ {i} (x ^ {t}) \right\| ^ {2} \right] = \mathbb {E} \left[ \left\| v _ {i} ^ {t - 1} - \nabla f _ {i} (x ^ {t}) + \eta (\nabla f _ {i} (x ^ {t}, \xi_ {i} ^ {t}) - v _ {i} ^ {t - 1}) \right\| ^ {2} \right] \\ = \mathbb {E} \left[ \mathbb {E} _ {\xi_ {i} ^ {t}} \left[ \left\| (1 - \eta) \left(v _ {i} ^ {t - 1} - \nabla f _ {i} \left(x ^ {t}\right)\right) + \eta \left(\nabla f _ {i} \left(x ^ {t}, \xi_ {i} ^ {t}\right) - \nabla f _ {i} \left(x ^ {t}\right)\right) \right\| ^ {2} \right] \right] \\ = (1 - \eta) ^ {2} \mathbb {E} \left[ \left\| v _ {i} ^ {t - 1} - \nabla f _ {i} \left(x ^ {t}\right) \right\| ^ {2} \right] + \eta^ {2} \mathbb {E} \left[ \left\| \nabla f _ {i} \left(x ^ {t}, \xi_ {i} ^ {t}\right) - \nabla f _ {i} \left(x ^ {t}\right) \right\| ^ {2} \right] \\ \leq \left. (1 - \eta) ^ {2} \left(1 + \frac {\eta}{2}\right) \mathbb {E} \left[ \left| \left| v _ {i} ^ {t - 1} - \nabla f _ {i} \left(x ^ {t - 1}\right) \right| \right| ^ {2} \right] \right. \\ + \left(1 + \frac {2}{\eta}\right) \mathbb {E} \left[ \left\| \nabla f _ {i} (x ^ {t - 1}) - \nabla f _ {i} (x ^ {t}) \right\| ^ {2} \right] + \eta^ {2} \sigma^ {2} \\ \leq (1 - \eta) \mathbb {E} \left[ \left\| v _ {i} ^ {t - 1} - \nabla f _ {i} (x ^ {t - 1}) \right\| ^ {2} \right] + \frac {3 L _ {i} ^ {2}}{\eta} \mathbb {E} \left[ \left\| x ^ {t} - x ^ {t + 1} \right\| ^ {2} \right] + \eta^ {2} \sigma^ {2}, \\ \end{array}
$$
where the first inequality holds by Young's inequality, and the last step uses smoothness of  $f_{i}(\cdot)$  (Assumption 1), which concludes the proof of (24).
For each  $t = 0, \ldots, T - 1$ , define a random vector  $\xi^t \coloneqq (\xi_1^t, \ldots, \xi_n^t)$  and denote by  $\nabla f(x^t, \xi^t) \coloneqq \frac{1}{n} \sum_{i=1}^{n} \nabla f_i(x^t, \xi_i^t)$ . Note that the entries of the random vector  $\xi^t$  are independent and  $\mathbb{E}_{\xi^t}[\nabla f(x^t, \xi^t)] = \nabla f(x^t)$ , then we have
$$
v ^ {t} = v ^ {t - 1} + \eta (\nabla f (x ^ {t}, \xi^ {t}) - v ^ {t - 1}),
$$
where  $v^t \coloneqq \frac{1}{n} \sum_{i=1}^{n} v_i^t$  is an auxiliary sequence. Therefore, we can similarly derive
$$
\begin{array}{l} \mathbb {E} \left[ \left\| v ^ {t} - \nabla f (x ^ {t}) \right\| ^ {2} \right] = \mathbb {E} \left[ \left\| v ^ {t - 1} - \nabla f (x ^ {t}) + \eta (\nabla f (x ^ {t}, \xi^ {t}) - v ^ {t - 1}) \right\| ^ {2} \right] \\ = \mathbb {E} \left[ \mathbb {E} _ {\xi^ {t}} \left[ \left\| (1 - \eta) \left(v ^ {t - 1} - \nabla f (x ^ {t})\right) + \eta \left(\nabla f (x ^ {t}, \xi^ {t}) - \nabla f (x ^ {t})\right) \right\| ^ {2} \right] \right] \\ \end{array}
$$
$$
\begin{array}{l} { = } { ( 1 - \eta ) ^ { 2 } \mathbb { E } \left[ \left\| v ^ { t - 1 } - \nabla f ( x ^ { t } ) \right\| ^ { 2 } \right] + \eta ^ { 2 } \mathbb { E } \left[ \left\| \nabla f ( x ^ { t } , \xi ^ { t } ) - \nabla f ( x ^ { t } ) \right\| ^ { 2 } \right] } \\ \leq (1 - \eta) ^ {2} \left(1 + \frac {\eta}{2}\right) \mathbb {E} \left[ \left\| v ^ {t - 1} - \nabla f \left(x ^ {t - 1}\right) \right\| ^ {2} \right] \\ + \left(1 + \frac {2}{\eta}\right) \mathbb {E} \left[ \left\| \nabla f (x ^ {t - 1}) - \nabla f (x ^ {t}) \right\| ^ {2} \right] + \frac {\eta^ {2} \sigma^ {2}}{n} \\ \leq (1 - \eta) \mathbb {E} \left[ \left\| v ^ {t - 1} - \nabla f (x ^ {t - 1}) \right\| ^ {2} \right] + \frac {3 L ^ {2}}{\eta} \mathbb {E} \left[ \left\| x ^ {t} - x ^ {t - 1} \right\| ^ {2} \right] + \frac {\eta^ {2} \sigma^ {2}}{n}, \\ \end{array}
$$
where the last step uses smoothness of  $f(\cdot)$  (Assumption 1), which concludes the proof of (25).
F.2 Controlling the error of contractive compression and momentum estimator
Lemma 3. Let Assumption 1 be satisfied, and suppose  $\mathcal{C}$  is a contractive compressor with  $\alpha \leq \frac{1}{2}$ . For every  $i = 1, \ldots, n$ , let the sequences  $\{v_i^t\}_{t \geq 0}$  and  $\{g_i^t\}_{t \geq 0}$  be updated via
$$
v _ {i} ^ {t} = v _ {i} ^ {t - 1} + \eta (\nabla f _ {i} (x ^ {t}, \xi_ {i} ^ {t}) - v _ {i} ^ {t - 1}),
$$
$$
g _ {i} ^ {t} = g _ {i} ^ {t - 1} + \mathcal {C} \left(v _ {i} ^ {t} - g _ {i} ^ {t - 1}\right),
$$
Then for every  $i = 1, \ldots, n$  and  $t \geq 0$  it holds
$$
\begin{array}{l} \mathbb {E} \left[ \left\| g _ {i} ^ {t} - v _ {i} ^ {t} \right\| ^ {2} \right] \leq \left(1 - \frac {\alpha}{2}\right) \mathbb {E} \left[ \left\| g _ {i} ^ {t - 1} - v _ {i} ^ {t - 1} \right\| ^ {2} \right] + \frac {4 \eta^ {2}}{\alpha} \mathbb {E} \left[ \left\| v _ {i} ^ {t - 1} - \nabla f _ {i} (x ^ {t - 1}) \right\| ^ {2} \right] \\ + \frac {4 L _ {i} ^ {2} \eta^ {2}}{\alpha} \mathbb {E} \left[ \left\| x ^ {t} - x ^ {t - 1} \right\| ^ {2} \right] + \eta^ {2} \sigma^ {2}. \tag {26} \\ \end{array}
$$
Proof. By the update rules of  $g_{i}^{t}$  and  $v_{i}^{t}$ , we derive
$$
\begin{array}{l} \mathbb {E} \left[ \left\| g _ {i} ^ {t} - v _ {i} ^ {t} \right\| ^ {2} \right] = \mathbb {E} \left[ \left\| g _ {i} ^ {t - 1} - v _ {i} ^ {t} + \mathcal {C} \left(v _ {i} ^ {t} - g _ {i} ^ {t - 1}\right) \right\| ^ {2} \right] \\ = \mathbb {E} \left[ \mathbb {E} _ {\mathcal {C}} \left[ \left\| \mathcal {C} \left(v _ {i} ^ {t} - g _ {i} ^ {t - 1}\right) - \left(v _ {i} ^ {t} - g _ {i} ^ {t - 1}\right) \right\| ^ {2} \right] \right] \\ \stackrel {(i)} {\leq} \quad (1 - \alpha) \mathbb {E} \left[ \left\| v _ {i} ^ {t} - g _ {i} ^ {t - 1} \right\| ^ {2} \right] \\ \stackrel {(i i)} {=} (1 - \alpha) \mathbb {E} \left[ \left\| v _ {i} ^ {t - 1} - g _ {i} ^ {t - 1} + \eta (\nabla f _ {i} (x ^ {t}, \xi_ {i} ^ {t}) - v _ {i} ^ {t - 1}) \right\| ^ {2} \right] \\ = (1 - \alpha) \mathbb {E} \left[ \mathbb {E} _ {\xi_ {i} ^ {t}} \left[ \left\| v _ {i} ^ {t - 1} - g _ {i} ^ {t - 1} + \eta \left(\nabla f _ {i} \left(x ^ {t}, \xi_ {i} ^ {t}\right) - v _ {i} ^ {t - 1}\right) \right\| ^ {2} \right] \right] \\ = (1 - \alpha) \mathbb {E} \left[ \left\| v _ {i} ^ {t - 1} - g _ {i} ^ {t - 1} + \eta (\nabla f _ {i} (x ^ {t}) - v _ {i} ^ {t - 1}) \right\| ^ {2} \right] \\ + (1 - \alpha) \eta^ {2} \mathbb {E} \left[ \left\| \nabla f _ {i} (x ^ {t}, \xi_ {i} ^ {t}) - \nabla f _ {i} (x ^ {t}) \right\| ^ {2} \right] \\ \end{array}
$$
$$
\begin{array}{l} \stackrel {(i i i)} {\leq} (1 - \alpha) (1 + \rho) \mathbb {E} \left[ \left\| v _ {i} ^ {t - 1} - g _ {i} ^ {t - 1} \right\| ^ {2} \right] + (1 - \alpha) (1 + \rho^ {- 1}) \eta^ {2} \mathbb {E} \left[ \left\| v _ {i} ^ {t - 1} - \nabla f _ {i} (x ^ {t}) \right\| ^ {2} \right] \\ + (1 - \alpha) \eta^ {2} \sigma^ {2} \\ \end{array}
$$
$$
\stackrel {(i v)} {=} (1 - \theta) \mathbb {E} \left[ \left\| g _ {i} ^ {t - 1} - v _ {i} ^ {t - 1} \right\| ^ {2} \right] + \beta \eta^ {2} \mathbb {E} \left[ \left\| v _ {i} ^ {t - 1} - \nabla f _ {i} (x ^ {t}) \right\| ^ {2} \right] + (1 - \alpha) \eta^ {2} \sigma^ {2}
$$
$$
\begin{array}{l} \stackrel {(v)} {\leq} (1 - \theta) \mathbb {E} \left[ \left\| g _ {i} ^ {t - 1} - v _ {i} ^ {t - 1} \right\| ^ {2} \right] + 2 \beta \eta^ {2} \mathbb {E} \left[ \left\| v _ {i} ^ {t - 1} - \nabla f _ {i} (x ^ {t - 1}) \right\| ^ {2} \right] \\ + 2 \beta \eta^ {2} \mathbb {E} \left[ \left\| \nabla f _ {i} (x ^ {t}) - \nabla f _ {i} (x ^ {t - 1}) \right\| ^ {2} \right] + \eta^ {2} \sigma^ {2} \\ \end{array}
$$
$$
\begin{array}{l} \leq (1 - \theta) \mathbb {E} \left[ \left\| g _ {i} ^ {t - 1} - v _ {i} ^ {t - 1} \right\| ^ {2} \right] + 2 \beta \eta^ {2} \mathbb {E} \left[ \left\| v _ {i} ^ {t - 1} - \nabla f _ {i} (x ^ {t - 1}) \right\| ^ {2} \right] \\ + 2 \beta L _ {i} ^ {2} \eta^ {2} \mathbb {E} \left[ \left\| x ^ {t} - x ^ {t - 1} \right\| ^ {2} \right] + \eta^ {2} \sigma^ {2}, \\ \end{array}
$$
where  $(i)$  is due to the definition of a contractive compressor (Definition 1),  $(ii)$  follows by the update rule of  $v_{i}^{t}$ ,  $(iii)$  and  $(v)$  hold by Young's inequality for any  $\rho >0$ . In  $(iv)$ , we introduced the notation
$\theta \coloneqq 1 - (1 - \alpha)(1 + \rho)$ , and  $\beta \coloneqq (1 - \alpha)(1 + \rho^{-1})$ . The last step follows by smoothness of  $f_{i}(\cdot)$  (Assumption 1). The proof is complete by the choice  $\rho = \alpha /2$ , which guarantees  $1 - \theta \leq 1 - \alpha /2$  and  $2\beta \leq 4 / \alpha$ .
G Further Improvement Using Double Momentum (Proof of Corollary 3)
Algorithm 3 EF21-SGD2M (Error Feedback 2021 Enhanced with Double Momentum)
1: Input: starting point  $x^0$ , step-size  $\gamma > 0$ , parameter  $\eta \in (0,1]$ , initial batch size  $B_{\mathrm{init}}$   
2: Initialize  $u_i^0 = v_i^0 = g_i^0 = \frac{1}{B_{\mathrm{init}}} \sum_{j=1}^{B_{\mathrm{init}}} \nabla f_i(x^0, \xi_{i,j}^0)$  for  $i = 1, \dots, n$ ;  $g^0 = \frac{1}{n} \sum_{i=1}^{n} g_i^0$   
3: for  $t = 0,1,2, \ldots, T-1$  do  
4: Master computes  $x^{t+1} = x^t - \gamma g^t$  and broadcasts  $x^{t+1}$  to all nodes  
5: for all nodes  $i = 1, \ldots, n$  in parallel do  
6: Compute the first momentum estimator  $v_i^{t+1} = (1 - \eta)v_i^t + \eta \nabla f_i(x^{t+1}, \xi_i^{t+1})$   
7: Compute the second momentum estimator  $u_i^{t+1} = (1 - \eta)u_i^t + \eta v_i^{t+1}$   
8: Compress  $c_i^{t+1} = C(u_i^{t+1} - g_i^t)$  and send  $c_i^{t+1}$  to the master  
9: Update local state  $g_i^{t+1} = g_i^t + c_i^{t+1}$   
10: end for  
11: Master computes  $g^{t+1} = \frac{1}{n} \sum_{i=1}^{n} g_i^{t+1}$  via  $g^{t+1} = g^t + \frac{1}{n} \sum_{i=1}^{n} c_i^{t+1}$   
12: end for
In this section, we state the detailed version of Corollary 3 in Theorem 5, followed by its formal proof. Notice that the key reason for the sample complexity improvement of the double momentum variant compared to EF21-SGDM (Theorem 3) is that in (27), one of the terms has better dependence on  $\eta$  compared to (9) in Theorem 3, i.e.,  $\eta^4\sigma^2/\alpha$  instead of  $\eta^2\sigma^2/\alpha$ . As a result, this term is dominated by other terms and vanishes in Corollary 3.
Theorem 5. Let Assumptions 1 and 2 hold. Let  $\hat{\boldsymbol{x}}^T$  be sampled uniformly at random from the iterates of the method. Let Algorithm 3 run with a contractive compressor. For all  $\eta \in (0,1]$  and  $B_{\mathrm{init}} \geq 1$ , with  $\gamma \leq \min \left\{\frac{\alpha}{60\widetilde{L}}, \frac{\eta}{16L}\right\}$ , we have
$$
\mathbb {E} \left[ \left\| \nabla f (\hat {x} ^ {T}) \right\| ^ {2} \right] \leq \mathcal {O} \left(\frac {\Psi_ {0}}{\gamma T} + \frac {\eta^ {3} \sigma^ {2}}{\alpha^ {2}} + \frac {\eta^ {4} \sigma^ {2}}{\alpha} + \frac {\eta \sigma^ {2}}{n}\right), \tag {27}
$$
where  $\Psi_0\coloneqq \delta_0 + \frac{\gamma}{\eta}\mathbb{E}\left[\left\| v^0 -\nabla f(x^0)\right\| ^2\right] + \frac{\gamma\eta^4}{\alpha^2}\frac{1}{n}\sum_{i = 1}^{n}\mathbb{E}\left[\left\| v_i^0 -\nabla f_i(x^0)\right\| ^2\right]$ . Setting initial batch size  $B_{\mathrm{init}} = \left\lceil \frac{\sigma^2}{L\delta_0}\right\rceil$  step-size and momentum parameters
$$
\gamma = \min  \left\{\frac {\alpha}{6 0 \widetilde {L}}, \frac {\eta}{1 6 L} \right\}, \quad \eta = \min  \left\{1, \left(\frac {L \delta_ {0} \alpha^ {2}}{\sigma^ {2} T}\right) ^ {1 / 4}, \left(\frac {L \delta_ {0} n}{\sigma^ {2} T}\right) ^ {1 / 2}, \frac {\alpha \sqrt {L \delta_ {0} B _ {\mathrm {i n i t}}}}{\sigma} \right\}, \tag {28}
$$
$$
\frac {1}{T} \sum_ {t = 0} ^ {T - 1} \mathbb {E} \left[ \left\| \nabla f (x ^ {t}) \right\| ^ {2} \right] \leq \mathcal {O} \left(\frac {\widetilde {L} \delta_ {0}}{\alpha T} + \left(\frac {L \delta_ {0} \sigma^ {2 / 3}}{\alpha^ {2 / 3} T}\right) ^ {3 / 4} + \left(\frac {L \delta_ {0} \sigma^ {2}}{n T}\right) ^ {1 / 2}\right).
$$
Proof. In order to control the error between  $g^{t}$  and  $\nabla f(x^{t})$ , we decompose it into three terms
$$
\begin{array}{l} {\left\| g ^ {t} - \nabla f (x ^ {t}) \right\| ^ {2}} \leq {3 \left\| g ^ {t} - u ^ {t} \right\| ^ {2} + 3 \left\| u ^ {t} - v ^ {t} \right\| ^ {2} + 3 \left\| v ^ {t} - \nabla f (x ^ {t}) \right\| ^ {2}} \\ \leq 3 \frac {1}{n} \sum_ {i = 1} ^ {n} \left\| g _ {i} ^ {t} - u _ {i} ^ {t} \right\| ^ {2} + 3 \left\| u ^ {t} - v ^ {t} \right\| ^ {2} + 3 \left\| v ^ {t} - \nabla f (x ^ {t}) \right\| ^ {2}, \\ \end{array}
$$
where we define the sequences  $v^{t} \coloneqq \frac{1}{n}\sum_{i = 1}^{n}v_{i}^{t}$  and  $u^t \coloneqq \frac{1}{n}\sum_{i = 1}^{n}u_i^t$ . In the following, we develop a recursion for each term above separately.
Part I. Controlling the error of momentum estimator for each  $v_{i}^{t}$  and on average for  $v^{t}$ . Denote  $P_{t} \coloneqq \mathbb{E}\left[\| v^{t} - \nabla f(x^{t})\|^{2}\right]$ ,  $\widetilde{P}_{t} \coloneqq \frac{1}{n}\sum_{i = 1}^{n}\mathbb{E}\left[\| v_{i}^{t} - \nabla f_{i}(x^{t})\|^{2}\right]$ ,  $R_{t} \coloneqq \mathbb{E}\left[\| x^{t} - x^{t + 1}\|^{2}\right]$ . Similarly to Part I of the proof of Theorem 3, we have
$$
\frac {1}{T} \sum_ {t = 0} ^ {T - 1} \widetilde {P} _ {t} \leq \frac {3 \widetilde {L} ^ {2}}{\eta^ {2}} \frac {1}{T} \sum_ {t = 0} ^ {T - 1} R _ {t} + \eta \sigma^ {2} + \frac {1}{\eta T} \widetilde {P} _ {0}, \tag {29}
$$
$$
\frac {1}{T} \sum_ {t = 0} ^ {T - 1} P _ {t} \leq \frac {3 L ^ {2}}{\eta^ {2}} \frac {1}{T} \sum_ {t = 0} ^ {T - 1} R _ {t} + \frac {\eta \sigma^ {2}}{n} + \frac {1}{\eta T} P _ {0}. \tag {30}
$$
Part II (a). Controlling the error of the second momentum estimator for each  $u_{i}^{t}$ . Recall that by Lemma 4-(37), we have for each  $i = 1, \dots, n$ , and any  $0 < \eta \leq 1$  and  $t \geq 0$
$$
\begin{array}{l} \mathbb {E} \left[ \left\| u _ {i} ^ {t + 1} - v _ {i} ^ {t + 1} \right\| ^ {2} \right] \leq (1 - \eta) \mathbb {E} \left[ \left\| u _ {i} ^ {t} - v _ {i} ^ {t} \right\| ^ {2} \right] + 6 \eta \mathbb {E} \left[ \left\| v _ {i} ^ {t} - \nabla f _ {i} (x ^ {t}) \right\| ^ {2} \right] \\ + 6 L _ {i} ^ {2} \eta \mathbb {E} \left[ \left\| x ^ {t + 1} - x ^ {t} \right\| ^ {2} \right] + \eta^ {2} \sigma^ {2}, \tag {31} \\ \end{array}
$$
Averaging inequalities (31) over  $i = 1,\dots ,n$  and denoting  $\widetilde{Q}_t\coloneqq \frac{1}{n}\sum_{i = 1}^n\mathbb{E}\left[\| u_i^t -v_i^t\| ^2\right]$ , we have
$$
\widetilde {Q} _ {t + 1} \leq (1 - \eta) \widetilde {Q} _ {t} + 6 \eta \widetilde {P} _ {t} + 6 \widetilde {L} ^ {2} \eta R _ {t} + \eta^ {2} \sigma^ {2}.
$$
Summing up the above inequalities for  $t = 0,\dots ,T - 1$  , we derive
$$
\begin{array}{l} \frac {1}{T} \sum_ {t = 0} ^ {T - 1} \widetilde {Q} _ {t} \leq \frac {6}{T} \sum_ {t = 0} ^ {T - 1} \widetilde {P} _ {t} + 6 \widetilde {L} ^ {2} \frac {1}{T} \sum_ {t = 0} ^ {T - 1} R _ {t} + \eta \sigma^ {2} + \frac {1}{\eta T} \widetilde {Q} _ {0} \\ \leq \left(\frac {6 \cdot 3 \widetilde {L} ^ {2}}{\eta^ {2}} + 6 \widetilde {L} ^ {2}\right) \frac {1}{T} \sum_ {t = 0} ^ {T - 1} R _ {t} + 7 \eta \sigma^ {2} + \frac {1}{\eta T} \widetilde {Q} _ {0} + \frac {6}{\eta T} \widetilde {P} _ {0} \\ \leq \frac {1 9 \widetilde {L} ^ {2}}{\eta^ {2}} \frac {1}{T} \sum_ {t = 0} ^ {T - 1} R _ {t} + 7 \eta \sigma^ {2} + \frac {6}{\eta T} \widetilde {P} _ {0}, \tag {32} \\ \end{array}
$$
where we used (29), the bound  $\eta \leq 1$ , and  $u_i^0 = v_i^0$  for  $i = 1, \ldots, n$ .
Part II (b). Controlling the error of the second momentum estimator  $u^t$  (on average). Similarly by Lemma 4-(38), we have for any  $0 < \eta \leq 1$  and  $t \geq 0$
$$
\begin{array}{l} \mathbb {E} \left[ \left\| u ^ {t + 1} - v ^ {t + 1} \right\| ^ {2} \right] \leq (1 - \eta) \mathbb {E} \left[ \left\| u ^ {t} - v ^ {t} \right\| ^ {2} \right] + 6 \eta \mathbb {E} \left[ \left\| v ^ {t} - \nabla f (x ^ {t}) \right\| ^ {2} \right] \\ + 6 L ^ {2} \eta \mathbb {E} \left[ \left\| x ^ {t + 1} - x ^ {t} \right\| ^ {2} \right] + \frac {\eta^ {2} \sigma^ {2}}{n}, \\ \end{array}
$$
Summing up the above inequalities for  $t = 0, \ldots, T - 1$ , and denoting  $Q_{t} \coloneqq \mathbb{E}\left[\| u^{t} - v^{t}\|^{2}\right]$ , we derive
$$
\begin{array}{l} \frac {1}{T} \sum_ {t = 0} ^ {T - 1} Q _ {t} \leq \frac {6}{T} \sum_ {t = 0} ^ {T - 1} P _ {t} + 6 L ^ {2} \frac {1}{T} \sum_ {t = 0} ^ {T - 1} R _ {t} + \eta \sigma^ {2} + \frac {1}{\eta T} Q _ {0} \\ \leq \left(\frac {6 \cdot 3 L ^ {2}}{\eta^ {2}} + 6 L ^ {2}\right) \frac {1}{T} \sum_ {t = 0} ^ {T - 1} R _ {t} + \frac {7 \eta \sigma^ {2}}{n} + \frac {1}{\eta T} Q _ {0} + \frac {6}{\eta T} P _ {0} \\ \leq \frac {1 9 L ^ {2}}{\eta^ {2}} \frac {1}{T} \sum_ {t = 0} ^ {T - 1} R _ {t} + \frac {7 \eta \sigma^ {2}}{n} + \frac {6}{\eta T} P _ {0}, \tag {33} \\ \end{array}
$$
where we used (30), the bound  $\eta \leq 1$ , and  $u^0 = v^0$ .
Part III. Controlling the error of contractive compressor and the double momentum estimator.
By Lemma 5 we have for each  $i = 1, \dots, n$ , and any  $0 < \eta \leq 1$  and  $t \geq 0$
$$
\begin{array}{l} \mathbb {E} \left[ \left\| g _ {i} ^ {t + 1} - u _ {i} ^ {t + 1} \right\| ^ {2} \right] \leq \left(1 - \frac {\alpha}{2}\right) \mathbb {E} \left[ \left\| g _ {i} ^ {t} - u _ {i} ^ {t} \right\| ^ {2} \right] + \frac {6 \eta^ {2}}{\alpha} \mathbb {E} \left[ \left\| u _ {i} ^ {t} - v _ {i} ^ {t} \right\| ^ {2} \right] \tag {34} \\ + \frac {6 \eta^ {4}}{\alpha} \mathbb {E} \left[ \left\| v _ {i} ^ {t} - \nabla f _ {i} (x ^ {t}, \xi_ {i} ^ {t}) \right\| ^ {2} \right] + \frac {6 L _ {i} ^ {2} \eta^ {4}}{\alpha} \mathbb {E} \left[ \left\| x ^ {t} - x ^ {t + 1} \right\| ^ {2} \right] + \eta^ {4} \sigma^ {2}. \\ \end{array}
$$
Averaging inequalities (34) over  $i = 1, \ldots, n$ , denoting  $\widetilde{V}_t := \frac{1}{n} \sum_{i=1}^n \mathbb{E}\left[\|g_i^t - u_i^t\|^2\right]$ , and summing up the resulting inequality for  $t = 0, \ldots, T-1$ , we obtain
$$
\frac {1}{T} \sum_ {t = 0} ^ {T - 1} \widetilde {V} _ {t} \leq \frac {1 2 \eta^ {2}}{\alpha^ {2}} \frac {1}{T} \sum_ {t = 0} ^ {T - 1} \widetilde {Q} _ {t} + \frac {1 2 \eta^ {4}}{\alpha^ {2}} \frac {1}{T} \sum_ {t = 0} ^ {T - 1} \widetilde {P} _ {t} + \frac {1 2 \widetilde {L} ^ {2} \eta^ {4}}{\alpha^ {2}} \frac {1}{T} \sum_ {t = 0} ^ {T - 1} R _ {t} + \frac {2 \eta^ {4} \sigma^ {2}}{\alpha}
$$
$$
\begin{array}{l} \leq \frac {1 2 \eta^ {2}}{\alpha^ {2}} \left(\frac {1 9 \widetilde {L} ^ {2}}{\eta^ {2}} \frac {1}{T} \sum_ {t = 0} ^ {T - 1} R _ {t} + 7 \eta \sigma^ {2}\right) + \frac {1 2 \eta^ {4}}{\alpha^ {2}} \left(\frac {3 \widetilde {L} ^ {2}}{\eta^ {2}} \frac {1}{T} \sum_ {t = 0} ^ {T - 1} R _ {t} + \eta \sigma^ {2}\right) \\ + \frac {1 2 \widetilde {L} ^ {2} \eta^ {4}}{\alpha^ {2}} \frac {1}{T} \sum_ {t = 0} ^ {T - 1} R _ {t} + \frac {2 \eta^ {4} \sigma^ {2}}{\alpha} + \frac {1 2 \eta^ {4}}{\alpha^ {2} T} \widetilde {P} _ {0} \\ \leq \frac {1 2 \cdot 1 9 \widetilde {L} ^ {2}}{\alpha^ {2}} \frac {1}{T} \sum_ {t = 0} ^ {T - 1} R _ {t} + \frac {1 2 \cdot 7 \eta^ {3} \sigma^ {2}}{\alpha^ {2}} + \frac {1 2 \cdot 3 \widetilde {L} ^ {2} \eta^ {2}}{\alpha^ {2}} \frac {1}{T} \sum_ {t = 0} ^ {T - 1} R _ {t} + \frac {1 2 \eta^ {5} \sigma^ {2}}{\alpha^ {2}} \\ + \frac {1 2 \widetilde {L} ^ {2} \eta^ {4}}{\alpha^ {2}} \frac {1}{T} \sum_ {t = 0} ^ {T - 1} R _ {t} + \frac {2 \eta^ {4} \sigma^ {2}}{\alpha} + \frac {1 2 \eta^ {4}}{\alpha^ {2} T} \widetilde {P} _ {0} \\ \leq \frac {2 7 6 \widetilde {L} ^ {2}}{\alpha^ {2}} \frac {1}{T} \sum_ {t = 0} ^ {T - 1} R _ {t} + \frac {8 4 \eta^ {3} \sigma^ {2}}{\alpha^ {2}} + \frac {1 2 \eta^ {5} \sigma^ {2}}{\alpha^ {2}} + \frac {2 \eta^ {4} \sigma^ {2}}{\alpha} + \frac {1 2 \eta^ {4}}{\alpha^ {2} T} \widetilde {P} _ {0} \tag {35} \\ \end{array}
$$
Part IV. Combining steps I, II and III with descent lemma. By smoothness (Assumption 1) of  $f(\cdot)$  it follows from Lemma 1 that for any  $\gamma \leq 1/(2L)$  we have
$$
\begin{array}{l} f \left(x ^ {t + 1}\right) \leq f \left(x ^ {t}\right) - \frac {\gamma}{2} \left\| \nabla f \left(x ^ {t}\right) \right\| ^ {2} - \frac {1}{4 \gamma} \left\| x ^ {t + 1} - x ^ {t} \right\| ^ {2} + \frac {\gamma}{2} \left\| g ^ {t} - \nabla f \left(x ^ {t}\right) \right\| ^ {2} \tag {36} \\ \leq f (x ^ {t}) - \frac {\gamma}{2} \| \nabla f (x ^ {t}) \| ^ {2} - \frac {1}{4 \gamma} \| x ^ {t + 1} - x ^ {t} \| ^ {2} \\ + \frac {3 \gamma}{2} \frac {1}{n} \sum_ {i = 1} ^ {n} \left\| g _ {i} ^ {t} - u _ {i} ^ {t} \right\| ^ {2} + \frac {3 \gamma}{2} \left\| u ^ {t} - v ^ {t} \right\| ^ {2} + \frac {3 \gamma}{2} \left\| v ^ {t} - \nabla f (x ^ {t}) \right\| ^ {2}. \\ \end{array}
$$
Subtracting  $f^*$  from both sides of (36), taking expectation and defining  $\delta_t \coloneqq \mathbb{E}\left[f(x^t) - f^*\right]$ , we derive
$$
\begin{array}{l} \mathbb {E} \left[ \left\| \nabla f (\hat {x} ^ {T}) \right\| ^ {2} \right] = \frac {1}{T} \sum_ {t = 0} ^ {T - 1} \mathbb {E} \left[ \left\| \nabla f (x ^ {t}) \right\| ^ {2} \right] \\ \leq \frac {2 \delta_ {0}}{\gamma T} + 3 \frac {1}{T} \sum_ {t = 0} ^ {T - 1} \widetilde {V} _ {t} + 3 \frac {1}{T} \sum_ {t = 0} ^ {T - 1} Q _ {t} + 3 \frac {1}{T} \sum_ {t = 0} ^ {T - 1} P _ {t} - \frac {1}{2 \gamma^ {2}} \frac {1}{T} \sum_ {t = 0} ^ {T - 1} R _ {t} \\ \stackrel {(i)} {\leq} \frac {2 \delta_ {0}}{\gamma T} + \frac {3 \cdot 2 7 6 \widetilde {L} ^ {2}}{\alpha^ {2}} \frac {1}{T} \sum_ {t = 0} ^ {T - 1} R _ {t} + \frac {3 \cdot 8 4 \eta^ {3} \sigma^ {2}}{\alpha^ {2}} + \frac {3 \cdot 1 2 \eta^ {5} \sigma^ {2}}{\alpha^ {2}} + \frac {3 \cdot 2 \eta^ {4} \sigma^ {2}}{\alpha} \\ + \frac {3 \cdot 1 9 L ^ {2}}{\eta^ {2}} \frac {1}{T} \sum_ {t = 0} ^ {T - 1} R _ {t} + \frac {3 \cdot 7 \eta \sigma^ {2}}{n} \\ + \frac {3 L ^ {2}}{\eta^ {2}} \frac {1}{T} \sum_ {t = 0} ^ {T - 1} R _ {t} + \frac {\eta \sigma^ {2}}{n} - \frac {1}{2 \gamma^ {2}} \frac {1}{T} \sum_ {t = 0} ^ {T - 1} R _ {t} \\ + \frac {3 6 \eta^ {4}}{\alpha^ {2} T} \widetilde {P} _ {0} + \frac {1 8}{\eta T} P _ {0} + \frac {3}{\eta T} P _ {0} \\ = \frac {2 \delta_ {0}}{\gamma T} + \frac {3 \cdot 8 4 \eta^ {3} \sigma^ {2}}{\alpha^ {2}} + \frac {3 \cdot 1 2 \eta^ {5} \sigma^ {2}}{\alpha^ {2}} + \frac {3 \cdot 2 \eta^ {4} \sigma^ {2}}{\alpha} + \frac {2 2 \eta \sigma^ {2}}{n} \\ + \left(\frac {6 0 L ^ {2}}{\eta^ {2}} + \frac {3 \cdot 2 7 6 \widetilde {L} ^ {2}}{\alpha^ {2}} - \frac {1}{2 \gamma^ {2}}\right) \frac {1}{T} \sum_ {t = 0} ^ {T - 1} R _ {t} \\ + \frac {3 6 \eta^ {4}}{\alpha^ {2} T} \widetilde {P} _ {0} + \frac {2 1}{\eta T} P _ {0} \\ = \frac {2 \delta_ {0}}{\gamma T} + \frac {2 8 8 \eta^ {3} \sigma^ {2}}{\alpha^ {2}} + \frac {6 \eta^ {4} \sigma^ {2}}{\alpha} + \frac {2 2 \eta \sigma^ {2}}{n} + \frac {3 6 \eta^ {4}}{\alpha^ {2} T} \widetilde {P} _ {0} + \frac {2 1}{\eta T} P _ {0}. \\ \end{array}
$$
where  $(i)$  holds due to (30), (35) and (33), the last two steps hold because of the assumption on the step-size, and  $\eta \leq 1$ , which completes the proof of the first part of Theorem.
Notice that it suffices to take the same initial batch-size as in the proof of the Theorem 3 in order to "remove"  $\widetilde{P}_0$  and  $P_0$  terms, since the power of  $\eta$  in front of  $\widetilde{P}_0$  is larger here compared to the proof of Theorem 3. The choice of the momentum parameter such that  $\eta \leq \left(\frac{L\delta_0\alpha^2}{\sigma^2T}\right)^{1/4}$ ,  $\eta \leq \left(\frac{L\delta_0n}{\sigma^2T}\right)^{1/2}$  ensures that  $\frac{\eta^3\sigma^2}{\alpha^2} \leq \frac{L\delta_0}{\eta T}$ , and  $\frac{\eta\sigma^2}{n} \leq \frac{L\delta_0}{\eta T}$ . Therefore, we can guarantee that the choice  $\eta = \min \left\{\frac{\alpha\sqrt{L\delta_0B_{\mathrm{init}}}}{\sigma}, \left(\frac{L\delta_0\alpha^2}{\sigma^2T}\right)^{1/4}, \left(\frac{L\delta_0n}{\sigma^2T}\right)^{1/2}\right\}$  ensures that
$$
\mathbb {E} \left[ \left\| \nabla f (\hat {x} ^ {T}) \right\| ^ {2} \right] \leq \mathcal {O} \left(\frac {\widetilde {L} \delta_ {0}}{\alpha T} + \left(\frac {L \delta_ {0} \sigma^ {2 / 3}}{\alpha^ {2 / 3} T}\right) ^ {3 / 4} + \left(\frac {L \delta_ {0} \sigma^ {2}}{n T}\right) ^ {1 / 2}\right).
$$
G.1 Controlling the error of second momentum estimator
Lemma 4. Let Assumption 1 be satisfied, and suppose  $0 < \eta \leq 1$ . For every  $i = 1, \dots, n$ , let the sequences  $\{v_i^t\}_{t \geq 0}$  and  $\{u_i^t\}_{t \geq 0}$  be updated via
$$
u _ {i} ^ {t} = u _ {i} ^ {t - 1} + \eta \left(v _ {i} ^ {t} - u _ {i} ^ {t - 1}\right).
$$
Define the sequences  $v^t \coloneqq \frac{1}{n} \sum_{i=1}^{n} v_i^t$  and  $u^t \coloneqq \frac{1}{n} \sum_{i=1}^{n} u_i^t$ . Then for every  $i = 1, \dots, n$  and  $t \geq 0$  it holds