SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Compared to EF21-SGD, given by (5a) + (5ab), we replace the previous state  $g_{i}^{t}$  by the exact gradient at the current iteration. Since EF21-SGD heavily relies on the approximation  $g_{i}^{t} \approx \nabla f_{i}(x^{t+1})$ , and according to the proof of convergence of EF21-SGD, such discrepancy tends to zero as  $t \to \infty$ , this change can only improve the method. While we admit this is a conceptual algorithm only (it does not lead to any communication or sample complexity reduction in practice), it serves us well to illustrate the drawbacks of EF21-SGD. We now establish the following negative result for EF21-SGD-ideal.
Theorem 1. Let  $L, \sigma > 0, 0 < \gamma \leq 1 / L$  and  $n = 1$ . There exists a convex,  $L$ -smooth function  $f: \mathbb{R}^2 \to \mathbb{R}$ , a contractive compressor  $\mathcal{C}(\cdot)$  satisfying Definition 1, and an unbiased stochastic gradient with bounded variance  $\sigma^2$  such that if the method EF21-SGD-ideal ((5a) + (5aa)) is run with 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\{\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{\sigma^2}{60\varepsilon^2}$ , then we have  $\mathbb{E}\left[\|\nabla f(x^T)\|\right] > \varepsilon$  for all  $T \geq 0$ .
The above theorem implies that the method (5a), (5aa), does not converge with small batch-size (e.g., equal to one) for any fixed step-size choice. Moreover, in distributed setting with  $n$  nodes, a mini-batch of order  $B = \Omega \left( \frac{\sigma^2}{\varepsilon^2} \right)$  is required for convergence. Notice that this batch-size is independent of  $n$ , which further implies that a linear speedup in the number of nodes  $n$  cannot be achieved for this method. While we only prove these negative results for an "idealized" version of EF21-SGD rather than for the method itself, in Figures 1a and 4a, we empirically verify that EF21-SGD also suffers from a similar divergence on the same problem instance provided in the proof of Theorem 1. Additionally, Figures 1b and 4b illustrate that the situation does not improve for EF21-SGD when increasing  $n$ .
3.2 Momentum for avoiding mega-batches
Let us now focus on the single node setting<sup>10</sup> and try to fix the divergence issue shown above. As we can learn from Theorem 1, the key reason for non-convergence of EF21-SGD is that even if the state vector  $g^{t}$  sufficiently approximates the current gradient, i.e.,  $g^{t} \approx \nabla f(x^{t + 1})$ , the design of
This is because full gradients would need to be computed and communicated for its implementation. Notice also that if  $\sigma = 0$ , this method becomes the exact distributed gradient descent.
<sup>9</sup>In fact, the example can be easily extended to the case of polynomially decaying step-size.
In this case, we can drop index  $i$  everywhere and write  $g_{i}^{t} = g^{t}$ ,  $\xi_{i}^{t} = \xi^{t}$  for all  $t\geq 0$ .
this method cannot guarantee that the quantity  $\| g^{t} - \nabla f(x^{t})\|^{2}\approx \| \mathcal{C}\left(\nabla f(x^{t},\xi^{t}) - \nabla f(x^{t})\right)\|^{2}$  is small enough. Indeed, the last term above can be bounded by  $2(2 - \alpha)\sigma^2$  in expectation, but it is not sufficient as formally illustrated in Theorem 1. To fix this problem, we propose to modify our "idealized" EF21-SGD-ideal method so that the compressed difference can be controlled and made arbitrarily small, which leads us to another (more advanced) conceptual algorithm,
EF21-SGDM-ideal:
$$
v ^ {t + 1} = \nabla f (x ^ {t + 1}) + \eta (\nabla f (x ^ {t + 1}, \xi^ {t + 1}) - \nabla f (x ^ {t + 1})),
$$
$$
g ^ {t + 1} = \nabla f (x ^ {t + 1}) + \mathcal {C} \left(v ^ {t + 1} - \nabla f (x ^ {t + 1})\right).
$$
In this method, instead of using  $v^{t + 1} = \nabla f(x^{t + 1},\xi^{t + 1})$  as in EF21-SGD-ideal, we introduce a correction, which allows to control variance of the difference  $\nabla f(x^{t + 1},\xi^{t + 1}) - \nabla f(x^{t + 1})$ . This allows us to derive the following convergence result. Let  $\delta_0\coloneqq f(x^0) - f^*$
Proposition 1. Let Assumptions 1, 2 hold, and let  $\mathcal{C}$  satisfy Definition 1. Let  $g^0 = 0$  and the step-size in method (5a), (6) be set as  $\gamma \leq 1 / L$ . Let  $\hat{x}^T$  be sampled uniformly at random from the iterates of the method. Then for any  $\eta > 0$  after  $T$  iterations, we have  $\mathbb{E}\left[\| \nabla f(\hat{x}^T)\|^2\right] \leq \frac{2\delta_0}{\gamma T} + 4\eta^2\sigma^2$ .
Notice that if  $\eta = 1$ , then algorithm EF21-SGDM-ideal (5a), (6) reduces to EF21-SGD-ideal method (5a), (5aa), and this result shows that the lower bound for the batch-size established in Theorem 1 is tight, i.e.,  $B = \Theta (\sigma^2 /\varepsilon^2)$  is necessary and sufficient<sup>11</sup> for convergence. For  $\eta < 1$ , the above theorem suggests that using a small enough parameter  $\eta$ , the variance term can be completely eliminated. This observation motivates us to design a practical variant of this method. Similarly to the design of EF21 mechanism (from EF21-SGD-ideal), we propose to do this by replacing the exact gradients  $\nabla f(x^{t + 1})$  by state vectors  $v^{t}$  and  $g^{t}$  as follows:
EF21-SGDM:
$$
v ^ {t + 1} = v ^ {t} + \eta (\nabla f (x ^ {t + 1}, \xi^ {t + 1}) - v ^ {t}),
$$
$$
g ^ {t + 1} = g ^ {t} + \mathcal {C} (v ^ {t + 1} - g ^ {t})
$$
Theorem 2. Let Assumptions 1, 2 hold, and let  $\mathcal{C}$  satisfy Definition 1. Let method (5a), (7) be run with  $g^0 = v^0 = \nabla f(x^0)$ , and  $\hat{x}^T$  be sampled uniformly at random from the iterates of the method. Then for all  $\eta \in (0,1]$  with  $\gamma \leq \gamma_0 = \min \left\{\frac{\alpha}{20L},\frac{\eta}{7L}\right\}$ , we have  $\mathbb{E}\left[\left\|\nabla f(\hat{x}^T)\right\|^2\right] \leq \mathcal{O}\left(\frac{\delta_0}{\gamma T} + \eta \sigma^2\right)$ . The choice  $\eta = \min \left\{1,\left(\frac{L\delta_0}{\sigma^2T}\right)^{1/2}\right\}$ ,  $\gamma = \gamma_0$  results in  $\mathbb{E}\left[\left\|\nabla f(\hat{x}^T)\right\|^2\right] \leq \mathcal{O}\left(\frac{L\delta_0}{\alpha T} + \left(\frac{L\delta_0\sigma^2}{T}\right)^{1/2}\right)$ .
Compared to Proposition 1, where  $\eta$  can be made arbitrarily small, Theorem 2 suggests that there is a trade-off for the choice of  $\eta \in (0,1]$  in algorithm (5a), (7). The above theorem implies that in single node setting EF21-SGDM has  $\mathcal{O}\left(\frac{L}{\alpha\xi^2} +\frac{L\sigma^2}{\varepsilon^4}\right)$  sample complexity. For  $\alpha = 1$ , this result matches with the sample complexity of SGD and is known to be unimprovable under Assumptions 1 and 2 [Arjevani et al., 2019]. Moreover, when  $\alpha = 1$ , our sample complexity matches with previous analysis of momentum methods in [Liu et al., 2020] and [Defazio, 2021]. However, even in this single node  $(n = 1)$ , uncompressed  $(\alpha = 1)$  setting our analysis is different from the previous work, in particular, our choice of momentum parameter and the Lyapunov function are different, see Appendix A and J. For  $\alpha < 1$ , the above result matches with sample complexity of EF14-SGD (single node setting) [Stich and Karimireddy, 2019], which was recently shown to be optimal [Huang et al., 2022] for biased compressors satisfying Definition 1. However, notice that the extension of the analysis by Stich and Karimireddy [2019] for EF14-SGD to distributed setting meets additional challenges and it is unclear whether it is possible without imposing additional BG or BGS assumptions as in [Koloskova et al., 2020]. We revisit this analysis in Appendix K to showcase the difficulty of removing BG/BGS. In the following we will demonstrate the benefit of our EF21-SGDM method by extending it to distributed setting without imposing any additional assumptions.
3.3 Distributed stochastic error feedback with momentum
Now we are ready to present a distributed variant of EF21-SGDM, see Algorithm 1. Letting  $\delta_t \coloneqq f(x^t) - f^*$ , our convergence analysis of this method relies on the monotonicity of the following Lyapunov function:
$$
\Lambda_ {t} := \delta_ {t} + \frac {\gamma}{\alpha n} \sum_ {i = 1} ^ {n} \left\| g _ {i} ^ {t} - v _ {i} ^ {t} \right\| ^ {2} + \frac {\gamma \eta}{\alpha^ {2} n} \sum_ {i = 1} ^ {n} \left\| v _ {i} ^ {t} - \nabla f _ {i} \left(x ^ {t}\right) \right\| ^ {2} + \frac {\gamma}{\eta} \left\| \sum_ {i = 1} ^ {n} \left(v _ {i} ^ {t} - \nabla f _ {i} \left(x ^ {t}\right)\right) \right\| ^ {2}. \tag {8}
$$
This follows by replacing  $\sigma^2$  in the batch-free algorithm by  $\sigma^2 / B$  if the batch-size of size  $B > 1$  is used.
Algorithm 1 EF21-SGDM (Error Feedback 2021 Enhanced with Polyak Momentum)
1: Input: starting point  $x^0$ , step-size  $\gamma > 0$ , momentum  $\eta \in (0,1]$ , initial batch size  $B_{\mathrm{init}}$   
2: Initialize  $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 momentum estimator  $v_i^{t+1} = (1 - \eta)v_i^t + \eta \nabla f_i(x^{t+1}, \xi_i^{t+1})$   
7: Compress  $c_i^{t+1} = C(v_i^{t+1} - g_i^t)$  and send  $c_i^{t+1}$  to the master  
8: Update local state  $g_i^{t+1} = g_i^t + c_i^{t+1}$   
9: end for  
10: 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}$   
11: end for
- Convergence of EF21-SGDM with contractive compressors. We obtain the following result:
Theorem 3. Let Assumptions 1 and 2 hold. Let  $\hat{x}^T$  be sampled uniformly at random from the  $T$  iterates of the method. Let EF21-SGDM (Algorithm 1) be run with a contractive compressor. For all  $\eta \in (0,1]$  and  $B_{\mathrm{init}} \geq 1$ , with  $\gamma \leq \min \left\{\frac{\alpha}{20\widetilde{L}}, \frac{\eta}{7L}\right\}$ , we have
$$
\mathbb {E} \left[ \left\| \nabla f (\hat {x} ^ {T}) \right\| ^ {2} \right] \leq \mathcal {O} \left(\frac {\Lambda_ {0}}{\gamma T} + \frac {\eta^ {3} \sigma^ {2}}{\alpha^ {2}} + \frac {\eta^ {2} \sigma^ {2}}{\alpha} + \frac {\eta \sigma^ {2}}{n}\right), \tag {9}
$$
where  $\Lambda_0$  is given by (8). Choosing the batch size  $B_{\mathrm{init}} = \left\lceil \frac{\sigma^2}{L\delta_0}\right\rceil$ , and stepsize  $\gamma = \min \left\{\frac{\alpha}{20\overline{L}},\frac{\eta}{7L}\right\}$  and momentum  $\eta = \min \left\{1,\left(\frac{L\delta_0\alpha^2}{\sigma^2T}\right)^{1 / 4},\left(\frac{L\delta_0\alpha}{\sigma^2T}\right)^{1 / 3},\left(\frac{L\delta_0n}{\sigma^2T}\right)^{1 / 2},\frac{\alpha\sqrt{L\delta_0B_{\mathrm{init}}}}{\sigma}\right\}$ , we get
$$
\mathbb {E} \left[ \left\| \nabla f (\hat {x} ^ {T}) \right\| ^ {2} \right] \leq \mathcal {O} \left(\frac {\tilde {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).
$$
Remark 1. Note that using large initial batch size  $B_{init} > 1$  is not necessary for convergence of EF21-SGDM. If we set  $B_{init} = 1$ , the above theorem still holds by replacing  $\delta_0$  with  $\Lambda_0$ .
Remark 2. In the single node setting ( $n = 1$ ), the above result recovers the statement of Theorem 2 (with the same choice of parameters) since by Young's inequality  $\left(\frac{L\delta_0\sigma^{2/3}}{\alpha^{2/3}T}\right)^{3/4} \leq \frac{1}{2}\frac{L\delta_0}{\alpha T} + \frac{1}{2}\left(\frac{L\delta_0\sigma^2}{T}\right)^{1/2}, \left(\frac{L\delta_0\sigma}{\sqrt{\alpha}T}\right)^{2/3} \leq \frac{1}{3}\frac{L\delta_0}{\alpha T} + \frac{2}{3}\left(\frac{L\delta_0\sigma^2}{T}\right)^{1/2}$ , and  $\widetilde{L} = L$ .
- Recovering previous rates in case of full gradients. Compared to the iteration complexity  $\mathcal{O}\left(\frac{L_{\max}G}{\alpha\varepsilon^3}\right)$  of EF14 [Koloskova et al., 2020], our result, summarized in
Corollary 1. If  $\sigma = 0$ , then  $\mathbb{E}\left[\left\|\nabla f(\hat{x}^T)\right\|\right] \leq \varepsilon$  after  $T = \mathcal{O}\left(\frac{\widetilde{L}}{\alpha\varepsilon^2}\right)$  iterations.
is better by an order of magnitude, and does not require the BG assumption. The result of Corollary 1 is the same as for EF21 method [Richtárik et al., 2021], and EF21-HB method [Fatkhullin et al., 2021]. Notice, however, that even in this deterministic setting  $(\sigma = 0)$  EF21-SGDM method is different from EF21 and EF21-HB: while the original EF21 does not use momentum, EF21-HB method incorporates momentum on the server side to update  $x^{t}$ , which is different from our Algorithm 1, where momentum is applied by each node. This iteration complexity  $\mathcal{O}\left(\frac{1}{\alpha\varepsilon^2}\right)$  is optimal in both  $\alpha$  and  $\varepsilon$ . The matching lower bound was recently established by Huang et al. [2022] for smooth nonconvex optimization in the class of centralized, zero-respecting algorithms with contractive compressors.
- Comparison to previous work. Our sample complexity<sup>13</sup> in
Corollary 2.  $\mathbb{E}\left[\| \nabla f(\hat{x}^T)\| \right]\leq \varepsilon$  after  $T = \mathcal{O}\left(\frac{\widetilde{L}}{\alpha\varepsilon^{2}} +\frac{L\sigma^{2 / 3}}{\alpha^{2 / 3}\varepsilon^{8 / 3}} +\frac{L\sigma}{\alpha^{1 / 2}\varepsilon^{3}} +\frac{L\sigma^{2}}{n\varepsilon^{4}}\right)$  iterations.
$^{12}$ In Appendix J, we show how to deal with time varying  $\gamma_{t}$  and  $\eta_{t}$ .
<sup>13</sup>Note that the initial batch size contributes to the sample complexity only an additive constant independent of  $\varepsilon$ . Moreover,  $B_{\mathrm{init}} = \left\lceil \frac{\sigma^2}{L\delta_0} \right\rceil \leq \left\lceil \frac{2\sigma^2}{\varepsilon^2} \right\rceil$  since, otherwise,  $\left\| \nabla f(x^0) \right\|^2 \leq 2L\delta_0 \leq \varepsilon^2$ , and  $x^0$  is a solution. In the following, we ignore the dependence on  $B_{\mathrm{init}}$  for a fair comparison with other works.
strictly improves over the complexity  $\mathcal{O}\left(\frac{GL_{\max}}{\alpha\varepsilon^3} +\frac{L_{\max}\sigma^2}{n\varepsilon^4}\right)$  of EF14-SGD by Koloskova et al. [2020], even in case when  $G < + \infty$  . Notice that it always holds that  $\sigma \leq G$  . If we assume that  $G\approx \sigma$  our three first terms in the complexity improve the first term from Koloskova et al. [2020] by the factor of  $\varepsilon /\sigma$ $(\varepsilon \alpha /\sigma)^{1 / 3}$  , or  $\alpha^{1 / 2}$  . Compared to the BEER algorithm of Zhao et al. [2022], with sample complexity  $\mathcal{O}\left(\frac{L_{\max}}{\alpha\varepsilon} +\frac{L_{\max}\sigma^2}{\alpha^2\varepsilon^4}\right)$  , the result of Corollary 2 is strictly better in terms of  $\alpha ,n,$  and the smoothness constants. In addition, we remove the large batch requirement for convergence compared to [Fatkhullin et al., 2021, Zhao et al., 2022]. Moreover, notice that Corollary 2 implies that EF21-SGDM achieves asymptotically optimal sample complexity  $\mathcal{O}\left(\frac{L\sigma^2}{n\varepsilon^4}\right)$  in the regime  $\varepsilon \rightarrow 0$
3.4 Further improvement using double momentum!
Unfortunately, in the non-asymptotic regime, our sample complexity does not match with the lower bound in all problem parameters simultaneously due to the middle term  $\frac{L\sigma^{2/3}}{\alpha^{2/3}\varepsilon^{8/3}} + \frac{L\sigma}{\alpha^{1/2}\varepsilon^3}$ , which can potentially dominate over  $\frac{L\sigma^2}{n\varepsilon^4}$  term for large enough  $n$  and  $\varepsilon$ , and small enough  $\alpha$  and  $\sigma$ . We propose a double-momentum method, which can further improve the middle term in the sample complexity of EF21-SGDM. We replace the momentum estimator  $v_i^t$  in line 6 of Algorithm 1 by the following two-step momentum update
$$
\text {E F 2 1 - S G D 2 M :} v _ {i} ^ {t + 1} = (1 - \eta) v _ {i} ^ {t} + \eta \nabla f _ {i} \left(x ^ {t + 1}, \xi_ {i} ^ {t + 1}\right), u _ {i} ^ {t + 1} = (1 - \eta) u _ {i} ^ {t} + \eta v _ {i} ^ {t + 1}. \tag {10}
$$
We formally present this method in Algorithm 3 in Appendix G. Compared to EF21-SGDM (Algorithm 1), the only change is that instead of compressing  $v_{i}^{t} - g_{i}^{t}$ , in EF21-SGD2M, we compress  $u_{i}^{t} - g_{i}^{t}$ , where  $u_{i}$  is a two step (double) momentum estimator. The intuition behind this modification is that a double momentum estimator  $u_{i}^{t}$  has richer "memory" of the past gradients compared to  $v_{i}^{t}$ . Notice that for each node, EF21-SGD2M requires to save 3 vectors  $(v_{i}^{t}, u_{i}^{t}, g_{i}^{t})$  instead of 2 in EF21-SGDM  $(v_{i}^{t}, g_{i}^{t})$  and EF14-SGD  $(e_{i}^{t}, g_{i}^{t})$ .<sup>15</sup> When interacting with biased compression operator  $\mathcal{C}(\cdot)$ , such effect becomes crucial in improving the sample complexity. For EF21-SGD2M, we derive
Corollary 3. Let  $v_{i}^{t}$  in Algorithm 1 be replaced by  $u_{i}^{t}$  given by (10) (Algorithm 3 in Appendix G). Then with appropriate choice of  $\gamma$  and  $\eta$  (given in Theorem 5), we have  $\mathbb{E}\left[\left\| \nabla f(\hat{x}^T)\right\|\right] \leq \varepsilon$  after  $T = \mathcal{O}\left(\frac{\widetilde{L}\delta_0}{\alpha\varepsilon^2} +\frac{L\delta_0\sigma^{2 / 3}}{\alpha^{2 / 3}\varepsilon^{8 / 3}} +\frac{L\delta_0\sigma^2}{n\varepsilon^4}\right)$  iterations.
We consider a nonconvex logistic regression problem:  $f_{i}(x_{1},\ldots ,x_{c}) = -\frac{1}{m}\sum_{j = 1}^{m}\log (\exp (a_{ij}^{\top}x_{y_{ij}}) / \sum_{y = 1}^{c}\exp (a_{ij}^{\top}x_{y}))$  with a nonconvex regularizer  $h(x_1,\dots,x_c) = \lambda \sum_{y = 1}^{c}\sum_{k = 1}^{l}[x_y]_k^2 /(1 + [x_y]_k^2)$  with  $\lambda = 10^{-3}$ , where  $x_{1},\ldots ,x_{c}\in \mathbb{R}^{l}$ ,  $[\cdot ]_k$  is an indexing operation of a vector,  $c\geq 2$  is the number of classes,  $l$  is the number of features,  $m$  is the size of a dataset,  $a_{ij}\in \mathbb{R}^l$  and  $y_{ij}\in \{1,\dots,c\}$  are features and labels. The datasets used are MNIST (with  $l = 784$ ,  $m = 60000$ ,  $c = 10$ ) and real-sim (with  $l = 20958$ ,  $m = 72309$ ,  $c = 2$ ) [LeCun et al., 2010, Chang and Lin, 2011]. The dimension of the problem is  $d = (l + 1)c$ , i.e.,  $d = 7850$  for MNIST and  $d = 41918$  for real-sim. In each experiment, we show relations between the total number of transmitted coordinates and gradient/function values. The stochastic gradients in each algorithm are replaced by a mini-batch estimator  $\frac{1}{B}\sum_{j = 1}^{B}\nabla f_i(x,\xi_{ij})$  with the same  $B\geq 1$  in each plot. Notice that all methods (except for NEOLITHIC) $^{16}$  calculate the same number of samples at each communication round, thus the dependence on the number of samples used will be qualitatively the same. In all algorithms, the step sizes are fine-tuned from a set  $\{2^k\mid k\in [-20,20]\}$  and the TopK compressor is used to compress information from the nodes to the master. For EF21-SGDM, we fix momentum parameter  $\eta = 0.1$  in all experiments. Prior to that, we tuned  $\eta \in \{0.01,0.1\}$  on the independent dataset w8a (with  $l = 300$ ,  $m = 49749$ ,  $c = 2$ ). We omit BEER method from the plots since it showed worse performance than EF21-SGD in all runs.
$^{14}L_{\mathrm{max}}\coloneqq \max_{i\in [n]}L_i$  . Notice that  $L\leq \widetilde{L}\leq L_{\mathrm{max}}$  and the inequalities are strict in heterogeneous setting.
<sup>15</sup>See Appendix K for details on EF14-SGD. In contrast, EF21-SGD needs to save only one vector  $(g_i^t)$ .
For NEOLITHIC, we use the parameter  $R = \lceil d / K \rceil$  following the requirement in their Theorem 3. Experiments in Huang et al. [2022] use a heuristic choice  $R = 4$ , and thus can show faster convergence.
4.1 Experiment 1: increasing batch-size
In this experiment, we use MNIST dataset and fix the number of transmitted coordinates to  $K = 10$  (thus  $\alpha \geq K / d \approx 10^{-3}$ ), and set  $n = 10$ . Figure 2 shows convergence plots for  $B \in \{1,32,128\}$ . EF21-SGDM and its double momentum version EF21-SGD2M have fast convergence and show a significant improvement when increasing batch-size compared to EF14-SGD. In contrast, EF21-SGD suffers from poor performance for small  $B$ , which confirms our observations in previous sections. NEOLITHIC has order times slower convergence rate due to the fact that it sends  $\lceil d / K \rceil$  compressed vectors in each iteration, while other methods send only one.
(a)  $B = 1$
(b)  $B = 32$
Figure 2: Experiment on MNIST dataset with  $n = 10$ , and Top10 compressor.
(c)  $B = 128$
4.2 Experiment 2: improving convergence with  $n$
This experiment uses real-sim dataset,  $K = 100$  (thus  $\alpha \geq K / d \approx 2 \cdot 10^{-3}$ ), and with  $B = 128 \ll m$ . We vary the number of nodes within  $n \in \{1, 10, 100\}$ , see Figure 3. In this case, EF21-SGDM and EF21-SGD2M have much faster convergence compared to other methods for all  $n$ . Moreover, the proposed algorithms show a significant improvement when  $n$  increases. We also observe that on this task, EF21-SGD2M performs slightly worse than EF21-SGDM for  $n = 10, 100$ , but it is still much faster than other other methods.
(a)  $n = 1$
Figure 3: Experiment on real-sim dataset with batch-size  $B = 128$ , and Top100 compressor.
(b)  $n = 10$
In Section C, we present extra simulations with different parameters for above experiments. Additionally, we include experiemnts on simple quadratic problems and perform training of larger image recognition models. In all cases, EF21-SGDM and EF21-SGD2M outperform other algorithms.
This work of I. Fatkhullin was supported by ETH AI Center doctoral fellowship. The work of P. Richtárik and A. Tyurin was supported by the KAUST Baseline Research Scheme (KAUST BRF) and the KAUST Extreme Computing Research Center (KAUST ECRC); P. Richtárik was also supported by the SDAIA-KAUST Center of Excellence in Data Science and Artificial Intelligence (SDAIA-KAUST AI).