Year: 2021

Abstract:
On-device memory concerns in distributed deep learning are becoming more severe due to i) the growth of model size in multi-GPU training, and ii) the adoption of neural networks for federated learning on IoT devices with limited storage. In such settings, this work deals with memory issues emerging with communication efficient methods. To tackle associated challenges, key advances are that i) instead of EFSGD that inefficiently manages memory, the sweet spot of convergence and memory usage can be attained via what is here termed contractive error feedback (ConEF); and, ii) communication efficiency in ConEF should be achieved by biased and allreducable gradient compression. ConEF is validated on various learning tasks that include image classification, language modeling, and machine translation. ConEF saves 80% – 90% of the extra memory in EFSGD with almost no loss on test performance, while also achieving 1.3x – 5x speedup of SGD.

heading: 1 INTRODUCTION
text: Many modern machine learning tasks can be cast as an optimization problem
min x∈Rd
f(x) := 1
N N∑ i=1 Eξ∼D [ fi(x, ξ) ] . (1)
Typically, the possibly nonconvex loss function fi(x, ξ) is continuously differentiable with respect to (w.r.t.) x. We assume that f∗ > −∞ is a global minimum of (1). In this work, we will focus on training neural networks (NN) by solving (1) in two important scenarios: i) data parallelism in a multi-GPU setting, and ii) federated learning where clients could be edge devices such as smart phones. In both cases, N denotes the number of workers (GPUs or devices) and f∗ is no smaller than 0 given common choices of loss functions such as cross entropy and mean square.
While training NNs is a resource-hungry task, bandwidth and memory prevent us from fully harvesting the merits of parallelism. Limited bandwidth slows down the speed of information (e.g., gradients) exchange among workers. In multi-GPU settings, it is observed that communication of gradients becomes a bottleneck that significantly drags down the training speed. While for federated learning, the system capacity is usually confined by communication latency between the server and workers. Memory constraints are less studied compared with bandwidth (Sohoni et al., 2019). However, given the trend of an increasing model size in both computer vision and natural language processing, e.g., ViT (0.6B) (Dosovitskiy et al., 2020), Megatron-LM (8.3B) (Shoeybi et al., 2019), T5 (11B) (Raffel et al., 2019), it evidences that memory is becoming an unavoidable constraint for training large models in current GPU generation with 16/32GB memory. For federated learning, it is also necessary to reduce the memory footprint on workers especially for those IoT applications.
In this work, we will focus on the parsimonious setting, jointly dealing with limited bandwidth and memory. The priority is still communication efficiency for speeding up training time, otherwise one can simply rely on SGD with almost no additional memory consumption. Many existing works are unable to cope with such a challenging problem because memory is not carefully taken care of.
Communication efficiency. A well-documented approach to alleviate communication bottleneck is to compress the (stochastic) gradients such that training can be made faster with reduced overhead (Seide et al., 2014; Alistarh et al., 2017; Karimireddy et al., 2019). Depending on whether the gradient compressor is biased, these schemes are categorized as follows.
Unbiased compressors maintain the unbiasedness of stochastic gradients at the price of enlarged variance. For example, uniform quantization of stochastic gradients is studied in QSGD (Alistarh et al., 2017). A concurrent work (Wen et al., 2017) focuses on a special case of QSGD by quantizing each entry of the gradient into {±1, 0}. Other variants include natural quantization (Horvath et al., 2019), non-uniform quantization (Ramezani-Kebrya et al., 2021) and adaptive quantization (Faghri et al., 2020). Another route to obtain an unbiased compressor is through (scaled) gradient sparsification (Wangni et al., 2018) or its generalized form, atomic decomposition (Wang et al., 2018). However, such methods may be practically slow with performance degradation (Vogels et al., 2019). Biased gradient compressors, on the other hand, have been more successful in practice since not only do they support a more impressive compression ratio, but also a test accuracy comparable with SGD can be obtained in many scenarios. Examples of such compressors contain top-k or (unscaled) random-k sparsification (Lin et al., 2018; Alistarh et al., 2019), signSGD (Bernstein et al., 2018a;b), and PowerSGD (Vogels et al., 2019). Due to the bias, such compressors typically rely on error feedback (EF) schemes to ensure convergence (Stich et al., 2018; Karimireddy et al., 2019). Memory concerns in communication efficient methods. Memory concerns are highly entangled with communication efficiency in distributed training. On the one hand, memory footprint can be alleviated through a smaller batchsize (Krause et al., 2016; Spring et al., 2019), leading to increased iterations as well as communication rounds per epoch. Hence, communication efficient methods are particularly useful for accelerating training time in such cases. On the other hand, existing communication efficient methods are challenged by limited memory due to several reasons. First, a smaller batch size, which enlarges the variance, degrades the applicability of unbiased gradient compressors. Confirmed in our experiments, the exploded variance typically calls for a small step size which decelerates convergence. Unbiased gradient compressors such as quantization, are also hindered by the need of AllGather (Xu et al., 2020), since peak memory proportional to the number of workers is required. This renders them inefficient for the parsimonious setting, especially when the number of workers is large. Memory is also not handled well in methods with biased gradient compressors due to the need of error feedback, where additional memory the same as model size has to be consumed to keep track of accumulated compression error. Note that although no memory footprint is introduced in Local SGD (Stich, 2019), it suffers from variance blown problem as well. More importantly, local SGD cannot be integrated with ZeRO 3 (Rajbhandari et al., 2020), state-ofthe-art method for coping with limited memory, due to partitioned optimizer states among workers. In this work, we will focus on EF-based methods because of the capability of adopting a biased gradient compressor, which is typically contractive thus more robust to the gradient variance in a small batchsize setting. The additional memory consumption of EFSGD over SGD is the local error vector, which is delicate and critical for ensuring convergence. Meanwhile, the designed manner for memory efficiency can only come with lightweight and negligible runtime overhead, otherwise it violates the ultimate goal of communication efficient methods. Contractive error feedback is thus introduced to address these challenges. In a nutshell, our contributions can be summarized as:
• To the best of our knowledge, this is the first work that systematically studies memory concerns in communication efficient methods. Contractive error feedback (ConEF), is introduced to effectively manage memory given the preference of an allreducable and biased gradient compressor.
• Convergence of ConEF is established. Our theoretical results suggest a tradeoff between memory efficiency and a faster convergence, and ConEF finds the sweet spot.
• ConEF is capable of saving 80% – 90% additional memory of EFSGD on various learning problems such as image classification, language modeling, and machine translation. With almost the same runtime of EFSGD, ConEF achieves 1.3x – 5x speedup over SGD. Though test performance slightly drops on smaller models, the proposed method is more useful and reliable for larger networks where improved test accuracy over EFSGD is observed.
Notation. Bold lowercase (uppercase) letters denote vectors (matrices); ‖x‖ stands for `2 norm of x; and 〈x,y〉 denotes the inner product of x and y. In addition, we use [x]i ([X]i,j) to denote the i-th entry of vector x (i, j-th entry of matrix X).

heading: 2 PRELIMINARIES
text: This section briefly recaps EFSGD (Stich et al., 2018; Karimireddy et al., 2019) listed under Alg. 1. We use git to denote the stochastic gradient at xt on worker i, and assume that the stochastic gradients are mutually independent among different workers per iteration. In line 5, the scaled
Algorithm 1 EFSGD 1: Initialize: x0 ∈ Rd, ei0 = 0 ∈ Rd,∀i, η 2: for t = 0, 1, . . . , T − 1 do 3: assert xt = xit for every worker i 4: for worker i = 1, . . . , N in parallel do 5: pit = ηg i t + e i t
6: ∆it = Q(pit) 7: ∆t = Aggregate(∆it,∀i) 8: xt+1 = xt −∆t 9: eit+1 = p i t −∆it
10: end for 11: end for
Algorithm 2 ConEF 1: Initialize: x0 ∈ Rd, ei0 = 0 ∈ Rd,∀i, η 2: for t = 0, 1, . . . , T − 1 do 3: assert xt = xit for every worker i 4: for worker i = 1, . . . , N in parallel do 5: pit = ηg i t + e i t
6: ∆it = Q(pit) 7: ∆t = Aggregate(∆it,∀i) 8: xt+1 = xt −∆t 9: eit+1 = C(pit −∆it)
10: end for 11: end for
stochastic gradient is augmented via accumulated compression error eit, then compressed and communicated. There is no restriction to the biasedness of gradient compressor Q. Examples of Q include (scaled) sign, random/top-k, and powerSGD. The “aggregate” in line 7 refers to ∆t = 1 N ∑ i ∆ i t, and the compression error e i t+1 is recomputed in line 9 after updating model xt. More implementation details on communication of gradients, and reasons for the preference of AllReduce, are deferred to Apdx. A.1 to save space. Besides the widely appreciated allreducable gradient compressors, those biased ones are more tailored for settings where the inclination is to rely on a smaller batchsize for memory saving. This is because biased gradient compressors are usually contractive; see more details shortly in Assumption 4. However, due to the need of error feedback to ensure convergence, the additional memory consumption of eit can confine the applicability of biased compressors as well as EFSGD in parsimonious setups, and this issue is somehow omitted by existing works. Other error feedback variants, e.g., (Wu et al., 2018; Basu et al., 2019; Xu et al., 2021; Richtárik et al., 2021), also rely on additional vectors for the compression error, hence may benefit from the proposed technique as well.

heading: 3 MEMORY SAVING VIA CONTRACTIVE ERROR FEEDBACK
text: To endow EFSGD with memory efficiency, contractive error feedback is introduced in this section. The key idea is to apply another compressor C on the error vector eit such that the memory footprint can be mitigated. The proposed method is summarized in Alg. 2, where the name contractive error feedback (ConEF) originates from the fact that the error vector is represented in its compressed format. Note that the encoding and decoding of compression in Alg. 2 are omitted for notational convenience. Q and C are adopted to denote the gradient and error compressors, respectively, to highlight their different roles. The gradient compressor Q is general and a biased one is more recommended for small batchsizes that can squeeze more memory out. An unbiased error compressor C is preferable since it is helpful for generalization as we shall discuss later in Section 4. Moreover, the implementation of C has to be lightweight to avoid slowing down runtime. We will focus on theoretical properties first, and practical guidances for the choice of C are deferred to Section 4.

heading: 3.1 CONVERGENCE
text: Convergence of Alg. 2 is established in this subsection under standard assumptions for communication efficient algorithms in the distributed setting (Alistarh et al., 2017; Karimireddy et al., 2019; Stich et al., 2018; Zheng et al., 2019; Abdi & Fekri, 2020; Basu et al., 2019; Alistarh et al., 2019). Assumption 1. The objective function f : Rd → R has L-Lipchitz continuous gradients; that is, ‖∇f(x)−∇f(y)‖ ≤ L‖x− y‖,∀x,y ∈ Rd. Assumption 2. The stochastic gradient git is unbiased with bounded variance E [ git|xt ] = ∇f(xt), and E[‖git −∇f(xt)‖2|x] ≤ σ2. The gradient is also upper bounded E[‖∇f(xt)‖2] ≤ G2. Assumption 3. (Unbiased error compressor C). For any given x, E[C(x)|x] = x. In addition, there exist θ ≥ 0 such that E[‖C(x)− x‖2|x] ≤ θ‖x‖2. Assumption 4. (Contractive gradient compressor Q). For any given x, there exists δ ∈ (0, 1) such that E[‖Q(x)− x‖2|x] ≤ δ‖x‖2. Holding true for many allreducable and biased gradient compressors such as unscaled random(block)-k and powerSGD, Assumption 4 implies contractivity – the compression error is always smaller than the norm square of the compressed vector. Such compressors are more preferable when working with small batchsizes (for memory saving) since they do not amplify the gradient variance
as explained in Apdx. A.2 and (Stich et al., 2018). Unbiased compressors satisfying Assumption 4 are usually not allreducable; see also Apdx. A.2. The convergence of Alg. 2 is established below. Theorem 1. Suppose that Assumptions 1 - 4 hold and T is large. Choosing η = O ( 1
L( √ T/N+T 1/3)
) ,
Alg. 2 guarantees that
1
T T−1∑ t=0 E [ ‖∇f(xt)‖2 ] ≤ O
( L ( f(x0)− f∗ ) + (σ2 +G2)√
NT
) +O ( δθ(σ2 +G2)√
NT
) . (2)
The dependence on N demonstrates that linear speedup is achieved by ConEF. The second term in the right hand side of (2) is the additional error term incurred by compression. Next, we compare it with existing works to gain deeper understanding of ConEF. Upon relating Theorem 1 with EFSGD (Karimireddy et al., 2019), the price paid for memory efficiency can be clearly visualized. Since only the single worker case is considered in the EFSGD paper, we let N = 1 in Theorem 1 for a fair comparison. From the additional error term of EFSGD, i.e., O( δ(σ
2+G2) T ), it is observed that
ConEF converges slower. However, ConEF still benefits from the locally compressed error vector when further compared to QSGD which does not employ error feedback schemes (Alistarh et al., 2017). Suppose that Q is unbiased with Assumption 3 satisfied, then the error term in QSGD is O ( (1+θ)(σ2+G2)√
NT
) . Therefore, ConEF improves the dependence on θ over QSGD. In addition, The-
orem 1 is also tighter than (Horváth & Richtárik, 2020) since their error term is O( (1+δθ)(σ 2+G2)√
NT ).
ConEF finds the sweet spot. The comparisons above suggest a tradeoff between memory efficiency and a fast convergence. QSGD is at one extreme where no extra memory is needed at the price of slower convergence due to the blown variance introduced by unbiased gradient compressors. Note also that in general biased gradient compressors cannot be applied in QSGD. EFSGD is at another extreme, where the extra memory consumption is induced by eit with the benefit of a small error term in the convergence rate. Our proposed Alg. 2 lives in between EFSGD and QSGD, it saves memory compared to EFSGD, but incurs an error term smaller than QSGD yet larger than EFSGD.
3.2 IMPROVED CONEF
Algorithm 3 iConEF 1: Initialize: x0, ẽi0 = 0,qi0 = 0, η 2: for t = 0, 1, . . . , T − 1 do 3: assert xt = xit for every worker i 4: for worker i = 1, . . . , N in parallel do 5: pit = ηg i t + ẽ i t + q i t
6: ∆it = Q(pit) 7: ∆t = Aggregate(∆it,∀i) 8: xt+1 = xt −∆t 9: if version 1 then:
10: ẽit+1 = C(pit −∆it) . iConEF-v1 11: else: 12: ẽit+1 = Q(pit −∆it) . iConEF-v2 13: end if 14: qit+1 = C(pit −∆it − ẽit+1) 15: end for 16: end for In this subsection, we show that when Q (or C) is accurate, the dependence on θ (resp. δ) in Theorem 1 can be improved. The basic idea is to apply contractive error feedback on top of contractive error feedback such that the compression error can be further reduced. The resultant algorithm, improved ConEF (iConEF), is summarized in Alg. 3. Depending on whether C orQ is used twice, there are two versions of iConEF as marked in lines 9 – 13. Equipping with an accurate C with θ < 1, the compression error can be reused to retain performance. Mathematically, this error pit −∆it − ẽit+1 is once again compressed and added back to pit+1 in the next iteration to reduce the variance of (ẽit+1 + q i t+1); see (20). However, this approach might require to store both ẽit+1 and q i t+1, thus less efficient in terms of memory consumption compared to Alg. 2 with the exception of cases where C is a linear compressor such as count sketch. Theorem 2. Suppose that Assumption 3 holds with 0 ≤ θ < 1. Given Assumptions 1, 2 and 4, and choosing η = O ( 1
L( √ T/N+T 1/3)
) , iConEF-v1 guarantees that
1
T T−1∑ t=0 E [ ‖∇f(xt)‖2 ] ≤ O
( L ( f(x0)− f∗ ) + (σ2 +G2)√
NT
) +O ( δθ2(σ2 +G2)√
NT
) .
The θ2 in the last term tightens Theorem 1 showcasing the benefit of the third compressor. Examples of an accurate C satisfying 0 ≤ θ < 1 include quantization working in “dense regime” (Alistarh
et al., 2017). In particular, when the quantization level is chosen as s = 2 √ d, we have θ = 14 . Note also that it is possible to encode a compressed vector with less than 3.44d+ 32 bits (Alistarh et al., 2017, Lemma A.6). This means the memory reduction is almost 90%. Another example is natural compression that can offer at least 3.5x memory saving with θ ≤ 18 (Horvath et al., 2019). Similarly, when Q satisfies Assumption 4, which amounts to most of commonly adopted biased gradient compressors with δ < 1, the dependence on δ can be improved through iConEF-v2. Theorem 3. Suppose that Assumptions 1 – 4 hold. Choosing η = O ( 1
L( √ T/N+T 1/3)
) , iConEF-v2
guarantees that
1
T T−1∑ t=0 E [ ‖∇f(xt)‖2 ] ≤ O
( L ( f(x0)− f∗ ) + (σ2 +G2)√
NT
) +O ( δ2θ(σ2 +G2)√
NT
) .

heading: 4 A CLOSER LOOK AT ERROR COMPRESSORS
text: Previous sections have coped with contractive error feedback in its general form. In what follows, we will discuss the error compressor C in detail to gain more insights and practical guidances.

heading: 4.1 GENERALIZATION PROPERTIES
text: The merits of an unbiased C over a biased one come from generalization properties. To see this point, consider overparameterized least square problems (Wilson et al., 2017) as an example
min x f(x) :=
1 2 ‖Ax− y‖2 (3)
where A ∈ Rn×d with d n is the data matrix and y ∈ Rd denotes the labels. Since d is larger than the number of data samples n, this overparameterized problem has multiple global minima with zero loss collected in the set X ∗ := {x|f(x) = 0}. Define R(X) as the range of X. Given the fact that any stochastic gradient of (3) lies in R(A>), it is straightforward that SGD converges to a point inR(A>) when initialized well. It is further shown that the solution SGD converges to has the smallest norm among X ∗, i.e., x∗ = arg minx∈X∗ = A>(AA>)−1y. This x∗ is also known as the maximum margin solution, which leads to potential generalization benefits (Valiant, 1984; Cortes & Vapnik, 1995). In fact, for problem (3), if the iterates of a converging algorithm lie in R(A>), finding the maximum margin solution is ensured (Karimireddy et al., 2019, Lemma 9).
EFSGD converges to a point that is sufficiently close to R(A>) (Karimireddy et al., 2019). This gives the potential generalization benefit when adopting biased gradient compressors Q compared with e.g., sign-SGD (Bernstein et al., 2018a;b). Next, we show that such a generalization merit is maintained in Alg. 2 and its variants when C is unbiased. Theorem 4. If Assumptions 1 - 4 hold, initialized at x0 ∈ R(A>), ConEF ensures E [ xt − 1 N ∑ i e i t ] ∈ R(A>); moreover, iConEF-v1 and v2 guarantee E [ xt− 1N ∑ i q i t− 1N ∑ i ẽ i t ] ∈ R(A>).
Note that for ConEF, we have ‖E[ 1N ∑N i=1 e i t]‖2 ≤ E[‖ 1N ∑N i=1 e i t‖2] = O(η2). This implies that xt is close to the maximum margin solution at convergence. We are unable to obtain such a property for a biased C. The potential generalization merits justify the unbiased choice for C.

heading: 4.2 CHOICES OF ERROR COMPRESSOR
text: Recall that a prudent error compressor C for ConEF should: i) satisfy Assumption 3; and, ii) have a lightweight implementation to avoid runtime overhead. Two representative examples, quantization and count sketch, will be discussed in detail in this subsection.

heading: 4.2.1 QUANTIZATION
text: Quantization is a large family of compression schemes. For the ease of presentation, here we only describe the basic version (Alistarh et al., 2017). Given s uniformly distributed quantization levels in [0, 1], this approach quantizes the i-the entry of e as
C([e]i) = ‖e‖ · sign([e]i) · ξsi ([e]i) (4)
where ξsi (·) is a random variable designed to ensure the unbiasedness of C([e]i). In particular, let integer l ∈ [0, s] such that |[e]i|‖e‖ ∈ [ ls , l+1s ], and then ξsi ([e]i) is set to l+1s with probability |[e]i| ‖e‖ s− l and ls otherwise. Assumption 3 is naturally satisfied according to the proposition below. Proposition 1. (Alistarh et al., 2017) For any e, the quantization in (4) satisfies: i) E[C(e)] = e; ii) E[‖C(e) − e‖2] ≤ min{ ds2 , √ d s }‖e‖2; and, iii) there exists an encoding scheme that requires
Õ(s2 + s √ d) bits in expectation to represent C(e).
When setting quantization level s = O(1), the memory saving is impressive since it only requires O( √ d) bits for a d-dimensional vector at the cost of additional implementation burden of e.g., Elias encoding. Stochastic quantization is also suitable for the error compressor C in iConEF-v1 since it can ensure θ < 1 when s = O( √ d).
Other quantization approaches can also be applied as the error compressor. For instance, one can use int-quantization (Mishchenko et al., 2021), where a fp32 number is compressed into int8 or int16; and non-uniform quantization (Ramezani-Kebrya et al., 2021) is also a valid choice.

heading: 4.2.2 COUNT SKETCH
text: While quantization is straightforward to come up with, it relies on bit operations for compression rather than working on float numbers directly. Therefore, we introduce count sketch (Charikar et al., 2002) in this subsection, whose implementation is simply incrementing entries in a matrix.
Count sketch compresses a vector e by maintaining a matrix S of size v × w < d, i.e., C(e) := S. Both v and w are chosen according to some desirable accuracy guarantees discussed shortly. Count sketch relies on hash functions hi : {1, 2, . . . , d} 7→ {1, 2, . . . , w} and random sign functions si : {1, 2, . . . , d} 7→ {+1,−1} in i-th row of S. Count sketches leverage [e]p to update each row i of S by incrementing hi(p)-th column by si(p)[e]p. Compression of e can be obtained by updating S with {(p, [e]p)}dp=1; see Fig. 1 (a). To decompress and estimate [e]p from S, one needs to calculate Median(si(p) · [S]i,hi(p),∀i). Replacing median with mean also works, but typically median performs better. Count sketch satisfies Assumption 3 as shown in the following proposition. Proposition 2. (Charikar et al., 2002) Count sketch is unbiased, i.e., E[C(e)] = e. With w = Θ( 1 2 ) and v = Θ(log dp ), count sketch ensures ‖e− C(e)‖∞ ≤ ‖e‖ with probability at least 1− p.
Proposition 2 (Spring et al., 2019) suggests that count sketch is unbiased. Using inequality ‖a‖∞ ≤ ‖a‖ ≤ √ d‖a‖∞, it is not difficult to verify ‖e− C(e)‖ ≤ √ d‖e‖. Although count sketch satisfies a high probability variant of Assumption 3 with θ = 2d, it works well in our experimental studies with the recommendation of v = 1 or v = 3. Hence, this error compressor is more practical relative to the theoretical bounds. Another useful property of count sketch is the linearity, i.e., C(e1 + e2) = C(e1) + C(e2). This is helpful since i) it opens the possibility to squeeze ẽt+1 and qt+1 in iConEF-v1 to further reduce the memory consumption by adding these two count sketches; and ii) it eases error reset (Basu et al., 2019; Xie et al., 2020), which synchronizes the local errors every a few iterations among all workers. The second merit is because count sketch is small in size and allreducable, and more on this can be found in Apdx. G.4. Empirically, count sketch performs best when the compressed vector follows power law1. We train ResNet18 on CIFAR10, and plot the magnitude of eit on the first worker at the end of epochs
1The magnitude-index plot is roughly a line in log-log scale.
{10, 20, 30, 40} in Fig. 1(b). It is observed that power law is not satisfied. However, in Fig. 1(c) we find that the error vectors in a variant of EFSGD, termed partial EFSGD (Abdi & Fekri, 2020) (see Alg. 4 in Apdx. F due to space limitation), follow broken power law, that is, a piece-wise linear function in the log-log plot. Partial EFSGD adds part of the compression error, i.e., (1 − β)eit for β ∈ [0, 1), to stochastic gradients before compressing. As we shall see in numerical experiments, the update in partial EFSGD is helpful to save more memory when working with count sketch.

heading: 5 NUMERICAL RESULTS
text: In this section, experimental studies are carried out to validate the proposed algorithms. Since our major goal is to push the limit of memory saving, we will mostly focus on ConEF, and experiments of iConEF can be found in Apdx. G.4. Overall, the experiments are designed to answer three research questions (RQs): 1) how should we choose hyperparameters in ConEF. Since quantization is well-known in community, we put more efforts on the less understood count sketch; 2) when the gradient is aggressively compressed, how much memory can be saved; and, 3) how will the proposed method perform on different learning tasks. Although allreducable gradient compressors are advocated, we also include numerical results for scaled-sign, which requires AllGather for communication, to mimic parameter-server scenarios in RQ1. Testing various gradient compressors is also helpful to demonstrate the generality of ConEF. We follow common practice where both gradient and error compressors are applied in their layerwise form. Our implementation utilizes torch.DistributedDataParallel so that the bucket trick can be leveraged to overlap gradient computation and communication, and NCCL is chosen as communication backend.

heading: 5.1 RQ1: GUIDANCE ON HYPERPARAMETER CHOICE
text: In this subsection, we focus on the parameter choices of count sketch and quantization in ConEF using scaled-sign gradient compressor Q(g) = ‖g‖1d sign(g) as an example. About 97% communication overhead can be reduced. As AllGather is required for communication, the peak memory consumption is increased. Such a gradient compressor is adopted to simulate the parameter-server setting, where there is enough memory capacity on the server side. We test this setting with 4 GPUs on an Amazon EC2 p3.8xlarge instance, where more implementation details are put into Apdx. G.1 given the page limitation. When using allreducable gradient compressors, the general guidance for parameter choices is similar, although the numbers might change; see Apdx. G.1 for additional illustration of this point. Most of experiments are tested on CIFAR10,2 which consists of 60000 32x32 images in 10 classes with 50000 training samples and 10000 test ones. Standard data augmentation and preprocessing techniques are employed (He et al., 2016). We start with comparing contractive error feedback for vanilla and partial EFSGD using count sketches. With v = 1, the sketch size has to be 0.95x model size to ensure a good convergence for compressing the error vector in vanilla EFSGD. On the other hand, the error vector of partial EFSGD follows the broken power law, and hence is more suitable for count sketches. We list the choice of β and the smallest sketch size for convergence in Table. 1. Note that only {0.9, 0.8, 0.7, 0.6, 0.5, 0.4, 0.3, 0.2, 0.1} are tested, though there is no barrier to choose other values. It is observed that a larger β leads to more saved memory. This is possibly because that a larger β helps to suppress the error accumulation in count sketches; see Apdx. F.
Next we focus on choices for v and w in count sketches. We implement the partial EFSGD variant
of ConEF (see Alg. 4 in Apdx. F) with memory budget as 0.6x model size. Two specific sketches, v = 1, w = 0.6x model size and v = 3, w = 0.2x model size, are tested. It can be seen in Fig. 2 that the first sketch with v = 1 performs better. This suggests that a larger w, leading to less collision of hash functions, is more numerically beneficial compared with an increased v. Similar observations also appear on allreducable compressors. Hnece, v = 1 is adopted in most of our experiments.
2https://www.cs.toronto.edu/ kriz/cifar.html
After clarifying the parameter choices, the effectiveness of ConEF in parameter-server settings is tested with ResNet18 on CIFAR10 for 150 epochs as shown in Fig. 3. The detailed parameter choices are deferred to Apdx. G.1. We use ConEF-CS, ConEF-Q16 and ConEF-Q8 to denote Alg. 2 with error compressors C being count sketch, 16-level quantization, and 8-level quantization, respectively. Note that when C is a stochastic quantizer, we follow common practice to replace the `2 norm in (4) with max-norm for superior performance (Alistarh et al., 2017). It is observed that on ResNet18, ConEF-Q16 achieves a test accuracy that is almost the same as SGD and EFSGD, demonstrating EFSGD uses unnecessary memory. On the other hand, 44Mb is saved on a 50Mb model with count sketch, while only incurring a slight drop of 0.24 on test accuracy compared with SGD. To further validate the potential of ConEF in smart phone based applications, we test MobileNetv2 (Sandler et al., 2018), a more tailored model for the targeted setting. We set w = 0.1x model size to save 8Mb on a 14Mb model compared with EFSGD, while almost losing no test accuracy. Another interesting observation is that the convergence of ConEF-CS tends to be faster in the first 80 epochs before decreasing the learning rate. This might be useful in federated learning with IoT devices where each worker may not participate in training for too many iterations.

heading: 5.2 RQ2: MEMORY SAVING WITH AGGRESSIVELY COMPRESSED GRADIENTS
text: Next we focus on allreducable gradient compressors. We use 16 GPUs on 4 Amazon EC2 p3.8xlarge instances for experiments.
To investigate the performance of ConEF under aggressively compressed gradients, powerSGD is selected (Vogels et al., 2019). PowerSGD finds a rank-r approximation to the gradient matrix G. It is allreducable and different levels of communication efficiency can be achieved by tuning r. This is a challenging setting for ConEF given that i) the performance of the heuristic powerSGD highly depends on eit, and ii) e i t plays a more important role when the gradient is aggressively compressed.
ResNet18 is trained on CIFAR10 with r = 4 in powerSGD, which reduces 98.5% communication overhead. Unlike (Vogels et al., 2019), a small batchsize, e.g., 16 per worker, is considered here for reducing memory requirement (Krause et al., 2016). Mixed precision trick is also adopted to facilitate memory saving, where we use fp16 numbers on count sketches instead of fp32 (Micikevicius et al., 2018). Two more benchmarks are compared with. qSGD (Vogels et al., 2019): SGD with an unbiased and allreducable gradient compressor that has the same communication overhead as powerSGD. In particular, the gradient matrix G ∈ Rn×m is approximated by (GU)U> for a random matrix U ∈ Rm×r. Small-sized matrices GU and U> are communicated. We term such a scheme as qSGD to distinguish it with QSGD (Alistarh et al., 2017). Another benchmark is hEFSGD, which is a heuristic method by keeping the top-k elements of eit in EFSGD for memory saving.
The results can be found in Fig. 4, where we mark the percentage of saved memory with ConEF in the figure legend. There are several observations. ConEF-CS has a similar runtime as EFSGD, and
both are 10% faster than SGD. This demonstrates the benefit of reduced communication overhead in distributed training. ConEFCS significantly outperforms qSGD, which confirms the necessity of additional memory to enable a good convergence in memory constraint setting with small batchsizes. When 60% memory is saved, ConEF-CS also outperforms hEFSGD by a large margin in terms of both runtime and test accuracy, validating the generalization merits of an unbiased error compressor (cf. Theorem 4). ConEF-CS has a comparable performance with EFSGD, and more saved memory leads to larger accuracy loss.

heading: 5.3 RQ3: MEMORY SAVING ON VARIOUS LEARNING TASKS
text: In this subsection, ConEF is tested on different learning tasks. The unscaled random block gradient compressor is adopted, where a continuous block of k coordinates are chosen in a uniformly random fashion to improve 90% communication efficiency. When using the same random seed among GPUs, AllReduce is readily applied. Such a gradient compressor is more suitable for ConEF, since the error vector eit has d − k zeros. Three different training tasks are considered: i) ResNet-18 on CIFAR10; ii) WideResNet-28-10 on CIFAR10; and iii) LSTM for language modeling on Wikitest-2 (Merity et al., 2016). To demonstrate the necessity of additional memory, an unbiased variant of random block gradient compressor enabled by scaling with importance sampling is also implemented. We term the corresponding method as rbSGD. A setting with small batchsizes is still the main focus here, and the results can be found in Fig. 5.
On ResNet-18, it is observed that EFGSD and ConEF save 22% and 21% training time compared with SGD. The percentage of time saving improves over powerSGD even with a smaller compression ratio here. This is because more time is spent on the complicated encoding in powerSGD. It is observed that ConEF significantly outperforms rbSGD, once again demonstrating the need of additional memory when training with a small batchsize of 16. Compared with EFSGD, ConEF-CS drops 0.04 and 0.5 on test accuracy to save 60% and 90% memory, respectively.
For WideResNet-28-10, EFSGD and ConEF are 49.0% and 48.6% faster than SGD. Algorithms with (contractive) error feedback again considerably outperform rbSGD. EFSGD does not catch up with the test accuracy of ConEF-CS with 60% memory saving at a difference of 0.6, suggesting that EFSGD consumes memory inefficiently. ConEF-CS with 90% memory saving only drops 0.1 test accuracy compared with EFSGD. In light of the results on ResNet-18, a larger model (WideResNet) appears to be more robust when used jointly with contractive error feedback.
Next we focus on language modeling with LSTM on WikiText-2. All tested algorithms reduce about 78% training time of SGD. rbSGD fails to match the performance of both ConEF-CS and EFSGD. When 60% memory is saved, ConEF-CS has a 0.3 lower perplexity compared to EFSGD. ConEF-CS with 80% memory saved performs slightly worse than EFSGD with 0.2 higher test perplexity.
We also train a transformer for machine translation on Multi30K (Elliott et al., 2016) in Apdx. G. Adam based heuristic methods with (contractive) error feedback are adopted. In such a case, even 90% memory saved ConEF-CS has a slightly better performance than EF-Adam.
Additional experiments. Performance of iConEF-v1 and v2 can be found in Fig. 6. In this experiment, another allreducable gradient compressor, unscaled random-k, is adopted to reduce 90% communication. Both versions of iConEF outperform ConEF, validating our findings in Theorems 2 and 3. Due to space limitation, more implementation details and numerical results regarding random-k gradient compressor are deferred to Apdx. G.4, where we observe that ConEF can reduce 95% additional memory of EFSGD on ResNet18.

heading: 6 CONCLUSION AND FUTURE WORK
text: Contractive error feedback (ConEF) is introduced to alleviate memory concerns in communication efficient methods. The key insight is that the error vector can be compressed and an effective error compressor, such as count sketch, saves memory with almost no accuracy drop. ConEF is tested on various image classification and language processing tasks, reducing 80% − 90% of additional memory footprint in EFSGD. Our next step will focus on saving memory more aggressively.

heading: ETHICS STATEMENT
text: The main topic of this work is speeding up distributed training on devices with limited memory. Since this work focuses on an algorithmic perspective, the impact would be mainly scientific rather than ethical.

heading: REPRODUCIBILITY STATEMENT
text: All missing proofs can be found in Appendix. Detailed choices for hyperparameters and additional experiments are deferred to Appendix G due to space limitation, and publicly available code will be provided after review process to ensure anonymity.

heading: Supplementary Document

heading: A MISSING DETAILS

heading: A.1 HOW GRADIENTS ARE COMMUNICATED
text: While the “aggregation” in line 7 of Alg. 1 is a highly abbreviated term, in practice (compressed) gradients are communicated through AllReduce, AllGather, or parameter-server, among which allreducable compressors are widely appreciated. In contrast to AllGather, AllReduce enjoys low communication cost that is nearly independent of the number of workers (Vogels et al., 2019; Peng et al., 2019). In addition, compressed gradients that are allreducable relieve runtime since it only needs to be decoded once. In AllGather, a worker receives N compressed gradients that not only take additional space but also need to be decompressed individually. This has a negative impact on memory and runtime scalability given the number of workers can be large Vogels et al. (2019). When compared to parameter-server, AllReduce bypasses the need of double compression, where compression is applied to the bidirectional communication of clients and server. Parameter-server structure can also introduce communication congestion due to the centralized aggregation at the server side (Yu et al., 2018).

heading: A.2 BIASED GRADIENT COMPRESSORS ARE MORE SUITABLE FOR SMALL BATCHSIZES
text: Suppose that g is a stochastic approximation of ∇f(x) satisfying Assumption 2. We show that compressors satisfying Assumption 4, typically biased, are more robust to the large variance (σ2) in the small batchsize settings. Consider gradient compressorQ1 satisfying Assumption 4 with δ < 1. The compression error of Q1 can be calculated as
E[‖Q1(g)− g‖2] ≤ δE[‖g‖2] ≤ δE[‖∇f(x)‖2] + δσ2. (5)
Let Q2 denote an unbiased compressor satisfying Assumption 3. Typically to get a good compression ratio, we have θ > 1. The compression error is bounded by
E[‖Q2(g)− g‖2] ≤ θE[‖g‖2] ≤ θE[‖∇f(x)‖2] + θσ2. (6) Given θ > 1, the bound in (6) amplifies the variance; while the bound in (5) shrinks the variance. Therefore, a contractive compressor satisfying Assumption 4, which is typically biased, is more suitable for a small batchsize setting due to the large σ2. Although some of unbiased gradient compressors, such as QSGD working in “dense regime” (Alistarh et al., 2017) and natural dithering (Horvath et al., 2019), ensure θ < 1, unfortunately they are not allreducable, and hence are not suitable for our setting.
Other works such as (Stich et al., 2018) also observe the variance blow-up in unbiased gradient compressors.

heading: A.3 APPLYING CONEF TO ZERO
text: ZeRO (Rajbhandari et al., 2020) can be used to address memory issues for training large models in a distributed fashion. Mathematically, ZeRO is equivalent to Adam or other chosen optimizers, while the memory issues are addressed by partitioning optimizer states, gradients, and even model parameters among different workers. Here we consider the case of ZeRO3, where all aforementioned parts are partitioned among different workers.
Being most memory efficient, ZeRO3 increases the communication overhead. Hence, it is helpful to use gradient compression to speed up training. However, EF based approach cannot be applied here due to the requirement of a model-sized local error vector that is burdensome to GPU memory. In addition, because of the partitioned gradients, it is impossible to employ the memory trick in (Ramesh et al., 2021) where one puts local error vector to gradient buffer to save space.
We believe ConEF is helpful to reduce the runtime of ZeRO3 while respect constraints on GPUs memory. In this paper, we first demonstrate the feasibility and convergence of ConEF to clear up the theoretical barrier of integrating ConEF to ZeRO3. The detailed implementation is left for future work.

heading: B PROOF OF THEOREM 1
text: Proof. For convenience, define the compression error of C on worker i as ζit = eit+1 − pit + ∆it. Following Assumption 3, we have
E [ ζit |pit,∆it ] = 0, (7)
E [ ‖ζit‖2|pit,∆it ] = E [ ‖C(pit −∆it)− pit+∆it‖2|pit,∆it ] ≤ θ‖pit −∆it‖2. (8)
Next, we rewrite local error eit+1 using ζ i t as
eit+1 = p i t −∆it + ζit (a) = ηgit + e i t −∆it + ζit
where (a) follows from pit = ηg i t + e i t. Then summing over N workers, we have
1
N N∑ i=1 eit+1
= η
N N∑ i=1 git + 1 N N∑ i=1 eit − 1 N N∑ i=1 ∆it + 1 N N∑ i=1 ζit
(b) =
η
N N∑ i=1 git + 1 N N∑ i=1 eit + xt+1 − xt + 1 N N∑ i=1 ζit
where (b) is because 1N ∑N i=1 ∆ i t = ∆t = xt − xt+1. Rearranging the terms and define zt :=
xt − 1N ∑N i=1 e i t, we arrive at
zt+1 = zt − η
N N∑ i=1 git − 1 N N∑ i=1 ζit . (9)
The analysis will rely on (9). By Assumption 1, we have
f(zt+1)− f(zt) (10)
≤ 〈∇f(zt), zt+1 − zt〉+ L
2 ‖zt+1 − zt‖2
= − 1 N N∑ i=1 〈∇f(zt), ζit〉 − η N N∑ i=1 〈∇f(zt),git〉+ L 2 ∥∥∥ η N N∑ i=1 git + 1 N N∑ i=1 ζit ∥∥∥2 = − 1
N N∑ i=1 〈∇f(zt), ζit〉 − η N N∑ i=1 〈∇f(zt)−∇f(xt),git〉
− η N N∑ i=1 〈∇f(xt),git〉+ L 2 ∥∥∥ η N N∑ i=1 git + 1 N N∑ i=1 ζit ∥∥∥2. Given (7) and taking expectation w.r.t. the randomness ofQ, we arrive at E[ζit |pit] = E[ζit |git, eit] = 0. Then conditioning on xt and taking expectation w.r.t. the randomness of git, we have E[ζit |xt, eit] = 0. Hence, we have
E [ 〈∇f(zt), ζit〉|xt, eit ] = 〈∇f(zt),E[ζit |xt, eit]〉 = 0
⇒ E [ 〈∇f(zt), ζit〉 ] = 0. (11)
We also have E [ 〈∇f(zt)−∇f(xt),git〉|xt, eit ] = E [ 〈∇f(zt)−∇f(xt),∇f(xt)〉|xt, eit ] .
Then taking expectation w.r.t. xt and eit, we have −ηE [ 〈∇f(zt)−∇f(xt),git〉 ] = −ηE [ 〈∇f(zt)−∇f(xt),∇f(xt)〉 ] (12)
(c) ≤ η 2 E [ ‖∇f(xt)‖2 ] + η 2 E [ ‖∇f(zt)−∇f(xt)‖2 ] (d)
≤ η 2 E [ ‖∇f(xt)‖2 ] + ηL2 2 E [∥∥∥ 1 N N∑ i=1 eit ∥∥∥2] (e)
≤ η 2 E [ ‖∇f(xt)‖2 ] + η3L2 2
c
(1− c)λ (σ 2 +G2)
where (c) is by Young’s inequality; (d) uses Assumption 1 and zt = xt− 1N ∑N i=1 e i t; and (e) comes from Lemma 1.
Next, we have ηE [ 〈∇f(xt),git〉 ] = ηE [ E [ 〈∇f(xt),git〉|xt ]] = ηE [ ‖∇f(xt)‖2 ] . (13)
To proceed, we have
L 2 E [∥∥∥ η N N∑ i=1 git + 1 N N∑ i=1 ζit ∥∥∥2] (14) ≤ Lη2E [∥∥∥ 1 N N∑ i=1 git −∇f(xt) +∇f(xt) ∥∥∥2]+ LE[∥∥∥ 1 N N∑ i=1 ζit
∥∥∥2] ≤ Lη2E [ ‖∇f(xt)‖2 ] + Lη2E [∥∥∥ 1 N N∑ i=1 git −∇f(xt) ∥∥∥2]+ LE[∥∥∥ 1 N N∑ i=1 ζit
∥∥∥2] ≤ Lη2E [ ‖∇f(xt)‖2 ] + Lη2σ2
N + LE [∥∥∥ 1 N N∑ i=1 ζit ∥∥∥2] ≤ Lη2E [ ‖∇f(xt)‖2 ] + Lη2σ2
N + 1− c+ c/λ 1− c 2Lδθη2 N (σ2 +G2)
where the last inequality comes from Lemma 2.
Taking expectation on both sides of (10), and plugging (11), (12), (13), and (14) in, we have E [ f(zt+1)− f(zt) ] ≤ (−η
2 + Lη2)E
[ ‖∇f(xt)‖2 ] + c
(1− c)λ L2η3(σ2 +G2) 2 + η2σ2L N + 1− c+ c/λ 1− c 2δθη2L(σ2 +G2) N .
Rearranging the terms and dividing both sides with η, we arrive at
( 1
2 − Lη)E
[ ‖∇f(xt)‖2 ] ≤ E
[ f(zt)− f(zt+1) ] η
+ c (1− c)λ L2η2(σ2 +G2) 2 + ησ2L N + 1− c+ c/λ 1− c 2δθηL(σ2 +G2) N .
Summing over t, and using the facts z0 = x0 and f(zT ) ≥ f∗, we arrive at
( 1 2 − Lη) 1 T T−1∑ t=0 E [ ‖∇f(xt)‖2 ] ≤ ( f(x0)− f∗ ) ηT
+ c (1− c)λ L2η2(σ2 +G2) 2 + ησ2L N + 1− c+ c/λ 1− c 2δθηL(σ2 +G2) N .
Let η = min {
1 4L ,
1 L( √
T N+T 1/3)
} . In this case, we have 12 − Lη ≥ 14 , η ≤ ( 1LT 1/3 and η ≤ √ N L √ T .
Using these three inequalities, we have
1
4T T−1∑ t=0 E [ ‖∇f(xt)‖2 ] ≤ 4L
( f(x0)− f∗) T + L ( f(x0)− f∗)√ NT + L ( f(x0)− f∗) T 2/3
+ c (1− c)λ σ2 +G2 2T 2/3 + σ2√ NT + 1− c+ c/λ 1− c 2δθ(σ2 +G2)√ NT .
If we consider a sufficiently large T , i.e, T > N3, we have √ NT < T 2/3. Plugging this inequality in and viewing both λ and c as constants, we have
1
T T−1∑ t=0 E [ ‖∇f(xt)‖2 ] ≤ O
( L ( f(x0)− f∗ ) √ NT ) +O ( σ2 +G2√ NT ) +O ( δθ(σ2 +G2)√ NT ) .
The proof is thus completed.
Lemma 1. Suppose that there exist c ∈ (0, 1) and λ > 0 such that δ = c (1+λ)(1+ √ θ)2 . It is guaranteed to have
E [∥∥∥ 1 N N∑ i=1 eit+1 ∥∥∥2] ≤ c (1− c)λη 2(σ2 +G2),∀t ≥ 0.
Proof. To start with, we have
E [∥∥∥ 1 N N∑ i=1 eit+1 ∥∥∥2] = 1 N2 E [∥∥∥ N∑ i=1 eit+1 ∥∥∥2] ≤ 1 N N∑ i=1 E [ ‖eit+1‖2 ] . (15)
Conditioning on git, the randomness only comes from two compressors. E [ ‖eit+1‖2 ∣∣git] = E[‖pit −∆it + ζit‖2|git] (16) (a)
≤ (1 + α)E [ ‖pit −∆it‖2|git ] + ( 1 + 1
α
) E [ ‖ζit‖2|git ] (b)
≤ ( 1 + α+ θ + θ
α
) E [ ‖pit −∆it‖2|git ] (c)
≤ ( 1 + √ θ )2E[‖pit −∆it‖2|git]
(d) ≤ δ ( 1 + √ θ )2E[‖pit‖2|git]
= δ ( 1 + √ θ )2E[‖ηgit + eit‖2|git]
(e) ≤ ( 1 + 1
λ
) δ ( 1 + √ θ )2 η2‖git‖2 + (1 + λ)δ ( 1 + √ θ )2E[‖eit‖2|git]
where (a) is by Young’s inequality, i.e., ‖a + b‖2 ≤ (1 + α)‖a‖2 + (1 + 1α )‖b‖2 with α > 0 to be specified shortly; (b) is the result of (8); (c) is by choosing α = √ θ, which minimizes the coefficient; (d) results from Assumption 4; and (e) follows again from Young’s inequality.
Then taking expectation w.r.t. git, we arrive at
E [ ‖eit+1‖2 ] ≤ ( 1 + 1
λ ) δ ( 1 + √ θ )2︸ ︷︷ ︸
:=Bλ
η2E [ ‖git‖2 ] + (1 + λ)δ ( 1 + √ θ )2︸ ︷︷ ︸
:=Aλ
E [ ‖eit‖2 ]
Summing over i, we get
1
N N∑ i=1 E [ ‖eit+1‖2 ] ≤ Bλη 2 N N∑ i=1 E [ ‖git‖2 ] + Aλ N N∑ i=1 E [ ‖eit‖2 ] (17)
(f) ≤ Bλη2(σ2 +G2) + Aλ N N∑ i=1 E [ ‖eit‖2 ] (g) ≤ Bλ 1−Aλ η2(σ2 +G2).
where (f) follows from Assumption 2; and (g) is obtained by recursively unrolling 1 N ∑N i=1 E [ ‖eit‖2 ] and using the fact that ei0 = 0. Given the parameter choice δ = c (1+λ)(1+ √ θ)2 for some c ∈ (0, 1), we have Aλ = c and Bλ = cλ . The proof can be completed by using (15).
Lemma 2. Choosing the parameters the same as those in Lemma 1, it is guaranteed to have
E [∥∥∥ 1 N N∑ i=1 ζit ∥∥∥2] ≤ 1− c+ c/λ 1− c 2δθη2(σ2 +G2) N .
Proof. Given (7), it is not hard to see that for i 6= j, E [ 〈ζit , ζjt 〉|xt ] = 0
which implies that
E [∥∥∥ 1 N N∑ i=1 ζit ∥∥∥2] = 1 N2 N∑ i=1 E [ ‖ζit‖2 ] . (18)
Then we have from (8) that
E [ ‖ζit‖2|pit,∆it ] ≤ θ‖pit −∆it‖2
Then taking expectation w.r.t. the randomness of Q, and by Assumption 4, we have E [ ‖ζit‖2|pit ] ≤ δθ‖pit‖2 = δθ‖ηgit + eit‖2.
Finally, taking expectation w.r.t. pit, we have E [ ‖ζit‖2 ] ≤ δθE [ ‖ηgit + eit‖2 ] ≤ 2δθη2E [ ‖git‖2] + 2δθE [ ‖eit‖2 ] .
Summing over i, we arrive at
1
N N∑ i=1 E [ ‖ζit‖2 ] ≤ 2δθη2 1 N N∑ i=1 E [ ‖git‖2] + 2δθ 1 N N∑ i=1 E [ ‖eit‖2 ] ≤ 2δθη2(σ2 +G2) + 2δθ c
(1− c)λη 2(σ2 +G2)
≤ 1− c+ c/λ 1− c 2δθη 2(σ2 +G2)
where we use Assumption 2 and (17) in Lemma 1 simultaneously in the second last inequality. The proof is completed by applying (18).

heading: C PROOF OF THEOREM 2
text: The proof idea is to define eit+1 := ẽ i t+1+q i t+1 and to view e i t+1 as a single error compressor so that we can reuse the proof of Theorem 1. To see this, define compression error as ζit = e i t+1−pit+∆it. Then we have
E [ ζit |pit,∆it ] = E [ eit+1|pit,∆it ] − (pit −∆it) (19)
= E [ ẽit+1|pit,∆it ] + E [ C(pit −∆it − ẽit+1)|pit,∆it ] − (pit −∆it) = 0.
where the last equation uses the fact E [ C(pit −∆it − ẽit+1)|pit,∆it, ẽit+1 ] = pit −∆it − ẽit+1. This means that eit+1 is still an unbiased estimator of p i t −∆it. However, by adding the compression error back, the variance of ζit can be reduced, i.e.,
E [ ‖ζit‖2|pit,∆it ] = E [ ‖ẽit+1 + C(pit −∆it − ẽit+1)− pit + ∆it‖2|pit,∆it ] (20)
≤ θE [ ‖ẽit+1 − (pit −∆it)‖2|pit,∆it ] = θE [ ‖C(pit −∆it)− (pit −∆it)‖2|pit,∆it
] ≤ θ2‖pit −∆it‖2.
Comparing with (8), one can see that θ dependence is indeed improved given an accurate C with θ < 1. Next, we modify Lemmas 1 and 2 to finish the proof. Lemma 3. Suppose that there exists c ∈ (0, 1) and λ > 0 such that δ = c(1+λ)(1+θ)2 . It is guaranteed to have
E [∥∥∥ 1 N N∑ i=1 eit+1 ∥∥∥2] ≤ c (1− c)λη 2(σ2 +G2),∀t ≥ 0.
Proof. The proof is basically the same as Lemma 1 except for using (20) rather than (8) in (16) to tighten the bound. Hence, we only highlight the difference here.
E [ ‖eit+1‖2 ∣∣git] = E[‖pit −∆it + ζit‖2|git] (21) ≤ (1 + α)E [ ‖pit −∆it‖2|git ] + ( 1 + 1
α
) E [ ‖ζit‖2|git ] (a)
≤ ( 1 + α+ θ2 + θ2
α
) E [ ‖pit −∆it‖2|git ] (b)
≤ ( 1 + θ )2E[‖pit −∆it‖2|git] ≤ δ ( 1 + θ
)2E[‖pit‖2|git] = δ ( 1 + θ
)2E[‖ηgit + eit‖2|git] (e)
≤ ( 1 + 1
λ
) δ ( 1 + θ )2 η2‖git‖2 + (1 + λ)δ ( 1 + θ )2E[‖eit‖2|git] where in (a) we uses (20); (b) chooses α = θ.
Lemma 4. Choosing parameters the same as those in Lemma 3, it is guaranteed to have
E [∥∥∥ 1 N N∑ i=1 ζit ∥∥∥2] ≤ 1− c+ c/λ 1− c 2δθ2η2(σ2 +G2) N .
Proof. The proof is exactly the same as Lemma 2 except for using (20) rather than (8). Hence it is omitted here.
Proof of Theorem 2. The rest of proof is the same as Theorem 1, except for i) applying Lemma 3 in (12); and ii) applying Lemma 4 in (14).

heading: D PROOF OF THEOREM 3
text: Proof. The basic idea is again to define eit+1 = ẽ i t+1 + q i t+1 and to view e i t+1 as a single error compressor so that we can reuse the proof of Theorem 1. Applying (Horváth & Richtárik, 2020, Theorem 3), we have eit+1 is unbiased with bounded variance. In particular, we have
E [ eit+1 − (pit −∆it)|pit,∆it ] = E [ (ẽit+1 + q i t+1)− (pit −∆it)|pit,∆it ] = 0 (22)
E [ ‖eit+1 − (pit −∆it)‖2|pit,∆it ] = E [ ‖(ẽit+1 + qit+1)− (pit −∆it)‖2|pit,∆it ] (23)
≤ δθ‖pit −∆it‖2.
Therefore, we can reuse the derivation in Theorem 1. The only change we need is to modify Lemmas 1 and 2 to Lemmas 5 and 6, respectively. After applying these modified lemmas, Theorem 3 directly follows.
Lemma 5. Suppose that there exists c ∈ (0, 1) and λ > 0 such that δ = c (1+λ)(1+ √ θ)2 . It is guaranteed to have
E [∥∥∥ 1 N N∑ i=1 eit+1 ∥∥∥2] ≤ c (1− c)λη 2(σ2 +G2),∀t ≥ 0.
Proof. The proof is basically the same as Lemma 1, hence we only highlight the different steps in (16) here. Define ζit = e i t+1 − pit + ∆it, we have
E [ ‖eit+1‖2 ∣∣git] = E[‖pit −∆it + ζit‖2|git] (24) ≤ (1 + α)E [ ‖pit −∆it‖2|git ] + ( 1 + 1
α
) E [ ‖ζit‖2|git ] (a)
≤ ( 1 + α+ δθ + δθ
α
) E [ ‖pit −∆it‖2|git ] (b)
≤ ( 1 + √ δθ )2E[‖pit −∆it‖2|git]
(c) ≤ ( 1 + √ θ )2E[‖pit −∆it‖2|git]
where (a) is the result of (23); (b) is by choosing α = √ δθ; and (c) follows from δ < 1.
Lemma 6. Choosing the parameters the same as those in Lemma 5, it is guaranteed in iConEF-v2 to have
E [∥∥∥ 1 N N∑ i=1 ζit ∥∥∥2] ≤ 1− c+ c/λ 1− c 2δ2θη2(σ2 +G2) N .
Proof. The proof is almost the same as that of Lemma 2, hence we only highlight the differences.
From (23), we have
E [ ‖ζit‖2|pit ] ≤ δ2θ‖pit‖2 = δ2θ‖ηgit + eit‖2.
Finally, taking expectation w.r.t. pt, we have E [ ‖ζit‖2 ] ≤ δ2θE [ ‖ηgit + eit‖2 ] ≤ 2δ2θη2E [ ‖git‖2] + 2δ2θE [ ‖eit‖2 ] ≤ 1− c+ c/λ
1− c 2δ 2θη2(σ2 +G2).
where we use Assumption 2 and Lemma 5 simultaneously in the last inequality With this modification, Lemma 6 follows directly.

heading: E PROOF OF THEOREM 4
text: We focus on ConEF. The proofs for iConEF-v1 and iConEF-v2 follow the same steps (simply by defining eit+1 = ẽ i t+1 + q i t+1) and hence are omitted here. Recall that in (9) we have zt+1 =
zt − ηN ∑N i=1 g i t − 1N ∑N i=1 ζ i t with zt = xt − 1N ∑N i=1 e i t. This implies that
xt+1 − 1
N N∑ i=1 eit+1 = xt − 1 N N∑ i=1 eit − η N N∑ i=1 git − 1 N N∑ i=1 ζit .
Conditioned on {xt, {eit}i, {git}i} := Pt and only take expectation w.r.t. the randomness of C, we have
E [ xt+1 − 1
N N∑ i=1 eit+1|Pt ] = E [ xt − 1 N N∑ i=1 eit − η N N∑ i=1 git ∣∣∣Pt]− E[ 1 N N∑ i=1 ζit ∣∣∣Pt] = E [ xt − 1
N N∑ i=1 eit − η N N∑ i=1 git ∣∣∣Pt]. Then we further take expectation w.r.t. Pt and unroll E [ xt − 1N ∑N i=1 e i t ] , we can get
E [ xt+1 − 1
N N∑ i=1 eit+1 ] = x0 − ηE [ t∑ τ=0 1 N N∑ i=1 giτ ] .
Given that x0 ∈ R(A>), and git ∈ R(A>),∀t, it is straightforward to see that E [ xt+1 − 1 N ∑N i=1 e i t+1 ] ∈ R(A>). Hence the proof is completed.

heading: F PARTIAL CONEF/EFSGD
text: To utilize memory in the most efficient manner when adopting count sketches as the error compressor, we introduce a variant ConEF, termed partial ConEF in Alg. 4. Partial ConEF is adapted from partial EFSGD (Abdi & Fekri, 2020), which can be recovered when C in line 9 is an identity mapping. The main difference with ConEF is the introduction of the scaler β ∈ [0, 1), highlighted in blue. Partial EFSGD/ConEF adds part of the compression error, i.e., (1− β)eit, to the stochastic gradient before compressing; see line 5. The remained local error βeit is kept and included in the update of eit+1 in line 9 for the use of next iteration. The convergence of partial ConEF is straight forward when combining our proof with (Abdi & Fekri, 2020).
Algorithm 4 Partial EFSGD / partial ConEF 1: Initialize: x0 ∈ Rd, ei0 = 0 ∈ Rd,∀i, η, β ∈ [0, 1) 2: for t = 0, 1, . . . , T − 1 do 3: assert xt = xit for every worker i 4: for worker i = 1, . . . N in parallel do 5: pit = ηg i t + (1− β)eit
6: ∆it = Q(pit) 7: ∆t = Aggregate(∆it,∀i) 8: xt+1 = xt −∆t 9: eit+1 = C(βeit + pit −∆it)
10: end for 11: end for
The partial ConEF update in line 9 is helpful when C is count sketch. In particular, we can rewrite line 9 as eit+1 = C(eit − (1− β)eit + pit −∆it). This means that the compression error in eit is also scaled down or subtracted. This is helpful for suppressing the accumulation of compression error.
Partial ConEF benefits from the linearity of count sketch since the decompression of eit can be avoided when updating eit+1 in line 9. One can simply compress p i t−∆it and add it to the scaled eit. Moreover, auxiliary variables in global memory can be avoided by taking advantage of the shared memory of the GPU. Compared to a nonlinear error compressor C, count sketch further reduces the burden on runtime.
G MORE ON NUMERICAL EXPERIMENTS
Table 2: β vs. sketch size for convergence for unscaled random block k gradient compressor
β 0.1 0.2 0.3 0.4 0.5 0.6 0.7 0.8 0.9
smallest sketch size − 0.9 0.9 0.8 0.7 0.5 0.3 0.3 0.1
0 30 60 90 120 150 epoch
70
80
90
100
te st
a cc
ur ac
y
column row
Figure 7: ConEF with row or column hashes has similar performance.
Experiments are repeated 3 times and averaged results are reported. In the implementation of count sketches, the tensor trick can be employed to avoid computing too many hashes, which improves the runtime. In the tensor trick, instead of using elementwise compressing, we view the error vector as a matrix, and treat one row or column as an “entry” of the compressed vector, resulting in a tensor sketch (matrix when v = 1). More specifically, suppose the error vector lives in Rd, whose matrix view is Rn×m with d = mn. It takes n (resp. m) hash computation in a row(resp. column-) tensor trick, whilemn hashes are needed without
this trick. Differences on test accuracy between row- or column-based tensor tricks are not observed; see Fig. 7.

heading: G.1 PARAMETER CHOICES FOR RQ1
text: Experiments in this subsection is carried out using a single Amazon EC2 p3.8xlarge instance with 4 GPUs. The goal here is to simulate a parameter-server setting since we do not have a real testbed.
ResNet18 on CIFAR10. The model is trained for 150 epochs with weight decay 10−4 and momentum 0.9. We decrease the learning rate by a factor of 10 at epochs 80 and 120. Overall batchsize 128 (32 per worker). And the initial learning rate for SGD is tuned from {0.1, 0.075, 0.05, 0.025}, with final choice of 0.1. EFSGD, ConEF, and QSGD use the same learning rate. In ConEF-CS we use β = 0.9.
MobileNetV2 on CIFAR10. The model is trained for 150 epochs with weight decay 4 × 10−5 and momentum 0.9. We decrease the learning rate by a factor of 10 at epochs 80 and 120. Overall batchsize is 128 (32 per worker). And the initial learning rate for SGD is tuned from {0.1, 0.075, 0.05, 0.025}, with final choice of 0.1. The same learning rate is used for EFSGD and ConEF. β = 0.9 is adopted in ConEF-CS.
Choices of β for allreducable gradient compressors. While the parameter guidances for scaledsign gradient compressor are discussed, we also include the choice of β and the smallest sketch size for convergence in Table. 2 when the gradient compressor Q is unscaled random-block-k, which is allreducable. We choose k such that 90% of communication overhead is reduced. Similar to scaledsign, it is observed that a larger β leads to more saved memory. Hence, we recommend to choose β = 0.9 or even larger.

heading: G.2 PARAMETER CHOICES FOR RQ2
text: The experiments are carried out using 4 Amazon EC2 p3.8xlarge instances (16 GPUs in total).
ResNet18 on CIFAR10. The model is trained for 150 epochs with weight decay 10−4 and momentum 0.9. We decrease the learning rate by a factor of 10 at epochs 80 and 120. We focus on small batchsize setting with 16 data sample per worker (256 in total). The initial learning rate for SGD is tuned from {0.1, 0.075, 0.05, 0.025}, finally chosen as 0.1. EFSGD, ConEF, and hEFSGD use the same learning rate. ConEF-CS with 60% memory saving uses β = 0.85; while its 90% memory saving relies on β = 0.9.

heading: G.3 PARAMETER CHOICES FOR RQ3
text: All experiments are carried out using 4 Amazon EC2 p3.8x large instances (16 GPUs in total).
ResNet18 on CIFAR10. We train this model for 120 epochs and decrease the learning rate by 10 at epochs 60 and 100. Momentum is chosen as 0.9 and weight decay is set to 10−4. To save memory, we focus on small batch setting with batchsize 16 per GPU (hence 256 in total). The learning rate for SGD is tuned from {0.05, 0.1, 0.2}, and 0.2 is chosen. We use the same learning rate in EFSGD
and ConEF-CS. For ConEF-CS with 60% saved memory, we choose β = 0.9; and for ConEF-CS with 90% memory saving we use β = 0.95. For rbSGD, we test step sizes from {0.05.0.1, 0.2} and 0.1 is chosen.
WideResNet28-10 on CIFAR10. The model is trained 150 epochs. Learning rate is decreased by 5 on epochs 50, 90, and 120. We set weight decay as 5 × 10−4 and momentum 0.9. Nesterov momentum is also adopted. A small batchsize 16 per GPU (256 in total) is considered for memory saving. The learning rate for SGD is tuned from {0.05, 0.1, 0.2}, where 0.1 performs the best. EFSGD and ConEF-CS adopt the same learning rate. For ConEF-CS with 60% saved memory, we choose β = 0.75; and for ConEF-CS with 90% memory saving we use β = 0.995. For rbSGD, we test step sizes from {0.001, 0.005, 0.05.0.1} and 0.005 is the final choice. LSTM on Wikitext-2. We use a 2-layer LSTM with 672 hidden units and 672-dimension word embeddings. We set BPTT as 35, and clip the gradient norm to 0.25. The momentum is chosen as 0.9, aided with Nesterov momentum as well. Dropout rate is chosen as 0.5. We also consider a small batchsize setting, e.g., 2 per GPU (32 in total). The model is trained for 60 epochs where the learning rate is decreased by 4 at epochs 15, 30, and 45. The initial learning rate for SGD is tuned from {0.5, 1, 2, 4}, and 2 is chosen. The same learning rate schedule is adopted for EFSGD and ConEF-CS. For ConEF-CS with 60% saved memory, we choose β = 0.8; and for ConEF-CS with 80% memory saving we use β = 0.9. For rbSGD, we test step sizes from {0.5, 1, 2, 4} and 1 is chosen.

heading: G.4 ADDITIONAL EXPERIMENTS
text: Adam type algorithms on a transformer. A single layer transformer for English-Germen machine translation is trained on Multi30k dataset (Elliott et al., 2016). In particular, we use a 512 dimensional embedding layer. Both encoder and decoder have 3 layers with hidden dimension 512 per layer and 8-head attention is adopted. The initial learning rate for ADAM and other algorithms is set as 10−4. The transformer is trained for 30 epochs, where the learning rate is decreased by 4 at the end of epochs 15 and 25. Batchsize is set to 128 and dropout ratio is set as 0.1. The biased gradient compressor is chosen as unscaled random block k to reduce 90% communication overhead.
Adam with an allreducable and unbiased gradient compressor (rb-Adam) struggles in terms of training loss. EF-Adam converges even slower than its ConEF counterpart. This once again demonstrates
that EF-Adam does not take full advantage of memory. ConEF with 60% and 90% memory saved obtain similar training loss both slightly larger than Adam, but there is a gap on validation.
Regarding validation loss, rb-Adam fails to match with other tested algorithms. With 90% memory saved, ConEF-CS still outperforms EF-Adam slightly. ConEF-CS with 60% saved memory has a comparable performance to vanilla Adam, and both are much better than EF-Adam. These observations suggest that ConEF is an improved alternative of error feedback to endow Adam type optimizers with communication and memory efficiency in distributed training. Lastly, the validation gap between ConEF with 60% and 90% memory saved intuitively makes sense, since as shown in Theorem 4, the best generalization error can be achieved when subtracting averaged local error vectors from current model. The count sketch in ConEF-CS introduces larger compression error when more memory is saved (small sketch size), thus losing more on its generalization behavior.
Random-k gradient compressor with more memory saving. We also test the unscaled random-k gradient compressor on ResNet-18 with CIFAR10, where 90% communication burden is reduced. As shown in Fig. 9, EFSGD has a test accuracy of 93.42, while ConEF-CS with w = 0.05 (95% memory saving) and w = 0.02 (98% memory saving) achieve a test accuracy of 92.70 and 92.68. Hence, it is possible to have an aggressive memory saving with less than 1 loss on test accuracy. In addition, it is observed that the compressed error vector is helpful to get a faster convergence initially. But after decreasing the step size, it has negative impact resulting in the accuracy drop.
iConEF with a random-k gradient compressor. To validate the improvement of iConEF over ConEF, we train a ResNet-18 on CIFAR10, and the results are reported in Fig. 6. A randomk gradient compressor reducing 90% communication overhead is adopted. For ConEF, we use a s = 64-level stochastic quantizer as the error compressor. iConEF-v2 is considered in this case, and
we use top-k compressor (5% elements) as the first error compressor, and s = 64-level stochastic quantizer as the second error compressor. As shown in Fig. 6, although ConEF has a similar performance as EFSGD, iConEF-v2 outperforms both of them, validating its efficiency. iConEF-v1 is also included in Fig. 6. In this case, both error compressors are chosen as s = 256-level quantizer. iConEF-v1 converges much faster comparing with EFSGD initially, and it has a similar test accuracy after learning rate is decreased.
ConEF-CS can be used easily with error reset (ER). As discussed in Section 4.2.2, error reset (Basu et al., 2019; Xie et al., 2020) can be applied directly when the error compressor C is count sketch. Thanks to the fact that count sketch is allreducable and small in size, ER can be done through allreducing the compressed error among all workers every a few iterations, which in our case is 512. ER slightly improves the test accuracy of ConEF-CS by 0.1 as observed in Fig. 10. More specifically, we train a ResNet-18 on CIFAR10 with unscaled random-k gradient compressor to reduce 90% communication overhead, and choose β = 0.9 with a count sketch that saves 80% memory. As shown here, both ConEF and ER-ConEF outperform EFSGD, and ER-ConEF has the best test accuracy.

Review summaries:
The paper introduces a two-stage error correction procedure for approximated (doubly compressed) gradients. Multiple theorems are derived to prove convergence and scaling properties. The algorithm can be combined with any compression algorithm for the gradient. The algorithm is applied to CIFAR-10 with results showing favorable results. 

Motivated by the memory limitations of large-scale training, this paper proposes a new modification to the error-feedback algorithm for parallel optimization. In particular, the authors suggest compressing the error vector with a separate compressor, which results in the ConEF algorithm. Under the assumption that this compressor is unbiased and has compression ratio theta, the authors prove that the convergence rate has some terms of the same order as in error feedback and others get multiplied by a factor of theta.  In addition, the authors propose the iConEF algorithms that introduce an additional error feedback for the extra compressor. This leads to yet another sequence of compressed vectors that need to be maintained. On the other hand, if theta<1 (or delta<1), the rate will become theta (or delta) times better.

The authors propose to compress the local error in communication-efficient distributed training to reduce the memory to store the local error.  The paper also proposes to compress the local error twice, which doubles the memory cost to store the local error but improves the convergence compared with compressing once.  Experiments show the training performance of the proposed method with different compression ratio. But the settings are unrealistic and can not validate the claims.