SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Ilyas Fatkhullin  
ETH AI Center & ETH Zurich
Alexander Tyurin  
KAUST*
Peter Rictárik  
KAUST
Due to the high communication overhead when training machine learning models in a distributed environment, modern algorithms invariably rely on lossy communication compression. However, when untreated, the errors caused by compression propagate, and can lead to severely unstable behavior, including exponential divergence. Almost a decade ago, Seide et al. [2014] proposed an error feedback (EF) mechanism, which we refer to as EF14, as an immensely effective heuristic for mitigating this issue. However, despite steady algorithmic and theoretical advances in the EF field in the last decade, our understanding is far from complete. In this work we address one of the most pressing issues. In particular, in the canonical nonconvex setting, all known variants of EF rely on very large batch sizes to converge, which can be prohibitive in practice. We propose a surprisingly simple fix which removes this issue both theoretically, and in practice: the application of Polyak's momentum to the latest incarnation of EF due to Richtárik et al. [2021] known as EF21. Our algorithm, for which we coin the name EF21-SGDM, improves the communication and sample complexities of previous error feedback algorithms under standard smoothness and bounded variance assumptions, and does not require any further strong assumptions such as bounded gradient dissimilarity. Moreover, we propose a double momentum version of our method that improves the complexities even further. Our proof seems to be novel even when compression is removed from the method, and as such, our proof technique is of independent interest in the study of nonconvex stochastic optimization enriched with Polyak's momentum.
Since the practical utility of modern machine learning models crucially depends on our ability to train them on large quantities of training data, it is imperative to perform the training in a distributed storage and compute environment. In federated learning (FL) [Konečný et al., 2016, Kairouz, 2019], for example, data is naturally stored in a distributed fashion across a large number of clients (who capture and own the data in the first place), and the goal is to train a single machine learning model from the wealth of all this distributed data, in a private fashion, directly on their devices.
1.1 Formalism. We consider the problem of collaborative training of a single model by several clients in a data-parallel fashion. In particular, we aim to solve the distributed nonconvex stochastic optimization problem
$$
\min  _ {x \in \mathbb {R} ^ {d}} \left[ f (x) := \frac {1}{n} \sum_ {i = 1} ^ {n} f _ {i} (x) \right], \quad f _ {i} (x) := \mathbb {E} _ {\xi_ {i} \sim \mathcal {D} _ {i}} \left[ f _ {i} (x, \xi_ {i}) \right], \quad i = 1, \dots , n, \tag {1}
$$
where  $n$  is the number of clients,  $x \in \mathbb{R}^d$  represents the parameters of the model we wish to train, and  $f_{i}(x)$  is the (typically nonconvex) loss of model parameterized by the vector  $x$  on the data  $\mathcal{D}_i$  owned by client  $i$ . Unlike most works in federated learning, we do not assume the datasets to be similar, i.e., we allow the distributions  $\mathcal{D}_1, \ldots, \mathcal{D}_n$  to be arbitrarily different.
*King Abdullah University of Science and Technology, Thuwal, Saudi Arabia.
We are interested in the fundamental problem of finding an approximately stationary point of  $f$  in expectation, i.e., we wish to find a (possibly random) vector  $\hat{x} \in \mathbb{R}^d$  such that  $\mathbb{E}\left[\|\nabla f(\hat{x})\|\right] \leq \varepsilon$ . In order to solve this problem, we assume that the  $n$  clients communicate via an orchestrating server. Typically, the role of the server is to first perform aggregation of the messages obtained from the workers, and to subsequently broadcast the aggregated information back to the workers. Following an implicit assumption made in virtually all theoretically-focused papers on communication-efficient training, we also assume that the speed of client-to-workers broadcast is so fast (compared to speed of workers-to-client communication) that the cost associated with broadcast can be neglected<sup>2</sup>.
1.2 Aiming for communication and computation efficiency at the same time. In our work, we pay attention to two key aspects of efficient distributed training—communication cost and computation cost for finding an approximate stationary point  $\hat{x}$ . The former refers to the number of bits that need to be communicated by the workers to the server, and the latter refers to the number of stochastic gradients that need to be sampled by each client. The rest of the paper can be summarized as follows: We pick one of the most popular communication-efficient gradient-type methods (the EF21 method of Richtárik et al. [2021] – the latest variant of error feedback pioneered by Seide et al. [2014]) and modify it in a way which provably preserves its communication complexity, but massively improves its computation/sample complexity, both theoretically and in practice.
2 Communication Compression, Error Feedback, and Sample Complexity
Communication compression techniques such as quantization [Alistarh et al., 2017, Horváth et al., 2019a] and sparsification [Seide et al., 2014, Beznosikov et al., 2020] are known to be immensely powerful for reducing the communication footprint of gradient-type<sup>3</sup> methods. Arguably the most studied, versatile and practically useful class of compression mappings are contractive compressors.
Definition 1 (Contractive compressors). We say that a (possibly randomized) mapping  $\mathcal{C}:\mathbb{R}^d\to \mathbb{R}^d$  is a contractive compression operator if there exists a constant  $0 < \alpha \leq 1$  such that
Inequality (2) is satisfied by a vast array of compressors considered in the literature, including numerous variants of sparsification operators [Alistarh et al., 2018, Stich et al., 2018], quantization operators [Alistarh et al., 2017, Horváth et al., 2019a], and low-rank approximation [Vogels et al., 2019, Safaryan et al., 2022] and more [Beznosikov et al., 2020, Safaryan et al., 2021]. The canonical examples are i) the Top  $K$  sparsifier, which preserves the  $K$  largest components of  $x$  in magnitude and sets all remaining coordinates to zero [Stich et al., 2018], and ii) the (scaled) Rand  $K$  sparsifier, which preserves a subset of  $K$  components of  $x$  chosen uniformly at random and sets all remaining coordinates to zero [Beznosikov et al., 2020]. In both cases, (2) is satisfied with  $\alpha = K / d$ .
2.1 Brief history of error-feedback
When greedy contractive compressors, such as Top  $K$ , are used in a direct way to compress the local gradients in distributed gradient descent (GD), the resulting method may diverge exponentially, even on strongly convex quadratics [Beznosikov et al., 2020]. Empirically, instability caused by such a naive application of greedy compressors was observed much earlier, and a fix was proposed in the form of the error feedback (EF) mechanism by Seide et al. [2014], which we henceforth call EF14 or EF14-SGD (in the stochastic case). To the best of our knowledge, the best sample complexity of EF14-SGD for finding a stationary point in the distributed nonconvex setting is given by Koloskova et al. [2020]: after  $\mathcal{O}(G\alpha^{-1}\varepsilon^{-3} + \sigma^2 n^{-1}\varepsilon^{-4})$  samples, EF14-SGD finds a point  $x$  such that  $\mathbb{E}[\| \nabla f(x)\| ]\leq \varepsilon$ , where  $\alpha$  is the contraction parameter (see Definition 1). However, such
2While this is a reasonable assumption in many practical situations [Mishchenko et al., 2019, Kairouz, 2019], some works consider the regime when the server-to-workers broadcast cannot be neglected [Horváth et al., 2019a, Tang et al., 2020, Philippenko and Dieuleveut, 2020, Kovalev et al., 2021, Fatkhullin et al., 2021, Gruntkowska et al., 2022].
3For Newton-type methods, see [Islamov et al., 2022] and references therein.
4 In Appendix A, we provide a more detailed discussion on theoretical developments for this method.
Here  $\sigma^2$  is the bound on the variance of stochastic gradients at each node, see Assumption 2. When referring the sample complexity we count the number of stochastic gradients used only at one of the  $n$  nodes rather than by all nodes in total. This is a meaningful notion because the computations are done in parallel.
(a) Divergence for  $n = 1$ .
(b) No improvement with  $n$
Figure 1: Divergence of EF21-SGD on the quadratic function  $f(x) = \frac{1}{2} \| x\|^2$ ,  $x \in \mathbb{R}^2$ , using the Top1 compressor. See the proof of Theorem 1 for details on the construction of the noise  $\xi$ ; we use  $\sigma = 1$ ,  $B = 1$ . The starting point is  $x^0 = (0, -0.01)^{\top}$ . Unlike EF21-SGD, our method EF21-SGDM does not suffer from divergence and is stable near optimum. Figure 1b shows that when increasing the number of nodes  $n$ , EF21-SGD applied with  $B = 1$  does not improve, and, moreover, diverges from the optimum even faster. All experiments use constant parameters  $\gamma = \eta = 0.1 / \sqrt{T} = 10^{-3}$ ; see Figure 4 for diminishing parameters. Each method is run 10 times and the plot shows the median performance alongside the 25% and 75% quantiles.
an analysis has two important deficiencies. First, in the deterministic case (when exact gradients are computable by each node), the analysis only gives the suboptimal  $\mathcal{O}(\varepsilon^{-3})$  iteration complexity, which is suboptimal compared to vanilla (i.e., non-compressed) gradient descent, whose iteration complexity is  $\mathcal{O}(\varepsilon^{-2})$ . Second, their analysis relies heavily on additional strong assumptions, such as the bounded gradient (BG) assumption,  $\mathbb{E}[\| \nabla f_i(x,\xi_i)\| ^2 ]\leq G^2$  for all  $x\in \mathbb{R}^d$ ,  $i\in [n]$ ,  $\xi_{i}\sim \mathcal{D}_{i}$ , or the bounded gradient similarity (BGS) assumption,  $\frac{1}{n}\sum_{i = 1}^{n}\| \nabla f_i(x) - \nabla f(x)\| ^2\leq G^2$  for all  $x\in \mathbb{R}^d$ . Such assumptions are restrictive and sometimes even unrealistic. In particular, both BG and BGS might not hold even in the case of convex quadratic functions. Moreover, it was recently shown that nonconvex analysis of stochastic gradient methods using a BG assumption may hide an exponential dependence on the smoothness constant in the complexity [Yang et al., 2023].
In 2021, these issues were partially resolved by Richtárik et al. [2021], who propose a modification of the EF mechanism, which they call EF21. They address both deficiencies of the original EF14 method: i) they removed the BG/BGS assumptions, and improved the iteration complexity to  $\mathcal{O}(\varepsilon^{-2})$  in the full gradient regime. Subsequently, the EF21 method was modified in several directions, e.g., extended to bidirectional compression, variance reduction and proximal setup [Fatkhullin et al., 2021], generalized from contractive to three-point compressors [Richtárik et al., 2022] and adaptive compressors [Makarenko et al., 2022], modified from dual (gradient) to primal (model) compression [Gruntkowska et al., 2022] and from centralized to decentralized setting [Zhao et al., 2022]. For further work, we refer to [Wang et al., 2022, Dorfman et al., 2023, Islamov et al., 2022].
2.2 Key issue: error feedback has an unhealthy appetite for samples!
Unfortunately, the current theory of EF21 with stochastic gradients has weak sample complexity guarantees. In particular, Fatkhullin et al. [2021] extended the EF21-GD method, which is the basic variant of EF21 using full gradient at the clients, to EF21-SGD, which uses a "large minibatch" of stochastic gradients instead. They obtained  $\mathcal{O}\left(\frac{1}{\alpha\varepsilon^2} + \frac{\sigma^2}{\alpha^3\varepsilon^4}\right)$  sample complexity for their method. Later, Zhao et al. [2022] improved this result slightly<sup>7</sup> to  $\mathcal{O}\left(\frac{1}{\alpha\varepsilon^2} + \frac{\sigma^2}{\alpha^2\varepsilon^4}\right)$ , shaving off one  $\alpha$  in the stochastic term. However, it is easy to notice several issues in these results, which generally feature the fundamental challenge of combining biased gradient methods with stochastic gradients.
- Mega-batches. These works require all clients to sample "mega-batches" of stochastic gradients/datatpoints in each iteration, of order  $\mathcal{O}(\varepsilon^{-2})$ , in order to control the variance coming from stochastic gradients. In Figure 1, we find that, in fact, a batch-free (i.e., with mini-batch size  $B = 1$ )
For example, one can consider  $f_{i}(x) = x^{\top}A_{i}x$  with  $A_{i}\in \mathbb{R}^{d\times d}$ , for which BG or BGS assumptions hold only in the trivial cases: matrices  $A_{i}$  are all zero or all equal to each other (homogeneous data regime).
The result was obtained under a more general setting of decentralized optimization over a network.
version of EF21-SGD diverges even on a very simple quadratic function. We also observe a similar behavior when a small batch  $B > 1$  is applied. This implies that there is a fundamental flaw in the EF21-SGD method itself, rather "just" a problem of the theoretical analysis. While mega-batch methods are common in optimization literature, smaller batches are often preferred whenever they "work". For example, the time/cost required to obtain such a large number of samples at each iteration might be unreasonably large compared to the communication time, which is already reduced using compression. Moreover, when dealing with medical data, large batches might simply be unavailable [Rieke et al., 2020]. In certain applications, such as federated reinforcement learning (RL) or multi-agent RL, it is often intractable to sample more than one trajectory of the environment in order to form a gradient estimator [Mitra et al., 2023, Doan et al., 2019, Jin et al., 2022, Khodadadian et al., 2022]. Further, a method using a mega-batch at each iteration effectively follows the gradient descent (GD) dynamics instead of the dynamics of (mini-batch) SGD, which may hinder the training and generalization performance of such algorithms since it is both empirically [Keskar et al., 2017, Kleinberg et al., 2018] and theoretically [Kale et al., 2021] observed that mini-batch SGD is superior to mega-batch SGD or GD in a number of machine learning tasks.
- Dependence on  $\alpha$ . The total sample complexity results derived by Fatkhullin et al. [2021], Zhao et al. [2022] suffer from poor dependence on the contraction parameter  $\alpha$ . Typically, EF methods are used with the Top  $K$  sparsifier, which only communicates  $K$  largest entries in magnitude. In this case,  $\alpha = K / d$ , and the stochastic part of sample complexity scales quadratically with dimension.
- No improvement with  $n$ . The stochastic term in the sample complexity of EF21-SGD does not improve when increasing the number of nodes. However, the opposite behavior is typically desired, and is present in several latest non-EF methods based on unbiased compressors, such as MARINA [Gorbunov et al., 2021] and DASHA [Tyurin and Richtárik, 2022]. We are not aware of any distributed algorithms utilizing the TopK compressor achieving linear speedup in  $n$  in the stochastic term without relying on restrictive BG or BGS assumptions.
These observations motivate our work with the following central questions:
Can we design a batch-free distributed SGD method utilizing contractive communication compression (such as TopK) without relying on restrictive BG/BGS assumptions? Is it possible to improve over the current state-of-the-art  $\mathcal{O}\left(\alpha^{-1}\varepsilon^{-2} + \sigma^2\alpha^{-2}\varepsilon^{-4}\right)$  sample complexity under the standard smoothness and bounded variance assumptions?
We answer both questions in the affirmative by incorporating a momentum update into EF21-SGD.
2.3 Mysterious effectiveness of momentum in nonconvex optimization
An immensely popular modification of SGD (and its distributed variants) is the use of momentum. This technique, initially inspired by the developments in convex optimization [Polyak, 1964], is often applied in machine learning for stabilizing convergence and speeding up the training. In particular, momentum is an important part of an immensely popular and empirically successful line of adaptive methods for deep learning, including ADAM [Kingma and Ba, 2015] and a plethora of variants. The classical SGD method with Polyak (i.e., heavy ball) momentum (SGDM) reads:
$$
x ^ {t + 1} = x ^ {t} - \gamma v ^ {t}, \quad v ^ {t + 1} = (1 - \eta) v ^ {t} + \eta \nabla f \left(x ^ {t + 1}, \xi^ {t + 1}\right), \tag {3}
$$
where  $\gamma > 0$  is a learning rate and  $\eta > 0$  is the momentum parameter.
We provide a concise walk through the key theoretical developments in the analysis of SGDM in stochastic nonconvex optimization in Appendix A; and only mention the most relevant works here. The most closely related works to ours are [Mishchenko et al., 2019], [Xie et al., 2020], and [Fatkhullin et al., 2021], which analyze momentum together with communication compression. The analysis in [Mishchenko et al., 2019, Xie et al., 2020] requires BG/BGS assumption, and does not provide any theoretical improvement over the variants without momentum. Finally, the analysis of Fatkhullin et al. [2021] is only established for deterministic case, and it is unclear if its extension to stochastic case can bring any convergence improvement over EF21-SGD. Recently, several other works attempt to explain the benefit of momentum [Plattner, 2022]; some consider structured nonconvex problems [Wang and Abernethy, 2021], and others focus on generalization [Jelassi and Li, 2022].
Method	Communication complexity	Asymptotic sample complexity	Batch-free?	No extra assumptions?
EF14-SGD [Koloskova et al., 2020]	KG/αε3	σ2/nε4	✓	X (a)
NEOLITHIC [Huang et al., 2022]	K/αε2 log (G/ε) (b)	σ2/nε4	X	X (c)
EF21-SGD [Fatkhullin et al., 2021]	K/αε2	σ2/α3ε4 (d)	X	✓
BEER [Zhao et al., 2022]	K/αε2	σ2/α2ε4 (d)	X	✓
EF21-SGDM Corollary 2	K/αε2	σ2/nε4	✓	✓
(a) Analysis requires a bound of the second moment of the stochastic gradients, i.e.,  $\mathbb{E}\left[\| \nabla f_i(x,\xi_i)\|^2\right]\leq G^2$  for all  $x\in \mathbb{R}^d$
(b) This complexity is achieved by using a large mini-batch and communicating  $\lceil K / \alpha \rceil$  coordinates per iteration, see Appendix A.
(c) Analysis requires a bounded gradient disimilarity assumption, i.e.,  $\frac{1}{n}\sum_{i = 1}^{n}\| \nabla f_i(x) - \nabla f(x)\| ^2\leq G^2$  for all  $x\in \mathbb{R}^d$
(d) Analysis requires a batch-size at least  $B \geq \frac{\sigma^2}{\alpha^2\varepsilon^2}$  for EF21-SGD and  $B \geq \frac{\sigma^2}{\alpha\varepsilon^2}$  for BEER.
Table 1: Summary of related works on distributed error compensated SGD methods using a Top  $K$  compressor under Assumptions 1 and 2. The goal is to find an  $\varepsilon$ -stationary point of a smooth nonconvex function of the form (1), i.e., a point  $x$  such that  $\mathbb{E}\left[\|\nabla f(x)\|\right] \leq \varepsilon$ . "Communication complexity": the total # of communicated bits if the method is applied with sufficiently large batch-size; see Table 2 for batch-size. "Asymptotic sample complexity": the total # of samples required at each node to find an  $\varepsilon$ -stationary point for batch-size  $B = 1$  in the regime  $\varepsilon \rightarrow 0$ . "No extra assumptions":  $\checkmark$  means that no additional assumption is required.
Summary of contributions. Despite the vast amount of work trying to explain the benefits of momentum, there is no work obtaining any theoretical improvement over vanilla SGD in the smooth nonconvex setting under the standard assumptions of smoothness and bounded variance.
- First, we establish a negative result for a simplified/idealized version of EF21-SGD, which shows that this algorithm does not converge with constant batch-size, and that a mega-batch of order  $\Omega (\sigma^2\varepsilon^{-2})$  is required. This provides a strong indication that EF21-SGD method is inherently sensitive to stochastic gradients, which is also confirmed by our numerical simulations.
- We propose a simple fix for this problem by incorporating momentum step into EF21-SGD, which leads to our one-batch EF21-SGDM algorithm. By leveraging our new Lyapunov function construction and new analysis, we establish  $\mathcal{O}\left(\alpha^{-1}\varepsilon^{-2} + \sigma^2\varepsilon^{-4}\right)$  sample complexity in the single node case.
- We extend our algorithm to the distributed setting and derive an improved sample complexity result compared to other methods using the Top  $K$  compressor without resorting to the BG/BGS assumptions. In particular, EF21-SGDM achieves asymptotically optimal  $\mathcal{O}\left(\sigma^{2}n^{-1}\varepsilon^{-4}\right)$  sample complexity. Moreover, when EF21-SGDM is applied with large enough batch size, we prove that it reaches the optimal communication complexity  $\mathcal{O}\left(K\alpha^{-1}\varepsilon^{-2}\right)$ ; see Tables 1 & 2 for more details.
- Finally, we propose a double momentum variant of EF21-SGDM, and find that it further improves the sample complexity of EF21-SGDM in the non-asymptotic regime.
We highlight that, interestingly, we prove that momentum helps: EF21-SGDM is theoretically better compared to its non-momentum variant – large-batch EF21-SGD. We believe that our new technique can be extended in many ways, e.g., to dealing with other biased compressors and other (biased) communication saving techniques such as lazy aggregation of gradients, model compression, bidirectional compression, partial participation, decentralized training, adaptive compression; other important optimization techniques involving biased updates such as proximal SGD with momentum, gradient clipping and adaptive step-size schedules. We also hope that our proof techniques can be useful to establish linear speedup for other classes of distributed methods, e.g., algorithms based on local training ProxSkip/Scaffnew [Mishchenko et al., 2022].
Additionally, we extend our results to the class of absolute compressors in Appendix H and study the variance reduced variant of error feedback in Appendix I.
Throughout the paper we work under the following standard assumptions.
Assumption 1 (Smoothness and lower boundedness). We assume that  $f$  has  $L$ -Lipschitz gradient, i.e.,  $\| \nabla f(x) - \nabla f(y) \| \leq L \| x - y \|$  for all  $x, y \in \mathbb{R}^d$ , and each  $f_i$  has  $L_i$ -Lipschitz gradient,
i.e.,  $\| \nabla f_i(x) - \nabla f_i(y)\| \leq L_i\| x - y\|$  for all  $i\in [n],x,y\in \mathbb{R}^d$  . We denote  $\widetilde{L}^2\coloneqq \frac{1}{n}\sum_{i = 1}^{n}L_i^2$  Moreover, we assume that  $f$  is lower bounded, i.e.,  $f^{*}\coloneqq \inf_{x\in \mathbb{R}^{d}}f(x) > - \infty$
Assumption 2 (Bounded variance (BV)). There exists  $\sigma >0$  such that
$$
\mathbb {E} \left[ \| \nabla f _ {i} (x, \xi_ {i}) - \nabla f _ {i} (x) \| ^ {2} \right] \leq \sigma^ {2}, \quad \forall x \in \mathbb {R} ^ {d}, \quad \forall i \in [ n ], \tag {4}
$$
where  $\xi_{i}\sim \mathcal{D}_{i}$  are i.i.d. random samples for each  $i\in [n]$ .
3.1 A deeper dive into the issues EF21 has with stochastic gradients
As remarked before, the current analysis of EF21 in the stochastic setting requires each client to sample a mega-batch in each iteration, and it is not clear how to avoid this. In order to understand this phenomenon, we propose to step back and examine an "idealized" version of EF21-SGD, which we call EF21-SGD-ideal, defined by the update rules (5a) + (5aa):
$$
x ^ {t + 1} = x ^ {t} - \gamma g ^ {t}, \quad g ^ {t} = \frac {1}{n} \sum_ {i = 1} ^ {n} g _ {i} ^ {t} \tag {5a}
$$
EF21-SGD-ideal:  $g_{i}^{t + 1} = \nabla f_{i}(x^{t + 1})$ $+C\left(\nabla f_i(x^{t + 1},\xi_i^{t + 1}) - \nabla f_i(x^{t + 1})\right),$  (5aa)
EF21-SGD:  $g_{i}^{t + 1} = g_{i}^{t}$ $+C\left(\nabla f_i(x^{t + 1},\xi_i^{t + 1}) - g_i^t\right).$  (5ab)