ATOMO: Communication-efficient Learning via Atomic Sparsification

Introduction

Distributed computing systems have become vital to the success of modern machine learning systems. Work in parallel and distributed optimization has shown that these systems can obtain massive speed up gains in both convex and non-convex settings . Several machine learning frameworks such as TensorFlow , MXNet , and Caffe2 , come with distributed implementations of popular training algorithms, such as mini-batch SGD. However, the empirical speed-up gains offered by distributed training, often fall short of the optimal linear scaling one would hope for. It is now widely acknowledged that communication overheads are the main source of this speedup saturation phenomenon .

Communication bottlenecks are largely attributed to frequent gradient updates transmitted between compute nodes. As the number of parameters in state-of-the-art models scales to hundreds of millions , the size of gradients scales proportionally. These bottlenecks become even more pronounced in the context of federated learning , where edge devices (e.g., mobile phones, sensors, etc) perform decentralized training, but suffer from low-bandwidth during up-link.

To reduce the cost of of communication during distributed model training, a series of recent studies propose communicating low-precision or sparsified versions of the computed gradients during model updates. Partially initiated by a 1-bit implementation of SGD by Microsoft in , a large number of recent studies revisited the idea of low-precision training as a means to reduce communication . Other approaches for low-communication training focus on sparsification of gradients, either by thresholding small entries or by random sampling . Several approaches, including QSGD and TernGrad, implicitly combine quantization and sparsification to maximize performance gains , while providing provable guarantees for convergence and performance. We note that quantization methods in the context of gradient based updates have a rich history, dating back to at least as early as the 1970s .

An atomic decomposition represents a vector as a linear combination of simple building blocks in an inner product space. In this work, we show that stochastic gradient sparsification and quantization are facets of a general approach that sparsifies a gradient in any possible atomic decomposition, including its entry-wise or singular value decomposition, its Fourier decomposition, and more. With this in mind, we develop Atomo, a general framework for atomic sparsification of stochastic gradients. Atomo sets up and optimally solves a meta-optimization that minimizes the variance of the sparsified gradient, subject to the constraints that it is sparse on the atomic basis, and also is an unbiased estimator of the input.

We show that 1-bit QSGD and TernGrad are in fact special cases of Atomo, and each is optimal (in terms of variance and sparsity), in different parameter regimes. Then, we argue that for some neural network applications, viewing the gradient as a concatenation of matrices (each corresponding to a layer), and applying atomic sparsification to their SVD is meaningful and well-motivated by the fact that these matrices are “nearly” low-rank, e.g., see Fig. 1. We show that Atomo on the SVD of each layer’s gradient, can lead to less variance, and faster training, for the same communication budget as that of QSGD or TernGrad. We present extensive experiments showing that using Atomo with SVD sparsification, can lead to up to $2\times$ faster training time (including the time to compute the SVD) compared to QSGD, on VGG and ResNet-18, and SVHN and CIFAR-10.

Relation to Prior Work

Atomo is closely related to work on communication-efficient distributed mean estimation in and . These works both note, as we do, that variance (or equivalently the mean squared error) controls important quantities such as convergence, and they seek to find a low-communication vector averaging scheme that minimizes it. Our work differs in two key aspects. First, we derive a closed-form solution to the variance minimization problem for all input gradients. Second, Atomo applies to any atomic decomposition, which allows us to compare entry-wise against singular value sparsification for matrices. Using this, we derive explicit conditions for which SVD sparsification leads to lower variance for the same sparsity budget.

The idea of viewing gradient sparsification through a meta-optimization lens was also used in . Our work differs in two key ways. First, consider the problem of minimizing the sparsity of a gradient for a fixed variance, while we consider the reverse problem, that is, minimizing the variance subject to a sparsity budget. The second more important difference is that while focuses on entry-wise sparsification, we consider a general problem where we sparsify according to any atomic decomposition. For instance, our approach directly applies to sparsifying the singular values of a matrix, which gives rise to faster training algorithms.

Finally, low-rank factorizations and sketches of the gradients when viewed as matrices were proposed in ; arguably most of these methods (with the exception of ) aimed to address the high flops required during inference by using low-rank models. Though they did not directly aim to reduce communication, this arises as a useful side effect.

Problem Setup

In machine learning, we often wish to find a model $w$ minimizing the empirical risk

where $x_{i}\in^{d}$ is the $i$-th data point. One way to approximately minimize $f(w)$ is by using stochastic gradient methods that operate as follows:

where $w_{0}$ is some initial model, $\gamma$ is the step size, and $\widehat{g}(w)$ is a stochastic gradient of $f(w)$, i.e.it is an unbiased estimate of the true gradient $g(w)=\nabla f(w)$. Mini-batch SGD, one of the most common algorithms for distributed training, computes $\widehat{g}$ as an average of $B$ gradients, each evaluated on randomly sampled data from the training set. Mini-batch SGD is easily parallelized in the parameter server (PS) setup, where a PS stores the global model, and $P$ compute nodes split the effort of computing the $B$ gradients. Once the PS receives these gradients, it applies them to the model, and sends it back to the compute nodes.

Since variance is a proxy for speed of convergence, in the context of communication-efficient stochastic gradient methods, one can ask: What is the smallest possible variance of an unbiased stochastic gradient that can be represented with $k$ bits? Note that under the unbiased assumption, minimizing variance is equivalent to minimizing the second moment of the random vector. This meta-optimization can be cast as the following meta-optimization:

Here, the expectation is taken over the randomness of $\widehat{g}$. We are interested in designing a stochastic approximation $\widehat{g}$ that “solves” this optimization. However, it seems difficult to design a formal, tractable version of the last constraint. In the next section, we replace this with a simpler constraint that instead requires that $\widehat{g}(w)$ is sparse with respect to a given atomic decomposition.

Atomo: Atomic Decomposition and Sparsification

Let $(V,\langle\cdot,\cdot\rangle)$ be an inner product space over and let $\|\cdot\|$ denote the induced norm on $V$. In what follows, you may think of $g$ as a stochastic gradient of the function we wish to optimize. An atomic decomposition of $g$ is any decomposition of the form $g=\sum_{a\in\mathcal{A}}\lambda_{a}a$ for some set of atoms $\mathcal{A}\subseteq V$. Intuitively, $\mathcal{A}$ consists of simple building blocks. We will assume that for all $a\in\mathcal{A}$, $\|a\|=1$, as this can be achieved by a positive rescaling of the $\lambda_{a}$.

An example of an atomic decomposition is the entry-wise decomposition $g=\sum_{i}g_{i}e_{i}$ where $\{e_{i}\}_{i=1}^{n}$ is the standard basis. More generally, any orthonormal basis of $V$ gives rise to a unique atomic decomposition of any $g\in V$. While we focus on finite-dimensional vectors, one could use Fourier and wavelet decompositions in this framework for infinite-dimensional spaces. When considering matrices, the singular value decomposition gives an atomic decomposition in the set of rank-1 matrices. More general atomic decompositions have found uses in a variety of situations, including solving linear inverse problems .

We are interested in finding an approximation to $g$ with fewer atoms. Our primary motivation is that this reduces communication costs, as we only need to send atoms with non-zero weights. We can use whichever decomposition is most amenable for sparsification. For instance, if $X$ is a low rank matrix, then its singular value decomposition is naturally sparse, so we can save communication costs by sparsifying its singular value decomposition instead of its entries.

where $t_{i}\sim\text{Bernoulli}(p_{i})$, for $0<p_{i}\leq 1$. We refer to this sparsification scheme as atomic sparsification. Note that the $t_{i}$’s are independent. Recall that we assumed above that $\|a_{i}\|^{2}=1$ for all $a_{i}$. We have the following lemma about $\widehat{g}$.

An equivalent form of this optimization problem was previously presented in (Section 6.1). The authors considered this problem for entry-wise sparsification and found a closed-form solution for $s\leq\|\lambda\|_{1}/\|\lambda\|_{\infty}$. We give a version of their result but extend this to a closed-form solution for all $s$. A similar optimization problem was given in , which instead minimizes sparsity subject to a variance constraint.

We will show that the Algorithm 1 produces a probability vector $p\in^{n}$ solving (3) for $0<s\leq n$. While we show in Appendix B that this result can be derived using the KKT conditions, we use an alternative method that focuses on a relaxation of (3) in order to better understand the structure of the problem. This approach has the added benefit of shedding light on what variance is achieved by solving (3).

Note that (3) has a non-empty feasible set only for $0<s\leq n$. Define $f(p):=\sum_{i=1}^{n}\lambda_{i}^{2}/p_{i}$. To understand how to solve (3), we first consider the following relaxation:

We have the following lemma about the solutions to (4), first shown in .

Any feasible vector $p$ to (4) satisfies $f(p)\geq\dfrac{1}{s}\|\lambda\|_{1}^{2}$. This is achieved iff

Lemma 2 implies that if we ignore the constraint that $p_{i}\leq 1$, then the optimal $p$ is achieved by setting $p_{i}=|\lambda_{i}|s/\|\lambda\|_{1}$. If the quantity in the right-hand side is greater than 1, this does not give us an actual probability. This leads to the following definition.

An atomic decomposition $g=\sum_{i=1}^{n}\lambda_{i}a_{i}$ is $s$-unbalanced at entry $i$ if $|\lambda_{i}|s>\|\lambda\|_{1}$.

Fix the atomic decomposition of $g$. If there are no $s$-unbalanced entries then we say that the $g$ is $s$-balanced. We have the following lemma which guarantees that $g$ is $s$-balanced for $s$ not too large.

An atomic decomposition $g=\sum_{i=1}^{n}\lambda_{i}a_{i}$ is $s$-balanced iff $s\leq\|\lambda\|_{1}/\|\lambda\|_{\infty}$.

Lemma 2 gives us the optimal way to sparsify $s$-balanced vectors, since the $p$ that is optimal for (4) is feasible for (3). Moreover, the iff condition in Lemma 2 implies that the optimal assignment of the $p_{i}$ are between 0 and 1 iff $v$ is $s$-balanced. Suppose now that $g$ is $s$-unbalanced at entry $j$. We cannot assign $p_{j}$ as in (5). We will show that setting $p_{j}=1$ is optimal in this setting. This comes from the following lemma.

Suppose that $g$ is $s$-unbalanced at entry $j$ and that $q$ is feasible in (3). Then $\exists p$ that is feasible in (3) such that $f(p)\leq f(q)$ and $p_{j}=1$.

Lemmas 2 and 4 imply the following theorem about solutions to (3).

Suppose we sparsify $g$ as in (2) with sparsity budget $s$.

with equality if and only if $p_{i}=|\lambda_{i}|s/\|\lambda\|_{1}$.

and is minimized by $p$ with $p_{j}=1$ where $j=\operatorname*{argmax}_{i=1,\ldots,n}|\lambda_{i}|$.

This theorem implies that Algorithm 1 produces a vector $p\in^{n}$ solving (3). Note that due to the sorting requirement in the input, the algorithm requires $O(n\log n)$ operations. As we discuss in Appendix B, we could instead do this in $O(sn)$ operations by, instead of sorting and iterating through the values in order, simply selecting the next unvisited index $i$ maximizing $|\lambda_{i}|$ and performing the same test/updates. As we show in Appendix B, we need to select at most $s$ indices before the if statement in Algorithm 1 holds. Whether to sort or do selection depends on the size of $s$ relative to $\log n$.

Relation to QSGD and TernGrad

In this section, we will discuss how Atomo is related to two recent quantization schemes, 1-bit QSGD and TernGrad . We will show that in certain cases, these schemes are versions of the Atomo for a specific sparsity budget $s$. Both schemes use the entry-wise atomic decomposition.

QSGD takes as input $g\in^{n}$ and $b\geq 1$. This $b$ governs the number of quantization buckets. When $b=1$, this is referred to as 1-bit QSGD. 1-bit QSGD produces a random vector $Q(g)$ defined by

Here, the $\zeta_{i}\sim\text{Bernoulli}(|g_{i}|/\|g\|_{2})$ are independent random variables. A straightforward computation shows that $Q(g)$ can be defined equivalently by

where $t_{i}\sim\text{Bernoulli}(|g_{i}|/\|g\|_{2})$. Therefore, 1-bit QSGD exactly uses the atomic sparsification framework in (2) with $p_{i}=|g_{i}|/\|g\|_{2}$. The total sparsity budget is therefore given by

By Lemma 3 any $g$ is $s$-balanced for this $s$. Therefore, Theorem 5 implies that the optimal way to assign $p_{i}$ with this given $s$ is $p_{i}=|g_{i}|/\|g\|_{2}$. Since this agrees with (6), this implies that 1-bit QSGD performs variance-optimal entry-wise sparsification for sparsity budget $s=\|g\|_{1}/\|g\|_{2}$.

2 TernGrad

Similarly, TernGrad takes as input $g\in^{n}$, and produces a sparsified version $T(g)$ given by

where $\zeta_{i}\sim\text{Bernoulli}(|g_{i}|/\|g\|_{\infty})$. A straightforward computation shows that $T(g)$ can be defined equivalently by

where $t_{i}\sim\text{Bernoulli}(|g_{i}|/\|g\|_{\infty})$. Therefore, TernGrad exactly uses the atomic sparsification framework in (2) with $p_{i}=|g_{i}|/\|g\|_{\infty}$. The total sparsity budget is given by

By Lemma 3, any $g$ is $s$-balanced for this $s$. Therefore, Theorem 5 implies that the optimal way to assign $p_{i}$ with this given $s$ is $p_{i}=|g_{i}|/\|g\|_{\infty}$. This agrees with (7). Therefore, TernGrad performs variance-optimal entry-wise sparsification for sparsity budget $s=\|g\|_{1}/\|g\|_{\infty}$.

where $\zeta_{i}\sim\text{Bernoulli}(|g_{i}|/\|g\|_{q})$. Note that for all $i$, $|g_{i}|\leq\|g\|_{\infty}\leq\|g\|_{q}$, so this does give us a valid probability. We can define $L_{q}(v)$ equivalently by

By Lemma 3, the optimal way to assign $p_{i}$ with this given $s$ is $p_{i}=|g_{i}|/\|g\|_{q}$. Since this agrees with (8), Theorem 5 implies the following theorem.

Spectral-Atomo: Sparsifying the Singular Value Decomposition

In this section we compare different methods for matrix sparsification. The first uses Atomo on the entry-wise decomposition of a matrix, and the second uses Atomo on the singular value decomposition (SVD) of a matrix. We refer to this second approach as Spectral-Atomo. We show that under concrete conditions, Spectral-Atomo incurs less variance than sparsifying entry-wise. We present these conditions and connect them to the equivalence of certain matrix norms.

For a rank $r$ matrix $X$, denote its singular value decomposition by

When $p=q=\infty$, we define this to be $\|X\|_{\max}$ where $\|X\|_{\max}:=\max_{i,j}|X_{i,j}|$.

Comparing matrix sparsification methods:

Suppose that $V$ is the vector space of real $n\times m$ matrices. Given $X\in V$, there are two standard atomic decompositions of $X$. The first is the entry-wise decomposition

The second is the singular value decomposition

If $r$ is small, it may be more efficient to communicate the $r(n+m)$ entries of the SVD, rather than the $nm$ entries of the matrix. Let $\widehat{X}$ and $\widehat{X}_{\sigma}$ denote the random variables in (2) corresponding to the entry-wise decomposition and singular value decomposition of $X$, respectively. We wish to compare these two sparsifications.

In Table 1, we compare the communication cost and second moment of these two methods. The communication cost is the expected number of non-zero elements (real numbers) that need to be communicated. For $\widehat{X}$, a sparsity budget of $s$ corresponds to $s$ non-zero entries we need to communicate. For $\widehat{X}_{\sigma}$, a sparsity budget of $s$ gives a communication cost of $s(n+m)$ due to the singular vectors. We compare the optimal second moment from Theorem 5.

To compare the second moment of these two methods under the same communication cost, we $s$ and suppose $X$ is $s$-balanced entry-wise. By Theorem 5 and Lemma 3, the second moment in Table 1 is achieved iff

To achieve the same communication cost with $\widehat{X}_{\sigma}$, we take a sparsity budget of $s^{\prime}=s/(n+m)$. By Theorem 5 and Lemma 3, the second moment in Table 1 is achieved iff

For any $n\times m$ matrix $X$ over $,\frac{1}{\sqrt{nm}}\|X\|_{1,1}\leq\|X\|_{*}\leq\|X\|_{1,1}$.

Experiments

We present an empirical study of Spectral-Atomo and compare it to the recently proposed QSGD , and TernGrad , on a different neural network models and data sets, under real distributed environments. Our main findings are as follows:

We observe that spectral-Atomo provides a useful alternative to entry-wise sparsification methods, it reduces communication compared to vanilla mini-batch SGD, and can reduce training time compared to QSGD and TernGrad by up to a factor of $2\times$ and $3\times$ respectively. For instance, on VGG11-BN trained on CIFAR-10, spectral-Atomo with sparsity budget 3 achieves $3.96\times$ speedup over vanilla SGD, while 4-bit QSGD achieves $1.68\times$ on a cluster of 16, g2.2xlarge instances. Both Atomo and QSGD greatly outperform TernGrad as well.

We observe that spectral-Atomo in distributed settings leads to models with negligible accuracy loss when combined with parameter tuning.

We compare spectral-Atomocode available at: https://github.com/hwang595/ATOMO with different sparsity budgets to $b$-bit QSGD across a distributed cluster with a parameter server (PS), implemented in mpi4py and PyTorch and deployed on multiple types of instances in Amazon EC2 (e.g.m5.4xlarge, m5.2xlarge, and g2.2xlarge), both PS and compute nodes are of the same type of instance. The PS implementation is standard, with a few important modifications. At the most basic level, it receives gradients from the compute nodes and broadcasts the updated model once a batch has been received.

In our experiments, we use data augmentation (random crops, and flips), and tuned the step-size for every different setup as shown in Table 5 in Appendix D. Momentum and regularization terms are switched off to make the hyperparamter search tractable and the results more legible. Tuning the step sizes for this distributed network for three different datasets and eight different coding schemes can be computationally intensive. As such, we only used small networks so that multiple networks could fit into GPU memory. To emulate the effect of larger networks, we use synchronous message communication, instead of asynchronous.

Each compute node evaluates gradients sampled from its partition of data. Gradients are then sparsified through QSGD or spectral-Atomo, and then are sent back to the PS. Note that spectral-Atomo transmits the weighted singular vectors sampled from the true gradient of a layer. The PS then combines these, and updates the model with the average gradient. Our entire experimental pipeline is implemented in PyTorch with mpi4py , and deployed on either g2.2xlarge, m5.2xlarge and m5.4xlarge instances in Amazon AWS EC2. We conducted our experiments on various models, datasets, learning tasks, and neural network models as detailed in Table 2.

Scalability

We study the scalability of these sparsification methods on clusters of different sizes. We used clusters with one PS and $n=2,4,8,16$ compute nodes. We ran ResNet-34 on CIFAR-10 using mini-batch SGD with batch size $512$ split among compute nodes. The experiment was run on m5.4xlarge instances of AWS EC2 and the results are shown in Figure 2.

While increasing the size of the cluster, decreases the computational cost per worker, it causes the communication overhead to grow. We denote as computational cost, the time cost required by each worker for gradient computations, while the communication overhead is represented by the amount time the PS waits to receive the gradients by the slowest worker. This increase in communication cost is non-negligible, even for moderately-sized networks with sparsified gradients. We observed a trade-off in both sparsification approaches between the information retained in the messages after sparsification and the communication overhead.

End-to-end convergence performance

We evaluate the end-to-end convergence performance on different datasets and neural networks, training with spectral-Atomo(with sparsity budget $s=1,2,3,4$), QSGD (with $n=1,2,4,8$ bits), and ordinary mini-batch SGD. The datasets and models are summarized in Table 2. We use ResNet-18 and VGG11-BN for CIFAR-10 and SVHN . Again, for each of these methods we tune the step size. The experiments were run on a cluster of 16 compute nodes instantiated on g2.2xlarge instances.

The gradients of convolutional layers are 4 dimensional tensors with shape of $[x,y,k,k]$ where $x,y$ are two spatial dimensions and $k$ is the size of the convolutional kernel. However, matrices are required to compute the SVD for spectral-Atomo, and we choose to reshape each layer into a matrix of size $[xy/2,2k^{2}]$. This provides more flexibility on the sparsity budget for the SVD sparsification. For QSGD, we use the bucketing and Elias recursive coding methods proposed in , with bucket size equal to the number of parameters in each layer of the neural network.

Figure 3 shows how the testing accuracy varies with wall clock time. Tables 3 and 4 give a detailed account of speedups of singular value sparsification compared to QSGD. In these tables, each method is run until a specified accuracy.

We observe that QSGD and Atomo speed up model training significantly and achieve similar accuracy to vanilla mini-batch SGD. We also observe that the best performance is not achieve with the most sparsified, or quantized method, but the optimal method lies somewhere in the middle where enough information is preserved during the sparsification. For instance, 8-bit QSGD converges faster than 4-bit QSGD, and spectral-Atomo with sparsity budget 3, or 4 seems to be the fastest. Higher sparsity can lead to a faster running time, but extreme sparsification can adversely affect convergence. For example, for a fixed number of iterations, 1-bit QSGD has the smallest time cost, but may converge much more slowly to an accurate model.

Conclusion

In this paper, we present and analyze Atomo, a general sparsification method for distributed stochastic gradient based methods. Atomo applies to any atomic decomposition, including the entry-wise and the SVD of a matrix. Atomo generalizes 1-bit QSGD and TernGrad, and provably minimizes the variance of the sparsified gradient subject to a sparsity constraint on the atomic decomposition. We focus on the use Atomo for sparsifying matrices, especially the gradients in neural network training. We show that applying Atomo to the singular values of these matrices can lead to faster training than both vanilla SGD or QSGD, for the same communication budget. We present extensive experiments showing that Atomo can lead to up to a $2\times$ speed-up in training time over QSGD and up to $3\times$ speed-up in training time over TernGrad.

In the future, we plan to explore the use of Atomo with Fourier decompositions, due to its utility and prevalence in signal processing. More generally, we wish to investigate which atomic sets lead to reduced communication costs. We also plan to examine how we can sparsify and compress gradients in a joint fashion to further reduce communication costs. Finally, when sparsifying the SVD of a matrix, we only sparsify the singular values. We also note that it would be interesting to explore jointly sparsification of the SVD and and its singular vectors, which we leave for future work.

Acknowledgement

This work was supported in part by AWS Cloud Credits for Research from Amazon.

References

Appendix A Proof of results

Suppose we have some $p$ satisfying the conditions in (4). We define two auxiliary vectors $\alpha,\beta\in^{n}$ by

Then note that using the fact that $\sum_{i}p_{i}=s$, we have

By the Cauchy-Schwarz inequality, this implies

This proves the first part of Lemma 2. In order to have $f(p)=\frac{1}{s}\|\lambda\|_{1}^{2}$, (9) implies that we need

By the Cauchy-Schwarz inequality, this occurs iff $\alpha$ and $\beta$ are linearly dependent. Therefore, $c\alpha=\beta$ for some constant $c$. Solving, this implies $p_{i}=c|\lambda_{i}|$. Since $\sum_{i=1}^{n}p_{i}=s$, we have

Therefore, $c=\|\lambda\|_{1}/s$, which implies the second part of the theorem. ∎

A.2 Proof of Lemma 4

Fix $q$ that is feasible in (3). To prove Lemma 4 we will require a lemma. Given the atomic decomposition $g=\sum_{i=1}^{n}\lambda_{i}a_{i}$, we say that $\lambda$ is $s$-unbalanced at $i$ if $|\lambda_{i}|s>\|\lambda\|_{1}$, which is equivalent to $g$ being unbalanced in this atomic decomposition at $i$. For notational simplicity, we will assume that $\lambda$ is $s$-unbalanced at $i=1$. Let $A\subseteq\{2,\ldots,n\}$. We define the following notation:

Note that under this notation, Lemma 2 implies that for all $p>0$,

Suppose that $q$ is feasible and that there is some set $A\subseteq\{2,\ldots,n\}$ such that

$\lambda_{A}$ is $(s_{A}+q_{1}-1)$-balanced.

$|\lambda_{1}|(s_{A}+q_{1}-1)>\|\lambda_{A}\|_{1}$.

Then there is a vector $p$ that is feasible satisfying $f(p)\leq f(q)$ and $p_{1}=1$.

Suppose that such a set $A$ exists. Let $B=\{2,\ldots,n\}\backslash A$. Note that we have

Note that by Assumption 1 and Lemma 2, we have

Since $p_{i}=q_{i}$ for $i\in B$, we have $f_{B}(p)=f_{B}(q)$. Therefore,

Combining this with Assumption 2, we have

To show that the RHS of (13) is at most , it suffices to show

However, note that since $0<q_{1}<1$, the RHS of (14) satisfies

Therefore, (14) holds, completing the proof. ∎

We can now prove Lemma 4. In the following, we will refer to Conditions 1 and 2, relative to some set $A$, as the conditions required by Lemma 9.

We first show this in the case that $n=2$. Here we have the atomic decomposition

The condition that $\lambda$ is $s$-unbalanced at $i=1$ implies

In particular, this implies $s>1$. For $A=\{2\}$, Condition 1 is equivalent to

Note that $s_{A}=q_{2}$ and that $q_{1}+q_{2}-2=s-2$ by assumption. Since $q_{i}\leq 1$, we know that $s-2\leq 0$ and so Condition 1 holds. Similarly, Condition 2 becomes

which holds by assumption. Therefore, Lemma 4 holds for $n=2$.

Now suppose that $n>2$, $q$ is some feasible probability vector, and that $\lambda$ is $s$-unbalanced at index $1$. We wish to find an $A$ satisfying Conditions 1 and 2. Consider $B=\{2,\ldots,n\}$. Note that for such $B$, $s_{B}+q_{1}-1=s-1$. By our unbalanced assumption, we know that Condition 2 holds for $B=\{2,\ldots,n\}$. If $\lambda_{B}$ is $(s-1)$-balanced, then Lemma 9 implies that we are done.

Assume that $\lambda_{B}$ is not $(s-1)$-balanced. After relabeling, we can assume it is unbalanced at $i=2$. Let $C=\{3,\ldots,n\}$. Therefore,

Combining this with the $s$-unbalanced assumption at $i=1$, we find

Let $D=\{1,3,4,\ldots,n\}=\{1,\ldots,n\}\backslash\{2\}$. Then note that (16) implies that $\lambda_{D}$ is $(s-q_{2})$-unbalanced at $i=1$. Inductively applying this theorem, this means that we can find a vector $p^{\prime}\in^{|D|}$ such that $p^{\prime}_{1}=1$ and $f_{D}(p^{\prime})\leq f_{D}(q)$. Moreover, $s_{D}(p^{\prime})=s-q_{2}$. Therefore, if we let $p$ be the vector that equals $p^{\prime}$ on $D$ and with $p_{2}=q_{2}$, we have

Appendix B Analysis of Atomo via the KKT Condtions

In this section we show how to derive Algorithm 1 using the KKT conditions. Recall that we wish to solve the following optimization problem:

We first note a few immediate consequences.

If $s>n$ then the problem is infeasible. Note that when $s\geq n$, the optimal thing to do is to set all $p_{i}=1$, in which case no sparsification takes place.

If $\lambda_{i}=0$, then $p_{i}=0$. This follows from the fact that this $p_{i}$ does not change the value of $f(p)$, and the objective could be decreased by allocating more to the $p_{j}$ associated to non-zero $\lambda_{j}$. Therefore we can assume that all $\lambda_{i}\neq 0$.

If $|\lambda_{i}|\geq|\lambda_{j}|>0$, then we can assume $p_{i}\geq p_{j}$. Otherwise, suppose $p_{j}>p_{i}$ but $|\lambda_{i}|\geq|\lambda_{j}|$. Let $p^{\prime}$ denote the vector with $p_{i},p_{j}$ switched. We then have

We therefore assume $0<s\leq n$ and $|\lambda_{1}|\geq|\lambda_{2}|\geq\ldots\geq|\lambda_{n}|>0$. As above we define $\lambda:=[\lambda_{1},\ldots,\lambda_{n}]^{T}$. While the formulation of (17) does not allow direct application of the KKT conditions, since we have a strict inequality of $0<p_{i}$, this is fixed with the following lemma.

The minimum of (17) is achieved by some $p^{*}$ satisfying

Define $\overline{p}$ by $\overline{p}_{i}=s/n$. This vector is clearly feasible in (17). Let $p$ be any feasible vector. If $f(p)\leq f(q)$ then for any $i\in[n]$ we have

Therefore, $p_{i}\geq\lambda_{i}^{2}/f(\overline{p})$. A straightforward computations shows that $f(\overline{p})=n\|\lambda\|_{2}^{2}/s$. Note that this implies that we can restrict to the feasbile set

This defines a compact region $C$. Since $f$ is continuous on this set, its maximum value is obtained at some $p^{*}$.∎

The KKT conditions then imply that at any point $p$ solving (17), we must have

for some $\mu\in$. Since $|\lambda_{i}|>0$ for all $i$, we actually must have $\mu>0$. We therefore have two conditions for all $i$.

$p_{i}<1\implies p_{i}=|\lambda_{i}|/\sqrt{\mu}$.

Note that in either case, to have $p_{1}$ feasible we must have $\mu\geq\lambda_{1}^{2}$. Combining this with the fact that we can always select $p_{1}\geq p_{2}\geq\ldots\geq p_{n}$, we obtain the following partial characterization of the solution to (17). For some $n_{s}\in[n]$, we have $p_{1},\ldots,p_{n_{s}}=1$ while $p_{i}=|\lambda_{i}|/\sqrt{\mu}\in(0,1)$ for $i=n_{s}+1,\ldots,n$. Combining this with the constraint that $\sum_{i=1}^{n}p_{i}=s$, we have

Thus, we need to select $n_{s}$ such that the $p_{i}$ in (21) are bounded above by 1. Let $n_{s}^{*}$ denote the first element of $[n]$ for which this holds. Then the condition that $p_{i}\leq 1$ for $i=n_{s}^{*}+1,\ldots,n$ is exactly the condition that $[\lambda_{n_{s}^{*}+1},\ldots,\lambda_{n}]$ is $(s-n_{s})$-balanced (see Definition 1. In particular, Lemma 2 implies that, fixing $p_{i}=1$ for $i=1,\ldots,n_{s}^{*}$, the optimal way to assign the remaining $p_{i}$ is by

This agrees with (21) for $n_{s}=n_{s}^{*}$. In particular, the minimal value of $f$ occurs at the first value of $n_{s}$ such that the $p_{i}$ in (21) are bounded above by 1.

Algorithm 1 scans through the sorted $\lambda_{i}$ and finds the first value of $n_{s}$ for which the probabilities in (21) are in $$, and therefore finds the optimal$p$for (17). The runtime is dominated by the$O(n\log n)$sorting cost. It is worth noting that we could perform the algorithm in$O(sn)$time as well. Instead of sorting and then iterating through the$\lambda_{i}$in order, at each step we could simply select the next largest$|\lambda_{i}|$not yet seen and perform an analogous test and update as in the above algorithm. Since we would have to do the selection step at most$s$times, this leads to an$O(sn)$ complexity algorithm.

Appendix C Equivalence of norms

We are often interested in comparing norms on vectors spaces. This naturally leads to the following definition.

As it turns out, norms on finite-dimensional vector spaces are always equivalent.

In order to compare norms, we often wish to determine the tightest constants which give equivalence between them. In Section 5, we are particularly interested in comparing the $\|X\|_{*}$ and $\|X\|_{1,1}$ on the space of $n\times m$ matrices. We have the following lemma.

Suppose that $X$ has the singular value decomposition

We will first show the left inequality. First, note that for any $n\times m$ matrix $A$, $\|A\|_{1,1}\leq\sqrt{nm}\|A\|_{F}$. This follows directly from the fact that for a $n$-dimensional vector $v$, $\|v\|_{1}\leq\sqrt{n}\|v\|_{2}$. We will also use the fact that for any vectors $u\in^{n},v\in^{m}$, $\|uv^{T}\|_{F}=\|u\|_{2}\|v\|_{2}$. We then have

For the right inequality, note that we have

where $e_{i}\in^{n}$ is the $i$-th standard basis vector, while $e_{j}\in^{m}$ is the $j$-th standard basis vector. We then have

In fact, these are the best constants possible. To see this, first consider the matrix $X$ with a 1 in the upper-left entry and 0 elsewhere. Clearly, $\|X\|_{*}=\|X\|_{1,1}=1$, so the right-hand inequality is tight. For the left-hand inequality, consider the all-ones matrix $X$. This has one singular value, $\sqrt{nm}$, so $\|X\|_{*}=\sqrt{nm}$. On the other hand, $\|X\|_{1,1}=nm$. Therefore, $\|X\|_{1,1}=\sqrt{nm}\|X\|_{*}$ in this case.

Appendix D Hyperparameter optimization

We firstly provide results of step size tunning, as shows in Table 5 we reported stepsize tunning results for all of our experiments. We tuned these step sizes by evaluating many logarithmically spaced step sizes (e.g., $2^{-10},\ldots,2^{0}$) and evaluated on validation loss.

This step sizes tuning, for 8 gradient coding methods and 3 datasets was only possible because fairly small networks were used.

Appendix E Additional Experiments

Runtime analysis: We empirically study runtime costs of spectral-Atomo with sparsity budget set at 1, 2, 3, 6 and made comparisons among $b$-bit QSGD and TernGrad. We deployed distributed training on ResNet-18 with batch size $B=256$ on the CIFAR-10 dataset run with m5.2xlarge instances. As shown in Figure 4, there is a trade-off between the amount of communication per iteration and the running time for both singular value sparsification and QSGD. In some scenarios, spectral-Atomo attains a higher compression ratio than QSGD and TernGrad. For example, singular value sparsification with sparsity budget 1 may communicate smaller messages than $\{2,4\}$-bit QSGD and Terngrad.