SlowGuess/NeurIPS_2023_OCR huggingface dataset:
A procedure to generate stochastic quadratic optimization problems. In this section, we present an algorithm that generates quadratic optimization tasks. The formal description is provided in Algorithm 2. The idea is to take a predefined tridiagonal matrix and add noises to simulate the heterogeneous setting. Algorithm 2 returns matrices  $\mathbf{Q}_1,\dots ,\mathbf{Q}_n$  , vectors  $b_{1},\dots ,b_{n}$  , and a starting
(a)  $B = 8$
Figure 8: ResNet-18 on CIFAR10 dataset with  $n = 5$ .
Figure 9: Accuracy on the CIFAR10 test split.
(b)  $B = 25$
Algorithm	Test Accuracy
SGD	81.5 %
EF21-SGD	82.5 %
EF14-SGD	83.1 %
EF21-SGDM	83.3 %
point  $x^0$  such that the matrix  $\mathbf{Q} = \frac{1}{n}\sum_{i=1}^{n}\mathbf{Q}_i$  has the minimum eigenvalue  $\lambda_{\min}(\mathbf{Q}) = \lambda$ , where  $\lambda \geq 0$  is a parameter. Next, we define the functions  $f_i$  and stochastic gradients in the following way:
$$
f _ {i} (x) := \frac {1}{2} x ^ {\top} \mathbf {Q} _ {i} x - x ^ {\top} b _ {i}
$$
$$
\nabla f _ {i} (x, \xi) := \nabla f _ {i} (x) + \xi_ {i},
$$
for all  $x \in \mathbb{R}^d$  and  $i \in [n]$ . The noises  $\xi_i$  are i.i.d. samples from  $\mathcal{N}(0, \sigma)$ , where  $\sigma$  is a parameter.
C.3 Experiment 4: training neural network
We test algorithms on an image recognition task, CIFAR10 [Krizhevsky et al., 2009], with the ResNet-18 [He et al., 2016] deep neural network (the number of parameters  $d \approx 10^{7}$ ). We split CIFAR10 among 5 nodes, and take  $K = 2 \times 10^{6}$  in TopK. In all methods we finetune the step sizes. One can see that our findings in the low-scale experiments translate into large-scale experiments in Figure 8. With different batch sizes, EF21-SGD converges slower than EF21-SGDM and EF14-SGD, and our new method EF21-SGDM improves over EF14-SGD in Figure 8b. We checked the accuracies on the test dataset (see Table 9) and observed the same relations between algorithms (note that accuracies are far from the real SOTA because we turned off all augmentations and regularizations in training).
D Descent Lemma
Let us state the following lemma that is used in the analysis of nonconvex optimization methods.
Lemma 1 ([Li et al., 2021]). Let the function  $f(\cdot)$  be  $L$ -smooth and let  $x^{t + 1} = x^t - \gamma g^t$  for some vector  $g^t \in \mathbb{R}^d$  and a step-size  $\gamma > 0$ . Then we have
$$
f \left(x ^ {t + 1}\right) \leq f \left(x ^ {t}\right) - \frac {\gamma}{2} \left\| \nabla f \left(x ^ {t}\right) \right\| ^ {2} - \left(\frac {1}{2 \gamma} - \frac {L}{2}\right) \left\| x ^ {t + 1} - x ^ {t} \right\| ^ {2} + \frac {\gamma}{2} \left\| g ^ {t} - \nabla f \left(x ^ {t}\right) \right\| ^ {2}. \tag {13}
$$
E EF21-SGDM-ideal (Proof of Theorem 1 and Proposition 1)
We now state a slightly more general result than Theorem 1, which holds for EF21-SGDM-ideal method with any  $\eta \in (0,1]$ . The statement of Theorem 1 follows by setting  $\eta = 1$ , since in that case EF21-SGDM-ideal coincides with EF21-SGD-ideal (5a), (5aa). Recall that EF21-SGDM-ideal (distributed variant) has the following update rule:
$$
x ^ {t + 1} = x ^ {t} - \gamma g ^ {t}, \quad g ^ {t} = \frac {1}{n} \sum_ {i = 1} ^ {n} g _ {i} ^ {t}, \tag {14}
$$
$$
v _ {i} ^ {t + 1} = \nabla f _ {i} \left(x ^ {t + 1}\right) + \eta \left(\nabla f _ {i} \left(x ^ {t + 1}, \xi_ {i} ^ {t + 1}\right) - \nabla f _ {i} \left(x ^ {t + 1}\right)\right), \tag {15}
$$
$$
g _ {i} ^ {t + 1} = \nabla f _ {i} (x ^ {t + 1}) + \mathcal {C} \left(v _ {i} ^ {t + 1} - \nabla f _ {i} (x ^ {t + 1})\right).
$$
Theorem 4. Let  $L, \sigma > 0, 0 < \gamma \leq 1 / L, 0 < \eta \leq 1$  and  $n = 1$ . There exists a convex,  $L$ -smooth function  $f(\cdot)$ , a contractive compressor  $\mathcal{C}(\cdot)$  satisfying Definition 1, and an unbiased stochastic gradient with bounded variance  $\sigma^2$  such that if the method (14), (15) is run with a step-size  $\gamma$ , then for all  $T \geq 0$  and for all  $x^0 \in \{(0, x_{(2)}^0)^\top \in \mathbb{R}^2 | x_{(2)}^0 < 0\}$ , we have
$$
\mathbb {E} \left[ \left\| \nabla f (x ^ {T}) \right\| ^ {2} \right] \geq \frac {1}{6 0} \min  \left\{\eta^ {2} \sigma^ {2}, \left\| \nabla f (x ^ {0}) \right\| ^ {2} \right\}.
$$
Fix  $0 < \varepsilon \leq L / \sqrt{60}$  and  $x^0 = (0, -1)^\top$ . Additionally assume that  $n \geq 1$  and the variance of unbiased stochastic gradient is controlled by  $\sigma^2 / B$  for some  $B \geq 1$ . If  $B < \frac{\eta^2 \sigma^2}{60 \varepsilon^2}$ , then we have  $\mathbb{E}\left[\|\nabla f(x^T)\|\right] > \varepsilon$  for all  $T \geq 0$ .
Proof of Theorem 1. Part I. Consider  $f(x) = \frac{L}{2} \| x \|^2, x \in \mathbb{R}^2$ . For each iteration  $t \geq 0$ , let the random vector  $\xi^{t+1}$  be sampled uniformly at random from the set of vectors:
$$
z _ {1} = \left( \begin{array}{c} 2 \\ 0 \end{array} \right) \sqrt {\frac {3 \sigma^ {2}}{1 0}}, \quad z _ {2} = \left( \begin{array}{c} 0 \\ 1 \end{array} \right) \sqrt {\frac {3 \sigma^ {2}}{1 0}}, \quad z _ {3} = \left( \begin{array}{c} - 2 \\ - 1 \end{array} \right) \sqrt {\frac {3 \sigma^ {2}}{1 0}}.
$$
Define the stochastic gradient as  $\nabla f(x^t,\xi^t)\coloneqq \nabla f(x^t) + \xi^t = Lx^t +\xi^t$  . Notice that  $\mathbb{E}\left[\nabla f(x^{t},\xi^{t})\right] = \nabla f(x^{t})$  , and  $\mathbb{E}\left[\| \nabla f(x^t,\xi^t) - \nabla f(x^t)\| ^2\right] = \sigma^2$  . The update rule of method (14), (15) with such estimator is
$$
x ^ {t + 1} = x ^ {t} - \gamma g ^ {t} = x ^ {t} - L \gamma x ^ {t} - \gamma \mathcal {C} (\eta \xi^ {t}),
$$
where we choose  $\mathcal{C}(\cdot)$  as a Top1 compressor. Notice that  $\mathbb{E}\left[\xi^t\right] = (0,0)^\top$ , but
$$
\mathbb {E} \left[ \mathcal {C} (\xi^ {t}) \right] = \eta \sqrt {\frac {3 \sigma^ {2}}{1 0}} (0, 1 / 3) ^ {\top} \neq (0, 0) ^ {\top}.
$$
By setting the initial iterate to  $x^0 = (0, x_{(2)}^0)^\top$  for any  $x_{(2)}^0 < 0$ , we can derive
$$
\begin{array}{l} \mathbb {E} \left[ x ^ {T} \right] = (1 - L \gamma) ^ {T} x ^ {0} - \eta \sqrt {\frac {3 \sigma^ {2}}{1 0}} \left( \begin{array}{c} 0 \\ \frac {1}{3} \end{array} \right) \gamma \sum_ {t = 0} ^ {T - 1} (1 - L \gamma) ^ {t} \\ = \left(1 - L \gamma\right) ^ {T} \binom {0} {x _ {(2)} ^ {0}} + \frac {\eta}{L} \sqrt {\frac {\sigma^ {2}}{3 0}} \binom {0} {- 1} \left(1 - \left(1 - L \gamma\right) ^ {T}\right) \neq \binom {0} {0} \tag {16} \\ \end{array}
$$
for any  $0 \leq \gamma \leq 1 / L$  and any  $x_{(2)}^0 < 0$ . The inequality in (16) is because the first vector has strictly negative component  $x_{(2)}^0$ , and the second vector has non-positive second component when  $\gamma > 0$  and  $\sigma^2 > 0$ . Therefore, since  $\| \nabla f(x) \|^2 = \| Lx \|^2$ , we have
$$
\begin{array}{l} \mathbb {E} \left[ \left\| \nabla f (x ^ {T}) \right\| ^ {2} \right] = \mathbb {E} \left[ \left\| L x ^ {T} \right\| ^ {2} \right] \\ = \left\| \mathbb {E} \left[ L x ^ {T} \right] \right\| ^ {2} + \mathbb {E} \left[ \left\| L x ^ {T} - \mathbb {E} \left[ L x ^ {T} \right] \right\| ^ {2} \right] \\ \geq \left\| \mathbb {E} \left[ L x ^ {T} \right] \right\| ^ {2} \\ \stackrel {(i)} {=} \left. \left((1 - L \gamma) ^ {T} \| L x ^ {0} \| + \eta \sqrt {\frac {\sigma^ {2}}{3 0}} (1 - (1 - L \gamma) ^ {T})\right) ^ {2} \right. \\ \stackrel {(i i)} {\geq} (1 - L \gamma) ^ {2 T} \left\| \nabla f (x ^ {0}) \right\| ^ {2} + \frac {\eta^ {2} \sigma^ {2}}{3 0} (1 - (1 - L \gamma) ^ {T}) ^ {2} \\ \geq \frac {\left\| \nabla f (x ^ {0}) \right\| ^ {2} \eta^ {2} \sigma^ {2}}{3 0 \left\| \nabla f (x ^ {0}) \right\| ^ {2} + \eta^ {2} \sigma^ {2}} \\ \end{array}
$$
for all  $T \geq 1$ , where in (i) we used the form of vector  $\mathbb{E}\left[x^T\right]$  in (16), in (ii) we drop a non-negative cross term, and use  $\nabla f(x^0) = Lx^0$ . The last inequality follows by lower bounding a univariate quadratic function with respect to  $z \coloneqq (1 - L\gamma)^T$  for  $0 \leq z \leq 1$ , where optimal choice is  $z = \eta^2\sigma^2 / (30\|\nabla f(x^0)\|^2 + \eta^2\sigma^2)$ . It is left to note that  $\frac{xy}{x + y} \geq \frac{1}{2}\min\{x,y\}$  for all  $x,y > 0$ .
Part II. Fix  $n \geq 1$  and  $B \geq 1$ . Let at each node  $i = 1, \ldots, n$ , the random vectors  $\xi_i^t$  be sampled independently and uniformly form the set of vectors:
$$
z _ {1} = \left( \begin{array}{c} 2 \\ 0 \end{array} \right) \sqrt {\frac {3 \sigma^ {2}}{1 0 B}}, \quad z _ {2} = \left( \begin{array}{c} 0 \\ 1 \end{array} \right) \sqrt {\frac {3 \sigma^ {2}}{1 0 B}}, \quad z _ {3} = \left( \begin{array}{c} - 2 \\ - 1 \end{array} \right) \sqrt {\frac {3 \sigma^ {2}}{1 0 B}}.
$$
Define a random matrix  $\xi^t \coloneqq (\xi_1^t, \ldots, \xi_n^t)^\top$ . Then  $\mathbb{E}\left[\| \nabla f(x^t, \xi^t) - \nabla f(x^t) \|^2\right] = \frac{\sigma^2}{B}$ . The update of the method (14), (15) on the same function instance will take the form
$$
x ^ {t + 1} = x ^ {t} - \gamma \frac {1}{n} \sum_ {i = 1} ^ {n} g _ {i} ^ {t} = x ^ {t} - L \gamma x ^ {t} - \gamma \frac {1}{n} \sum_ {i = 1} ^ {n} \mathcal {C} (\eta \xi_ {i} ^ {t}).
$$
Notice that in this case, we still have
$$
\mathbb {E} \left[ \frac {1}{n} \sum_ {i = 1} ^ {n} \mathcal {C} \left(\eta \xi_ {i} ^ {t}\right) \right] = \frac {1}{n} \sum_ {i = 1} ^ {n} \mathbb {E} \left[ \mathcal {C} \left(\eta \xi_ {i} ^ {t}\right) \right] = \eta \sqrt {\frac {3 \sigma^ {2}}{1 0}} (0, 1 / 3) ^ {\top} \neq (0, 0) ^ {\top},
$$
which is independent (!) of  $n$ . Therefore, by similar derivations, we can conclude that
$$
\mathbb {E} \left[ \left\| \nabla f (x ^ {T}) \right\| ^ {2} \right] \geq \frac {1}{6 0} \min  \left\{\frac {\eta^ {2} \sigma^ {2}}{B}, \left\| \nabla f (x ^ {0}) \right\| ^ {2} \right\} > \varepsilon^ {2},
$$
where we use that  $B < \frac{\eta^2\sigma^2}{60\varepsilon^2}, \varepsilon \leq L / \sqrt{60}$ , and  $x^0 = (0, -1)^\top$ .
Proof of Proposition 1. By smoothness (Assumption 1) of  $f(\cdot)$  it follows from Lemma 1 that for  $\gamma \leq 1 / L$  we have
$$
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 {\gamma}{2} \left\| g ^ {t} - \nabla f \left(x ^ {t}\right) \right\| ^ {2}. \tag {17}
$$
Now it remains to control the last term, which is due to the error introduced by a contractive compressor and stochastic gradients. We have
$$
\begin{array}{l} \mathbb {E} \left[ \left\| g ^ {t} - \nabla f (x ^ {t}) \right\| ^ {2} \right] \stackrel {(i)} {=} \mathbb {E} \left[ \left\| \mathcal {C} \left(v ^ {t} - \nabla f (x ^ {t})\right) \right\| ^ {2} \right] \stackrel {(i i)} {=} \mathbb {E} \left[ \left\| \mathcal {C} \left(\eta \left(\nabla f (x ^ {t}, \xi^ {t}) - \nabla f (x ^ {t})\right)\right) \right\| ^ {2} \right] \\ \stackrel {(i i i)} {\leq} 2 \mathbb {E} \left[ \left\| \mathcal {C} \left(\eta (\nabla f (x ^ {t}, \xi^ {t}) - \nabla f (x ^ {t}))\right) - \eta (\nabla f (x ^ {t}, \xi^ {t}) - \nabla f (x ^ {t})) \right\| ^ {2} \right] \\ \end{array}
$$
$$
\begin{array}{l} + 2 \eta^ {2} \mathbb {E} \left[ \left\| \nabla f (x ^ {t}, \xi^ {t}) - \nabla f (x ^ {t}) \right\| ^ {2} \right] \\ \leq 2 (2 - \alpha) \eta^ {2} \mathbb {E} \left[ \left\| \nabla f (x ^ {t}, \xi^ {t}) - \nabla f (x ^ {t}) \right\| ^ {2} \right] \\ \leq \quad 4 \eta^ {2} \sigma^ {2}, \\ \end{array}
$$
where  $(i)$  and  $(ii)$  use the update rule (6),  $(iii)$  holds by Young's inequality, and the last two steps hold by Definition 1 and Assumption 2.
Subtracting  $f^*$  from both sides of (17), taking expectation and defining  $\delta_t \coloneqq \mathbb{E}\left[f(x^t) - f^*\right]$ , we derive
$$
\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} + 4 \eta^ {2} \sigma^ {2}.
$$
F EF21-SGDM (Proof of Theorems 2 and 3)
The statement of Theorem 2 follows directly from Theorem 3 and Remark 2. Let us prove Theorem 3.
Proof of Theorem 3. In order to control the error between  $g^t$  and  $\nabla f(x^t)$ , we decompose it into two terms
$$
\left\| g ^ {t} - \nabla f (x ^ {t}) \right\| ^ {2} \leq 2 \left\| g ^ {t} - v ^ {t} \right\| ^ {2} + 2 \left\| v ^ {t} - \nabla f (x ^ {t}) \right\| ^ {2} \leq 2 \frac {1}{n} \sum_ {i = 1} ^ {n} \left\| g _ {i} ^ {t} - v _ {i} ^ {t} \right\| ^ {2} + 2 \left\| v ^ {t} - \nabla f (x ^ {t}) \right\| ^ {2},
$$
and develop a recursion for each term above separately.
Part I (a). Controlling the error of momentum estimator for each  $v_{i}^{t}$ . Recall that by Lemma 2-(24), we have for each  $i = 1, \dots, n$ , and any  $0 < \eta \leq 1$  and  $t \geq 0$
$$
\mathbb {E} \left[ \left\| v _ {i} ^ {t + 1} - \nabla f _ {i} \left(x ^ {t + 1}\right) \right\| ^ {2} \right] \leq (1 - \eta) \mathbb {E} \left[ \left\| v _ {i} ^ {t} - \nabla f _ {i} \left(x ^ {t}\right) \right\| ^ {2} \right] + \frac {3 L _ {i} ^ {2}}{\eta} \mathbb {E} \left[ \left\| x ^ {t + 1} - x ^ {t} \right\| ^ {2} \right] + \eta^ {2} \sigma^ {2}. \tag {18}
$$
Averaging inequalities (18) over  $i = 1,\ldots ,n$  and denoting  $\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]$  we have
$$
\widetilde {P} _ {t + 1} \leq (1 - \eta) \widetilde {P} _ {t} + \frac {3 \widetilde {L} ^ {2}}{\eta} R _ {t} + \eta^ {2} \sigma^ {2}.
$$
Summing up the above inequality for  $t = 0,\dots ,T - 1$  , we derive
$$
\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 {19}
$$
Part I (b). Controlling the error of momentum estimator for  $v^t$  (on average). Similarly by Lemma 2-(25), we have for any  $0 < \eta \leq 1$  and  $t \geq 0$
$$
\mathbb {E} \left[ \left\| v ^ {t + 1} - \nabla f (x ^ {t + 1}) \right\| ^ {2} \right] \leq (1 - \eta) \mathbb {E} \left[ \left\| v ^ {t} - \nabla f (x ^ {t}) \right\| ^ {2} \right] + \frac {3 L ^ {2}}{\eta} \mathbb {E} \left[ \left\| x ^ {t + 1} - x ^ {t} \right\| ^ {2} \right] + \frac {\eta^ {2} \sigma^ {2}}{n},
$$
where  $v^t \coloneqq \frac{1}{n} \sum_{i=1}^{n} v_i^t$  is an auxiliary sequence.
Summing up the above inequality for  $t = 0, \dots, T - 1$ , and denoting  $P_{t} \coloneqq \mathbb{E}\left[\| v^{t} - \nabla f(x^{t})\|^{2}\right]$  we derive
$$
\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 {20}
$$
Part II. Controlling the error of contractive compressor and momentum estimator. By Lemma 3 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} - v _ {i} ^ {t + 1} \right\| ^ {2} \right] \leq \left(1 - \frac {\alpha}{2}\right) \mathbb {E} \left[ \left\| g _ {i} ^ {t} - v _ {i} ^ {t} \right\| ^ {2} \right] + \frac {4 \eta^ {2}}{\alpha} \mathbb {E} \left[ \left\| v _ {i} ^ {t} - \nabla f _ {i} (x ^ {t}) \right\| ^ {2} \right] \\ + \frac {4 L _ {i} ^ {2} \eta^ {2}}{\alpha} \mathbb {E} \left[ \left\| x ^ {t + 1} - x ^ {t} \right\| ^ {2} \right] + \eta^ {2} \sigma^ {2}. \tag {21} \\ \end{array}
$$
Averaging inequalities (21) over  $i = 1, \ldots, n$ , denoting  $\widetilde{V}_t := \frac{1}{n} \sum_{i=1}^n \mathbb{E}\left[\|g_i^t - v_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 {8 \eta^ {2}}{\alpha^ {2}} \frac {1}{T} \sum_ {t = 0} ^ {T - 1} \widetilde {P} _ {t} + \frac {8 \widetilde {L} ^ {2} \eta^ {2}}{\alpha^ {2}} \frac {1}{T} \sum_ {t = 0} ^ {T - 1} R _ {t} + \frac {2 \eta^ {2} \sigma^ {2}}{\alpha} + \frac {2}{\alpha T} \widetilde {V} _ {0}. \tag {22}
$$
Part III. Combining steps I and II 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 {23} \\ \leq f (x ^ {t}) - \frac {\gamma}{2} \left\| \nabla f (x ^ {t}) \right\| ^ {2} - \frac {1}{4 \gamma} \left\| x ^ {t + 1} - x ^ {t} \right\| ^ {2} + \gamma \frac {1}{n} \sum_ {i = 1} ^ {n} \left\| g _ {i} ^ {t} - v _ {i} ^ {t} \right\| ^ {2} + \gamma \left\| v ^ {t} - \nabla f (x ^ {t}) \right\| ^ {2}. \\ \end{array}
$$
Subtracting  $f^*$  from both sides of (23), taking expectation and defining  $\delta_t \coloneqq \mathbb{E}\left[f(x^t) - f^*\right]$ , we derive
$$
\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} + 2 \frac {1}{T} \sum_ {t = 0} ^ {T - 1} \widetilde {V} _ {t} + 2 \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}
$$
$$
\begin{array}{l} \stackrel {(i)} {\leq} \quad \frac {2 \delta_ {0}}{\gamma T} + \frac {1 6 \eta^ {2}}{\alpha^ {2}} \frac {1}{T} \sum_ {t = 0} ^ {T - 1} \widetilde {P} _ {t} + 2 \frac {1}{T} \sum_ {t = 0} ^ {T - 1} P _ {t} + \frac {4 \eta^ {2} \sigma^ {2}}{\alpha} \\ - \frac {\frac {1}{2} - \frac {1 6 \gamma^ {2} \widetilde {L} ^ {2} \eta^ {2}}{\alpha^ {2}}}{\gamma^ {2}} \frac {1}{T} \sum_ {t = 0} ^ {T - 1} R _ {t} \\ \end{array}
$$
$$
\begin{array}{l} \stackrel {(i i)} {\leq} \quad \frac {2 \delta_ {0}}{\gamma T} + \frac {1 6 \eta^ {3} \sigma^ {2}}{\alpha^ {2}} + \frac {4 \eta^ {2} \sigma^ {2}}{\alpha} + \frac {2 \eta \sigma^ {2}}{n} + \frac {4}{\alpha T} \widetilde {V} _ {0} \\ - \frac {\frac {1}{2} - \frac {1 6 \gamma^ {2} \widetilde {L} ^ {2} \eta^ {2}}{\alpha^ {2}} - \frac {6 \gamma^ {2} L ^ {2}}{\eta^ {2}} - \frac {4 8 \gamma^ {2} \widetilde {L} ^ {2}}{\alpha^ {2}}}{\gamma^ {2}} \frac {1}{T} \sum_ {t = 0} ^ {T - 1} R _ {t} + \frac {2}{\eta T} P _ {0} + \frac {1 6 \eta}{\alpha^ {2} T} \widetilde {P} _ {0} \\ \end{array}
$$
$$
\begin{array}{l} \stackrel {(i i i)} {\leq} \quad \frac {2 \delta_ {0}}{\gamma T} + \frac {1 6 \eta^ {3} \sigma^ {2}}{\alpha^ {2}} + \frac {4 \eta^ {2} \sigma^ {2}}{\alpha} + \frac {2 \eta \sigma^ {2}}{n} - \frac {\frac {1}{2} - \frac {6 \gamma^ {2} L ^ {2}}{\eta^ {2}} - \frac {6 4 \gamma^ {2} \tilde {L} ^ {2}}{\alpha^ {2}}}{\gamma^ {2}} \frac {1}{T} \sum_ {t = 0} ^ {T - 1} R _ {t} \\ + \frac {2}{\eta T} P _ {0} + \frac {1 6 \eta}{\alpha^ {2} T} \widetilde {P} _ {0} + \frac {4}{\alpha T} \widetilde {V} _ {0} \\ \leq \frac {2 \delta_ {0}}{\gamma T} + \frac {1 6 \eta^ {3} \sigma^ {2}}{\alpha^ {2}} + \frac {4 \eta^ {2} \sigma^ {2}}{\alpha} + \frac {2 \eta \sigma^ {2}}{n} + \frac {2}{\eta T} P _ {0} + \frac {1 6 \eta}{\alpha^ {2} T} \widetilde {P} _ {0} + \frac {4}{\alpha T} \widetilde {V} _ {0}. \\ \end{array}
$$
where  $(i)$  holds due to (22),  $(ii)$  utilizes (20), and  $(iii)$  follows by  $\eta \leq 1$ , and the last step holds due to the assumption on the step-size. We proved (9).
We now find the particular values of parameters. Since  $g_{i} = v_{i}$  for all  $i \in [n]$ , we have  $\widetilde{V}_0 = 0$ . Using  $v_{i}^{0} = \frac{1}{B_{\mathrm{init}}} \sum_{j=1}^{B_{\mathrm{init}}} \nabla f_{i}(x^{0}, \xi_{i,j}^{0})$  for all  $i = 1, \ldots, n$ , we have
$$
P _ {0} = \mathbb {E} \left[ \left\| v ^ {0} - \nabla f (x ^ {0}) \right\| ^ {2} \right] \leq \frac {\sigma^ {2}}{n B _ {\mathrm {i n i t}}} \text {a n d} \widetilde {P} _ {0} = \frac {1}{n} \sum_ {i = 1} ^ {n} \mathbb {E} \left[ \left\| v _ {i} ^ {0} - \nabla f _ {i} (x ^ {0}) \right\| ^ {2} \right] \leq \frac {\sigma^ {2}}{B _ {\mathrm {i n i t}}}.
$$
We can substitute the choice of  $\gamma$  and obtain
$$
\mathbb {E} \left[ \left\| \nabla f (\hat {x} ^ {T}) \right\| ^ {2} \right] = \mathcal {O} \left(\frac {\widetilde {L} \delta_ {0}}{\alpha T} + \frac {L \delta_ {0}}{\eta T} + \frac {\eta^ {3} \sigma^ {2}}{\alpha^ {2}} + \frac {\eta^ {2} \sigma^ {2}}{\alpha} + \frac {\eta \sigma^ {2}}{n} + \frac {\sigma^ {2}}{\eta n B _ {\mathrm {i n i t}} T} + \frac {\eta \sigma^ {2}}{\alpha^ {2} B _ {\mathrm {i n i t}} T}\right).
$$
Since  $B_{\mathrm{init}} \geq \frac{\sigma^2}{L\delta_0n}$ , we have
$$
\mathbb {E} \left[ \left\| \nabla f (\hat {x} ^ {T}) \right\| ^ {2} \right] = \mathcal {O} \left(\frac {\widetilde {L} \delta_ {0}}{\alpha T} + \frac {L \delta_ {0}}{\eta T} + \frac {\eta^ {3} \sigma^ {2}}{\alpha^ {2}} + \frac {\eta^ {2} \sigma^ {2}}{\alpha} + \frac {\eta \sigma^ {2}}{n} + \frac {\eta \sigma^ {2}}{\alpha^ {2} B _ {\mathrm {i n i t}} T}\right).
$$
Notice that 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_0\alpha}{\sigma^2T}\right)^{1/3}$ ,  $\eta \leq \left(\frac{L\delta_0n}{\sigma^2T}\right)^{1/2}$  and  $\eta \leq \frac{\alpha\sqrt{L\delta_0B_{\mathrm{init}}}}{\sigma}$  ensures that  $\frac{\eta^3\sigma^2}{\alpha^2} \leq \frac{L\delta_0}{\eta T}$ ,  $\frac{\eta^2\sigma^2}{\alpha} \leq \frac{L\delta_0}{\eta T}$ ,  $\frac{\eta\sigma^2}{n} \leq \frac{L\delta_0}{\eta T}$ , and
$\frac{\eta\sigma^2}{\alpha^2B_{\mathrm{init}}T}\leq \frac{L\delta_0}{\eta T}.$  Therefore, we have
$$
\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}{\sqrt {\alpha T}}\right) ^ {2 / 3} + \left(\frac {L \delta_ {0} \sigma^ {2}}{n T}\right) ^ {1 / 2} + \frac {\sigma \sqrt {L \delta_ {0}}}{\alpha \sqrt {B _ {\mathrm {i n i t}}} T}\right).
$$
Using  $B_{\mathrm{init}} \geq \min \left\{\frac{\sigma^2 L}{\widetilde{L}^2 \delta_0}, \frac{\sigma}{\alpha \sqrt{L \delta_0 T}}, \frac{\sigma^{2/3}}{\alpha^{4/3} T^{2/3} (L \delta_0)^{1/3}}, \frac{n}{\alpha^2 T}\right\}$ , we obtain
$$
\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}{\sqrt {\alpha} T}\right) ^ {2 / 3} + \left(\frac {L \delta_ {0} \sigma^ {2}}{n T}\right) ^ {1 / 2}\right).
$$
It remains to notice that  $\left\lceil \max \left\{\min \left\{\frac{\sigma^2L}{L^2\delta_0},\frac{\sigma}{\alpha\sqrt{L\delta_0T}},\frac{\sigma^{2 / 3}}{\alpha^{4 / 3}T^{2 / 3}(L\delta_0)^{1 / 3}},\frac{n}{\alpha^2T}\right\},\frac{\sigma^2}{L\delta_0n}\right\} \right\rceil \leq \left\lceil \frac{\sigma^2}{L\delta_0}\right\rceil$
F.1 Controlling the error of momentum estimator
Lemma 2. Let Assumption 1 be satisfied, and suppose  $0 < \eta \leq 1$ . For every  $i = 1, \dots, n$ , let the sequence  $\{v_i^t\}_{t \geq 0}$  be updated via
$$
v _ {i} ^ {t} = v _ {i} ^ {t - 1} + \eta \left(\nabla f _ {i} (x ^ {t}, \xi_ {i} ^ {t}) - v _ {i} ^ {t - 1}\right),
$$
Define the sequence  $v^t \coloneqq \frac{1}{n} \sum_{i=1}^{n} v_i^t$ . Then for every  $i = 1, \ldots, n$  and  $t \geq 0$  it holds
$$
\mathbb {E} \left[ \left\| v _ {i} ^ {t} - \nabla f _ {i} \left(x ^ {t}\right) \right\| ^ {2} \right] \leq (1 - \eta) \mathbb {E} \left[ \left\| v _ {i} ^ {t - 1} - \nabla f _ {i} \left(x ^ {t - 1}\right) \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}, \tag {24}
$$
$$
\mathbb {E} \left[ \left\| v ^ {t} - \nabla f (x ^ {t}) \right\| ^ {2} \right] \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}. \tag {25}
$$
Proof. By the update rule of  $v_{i}^{t}$ , we have