ATOMO: Communication-efficient Learning via Atomic Sparsification

Hongyi Wang; Scott Sievert; Zachary Charles; Shengchao Liu; Stephen Wright; Dimitris Papailiopoulos

Source-linked AI summary

ATOMO: Communication-efficient Learning via Atomic Sparsification

Hongyi Wang, Scott Sievert, Zachary Charles, Shengchao Liu, Stephen Wright, Dimitris Papailiopoulos

arXiv:1806.04090v3Jun 11, 2018stat.MLcs.DCcs.LG

TL;DR

Distributed training is limited by the communication cost of frequent, increasingly large gradient updates. ATOMO formulates unbiased atomic sparsification as variance minimization under a sparsity budget, encompassing coordinate, singular-value, and other decompositions. Its SVD-based variant reports faster training than coordinate-based baselines, including up to 2× over QSGD and up to 3× over TernGrad.

Problem
Communication overhead from frequent gradient updates limits distributed-training speedups, especially as model gradients and federated uplink demands grow.
Method
ATOMO optimally chooses random atom-retention probabilities to produce an unbiased sparse gradient with minimum variance under a sparsity constraint.
Results
SVD-based ATOMO achieves up to 2× faster training than QSGD and up to 3× faster training than TernGrad at comparable communication budgets.
Takeaways & Limitations
Choosing an atomic decomposition aligned with gradient structure can make communication-efficient training faster than entry-wise sparsification.
Takeaways & Limitations
The paper leaves joint sparsification of singular values and singular vectors, as well as Fourier-based and joint compression methods, for future work.

Abstract

from arXiv · show

Distributed model training suffers from communication overheads due to frequent gradient updates transmitted between compute nodes. To mitigate these overheads, several studies propose the use of sparsified stochastic gradients. We argue that these are facets of a general sparsification method that can operate on any possible atomic decomposition. Notable examples include element-wise, singular value, and Fourier decompositions. We present ATOMO, a general framework for atomic sparsification of stochastic gradients. Given a gradient, an atomic decomposition, and a sparsity budget, ATOMO gives a random unbiased sparsification of the atoms minimizing variance. We show that recent methods such as QSGD and TernGrad are special cases of ATOMO and that sparsifiying the singular value decomposition of neural networks gradients, rather than their coordinates, can lead to significantly faster distributed training.

1 Introduction

Distributed training often saturates because communicating increasingly large, frequent gradients creates a bottleneck. ATOMO generalizes gradient sparsification across atomic decompositions and reports faster SVD-based training than coordinate-based methods.

1 Introduction: Communication overheads limit distributed-training speedups because frequent gradient updates grow with model size.The bottleneck is especially pronounced for federated learning over low-bandwidth uplinks.
1 Introduction: ATOMO treats stochastic-gradient sparsification and quantization as one framework applicable to entry-wise, singular-value, Fourier, and other atomic decompositions.Its estimator is unbiased, sparse on the atomic basis, and variance-minimizing under a sparsity constraint.
1 Introduction: 1-bit QSGD and TernGrad are special cases of ATOMO that are optimal in different variance–sparsity regimes.Both methods use the entry-wise atomic decomposition.
1 Introduction: Up to 2× faster training is reported for SVD-based ATOMO than QSGD on VGG and ResNet-18 using SVHN and CIFAR-10.The comparison includes the time required to compute the SVD and uses the same communication budget.
1 Introduction: Gradient singular values can decay sharply, with the top 3 standing out in the illustrated ResNet-18 CIFAR-10 convolutional layer.This motivates treating layer gradients as nearly low-rank matrices.

2 Problem Setup

The problem setup asks how to construct unbiased stochastic gradients with minimal variance under a communication or sparsity budget. Small second moments matter because they control expected optimization progress and convergence bounds.

2 Problem Setup: Mini-batch SGD parallelizes gradient computation across compute nodes, while a parameter server aggregates updates and communicates the model.Each node computes gradients on randomly sampled training data.
2 Problem Setup: Unbiased stochastic gradients with smaller second moments provide greater expected single-step progress toward optimization.Convergence analyses in convex and nonconvex settings commonly bound this second moment.
2 Problem Setup: The central question is the smallest variance achievable by an unbiased stochastic gradient represented with k bits.The setup connects variance minimization to communication-efficient stochastic-gradient design.
2 Problem Setup: Because an exact formal communication constraint is difficult to impose, the framework replaces it with sparsity relative to a chosen atomic decomposition.The sparsity constraint is imposed on the expected number of selected atoms.

3 Atomo: Atomic Decomposition and Sparsification

ATOMO represents gradients as combinations of normalized atoms and randomly retains atoms to form an unbiased sparse estimator. It chooses retention probabilities that minimize variance subject to an average sparsity budget.

3 Atomo: Atomic Decomposition and Sparsification: An atomic decomposition expresses a gradient as a linear combination of simple unit-norm building blocks.Entry-wise, orthonormal-basis, Fourier, wavelet, and singular-value decompositions fit this framework.
3 Atomo: Atomic Decomposition and Sparsification: Sparsifying singular values can reduce communication when a matrix is low rank because only nonzero atomic weights need to be transmitted.For an r-rank m×n matrix, communicating SVD factors may require r(n + m) entries instead of nm entries.
3 Atomo: Atomic Decomposition and Sparsification: Atomic sparsification retains each atom independently with a Bernoulli probability and rescales retained weights to preserve unbiasedness.The estimator is designed to have small variance while selecting at most an average sparsity budget.
3 Atomo: Atomic Decomposition and Sparsification: The optimal probabilities minimize expected squared norm, equivalently variance for an unbiased estimator, under the constraint that their sum equals s.The solution handles both s-balanced and s-unbalanced gradients, setting the largest unbalanced atom’s probability to 1.
3 Atomo: Atomic Decomposition and Sparsification: Computing the probability vector by sorting atom weights requires O(n log n) operations, while a selection-based implementation can require O(sn).The preferable implementation depends on the relative sizes of s and log n.

4 Relation to QSGD and TernGrad

ATOMO unifies entry-wise quantization methods by expressing their coordinate probabilities as atomic-sparsification probabilities. The ℓq family exposes a variance–sparsity trade-off that includes 1-bit QSGD and TernGrad.

4.1 1-bit QSGD: 1-bit QSGD is variance-optimal entry-wise sparsification for sparsity budget s = ∥g∥1/∥g∥2.Its coordinate probabilities are pi = |gi|/∥g∥2.
4.2 TernGrad: TernGrad is variance-optimal entry-wise sparsification for sparsity budget s = ∥g∥1/∥g∥∞.Its coordinate probabilities are pi = |gi|/∥g∥∞.
4.3 ℓq-quantization: ℓq-quantization generalizes both methods through probabilities pi = |gi|/∥g∥q and sparsity budget s = ∥g∥1/∥g∥q.The standard basis is used for this quantization family.
4.3 ℓq-quantization: As q increases, ℓq-quantization increases the sparsity budget while decreasing variance.Choosing q = 2 yields 1-bit QSGD, while q = ∞ yields TernGrad.

5 Spectral-Atomo: Sparsifying the Singular Value Decomposition

Spectral-Atomo applies Atomo to a matrix’s singular value decomposition and compares it with entry-wise sparsification under matched communication costs. The comparison uses second moment and matrix-norm conditions to characterize when singular-value sparsification is preferable.

Matrix sparsification methods: Spectral-Atomo applies Atomo to the singular value decomposition, while the alternative applies it to the entry-wise decomposition.These are the two standard atomic decompositions compared for matrix sparsification.
Communication cost: When the rank is small, communicating the SVD can be more efficient because it uses r(n + m) entries instead of nm matrix entries.The comparison denotes the resulting random sparsifications by b X and b Xσ.
Communication cost: The methods are compared using communication cost and the optimal second moment under a sparsity budget.For SVD sparsification, a budget of s corresponds to communication cost s(n + m).
Variance comparison: With matched communication cost, singular-value sparsification has no greater second moment than entry-wise sparsification exactly under the stated theorem condition.The matching uses sparsity budget s′ = s/(n + m) for the singular-value method.
Variance comparison: The comparison connects the variance condition to equivalence bounds between the ℓ1,1 and spectral norms.The bounds are tight in both directions, so either sparsification method can have the lower second moment depending on the matrix.

6 Experiments

The experiments evaluate Spectral-Atomo against QSGD, TernGrad, and vanilla SGD across distributed models and datasets, measuring timing, convergence, and speedups. Spectral-Atomo reduces communication and can substantially accelerate training, although intermediate sparsity or quantization levels perform best.

Experimental setup: The study compares Spectral-Atomo with QSGD and TernGrad in real distributed environments across multiple models and datasets.Experiments use distributed clusters with a parameter server and compute nodes, including ResNet-18, ResNet-34, and VGG11-BN.
Speed and communication: Spectral-Atomo reduces communication relative to vanilla mini-batch SGD and can reduce training time versus QSGD and TernGrad by up to 2× and 3×, respectively.On VGG11-BN trained on CIFAR-10, sparsity budget 3 achieved 3.96× speedup over vanilla SGD, compared with 1.68× for 4-bit QSGD.
Scalability: Increasing cluster size lowers per-worker computation cost but increases communication overhead, creating a trade-off between retained information and communication time.Figure 2 separates total iteration time into computation, encoding, and communication components for QSGD and Spectral-Atomo.
Convergence: Both QSGD and Atomo significantly speed up model training while achieving accuracy similar to vanilla mini-batch SGD.Convergence is evaluated using testing accuracy versus wall-clock time and speedups to specified accuracy targets.
Compression trade-offs: The fastest configuration is typically intermediate: Spectral-Atomo with sparsity budget 3 or 4 and 8-bit QSGD outperform more aggressively compressed settings.Extreme sparsification can reduce iteration time but adversely affect convergence; 1-bit QSGD may take much longer to reach an accurate model.

7 Conclusion

The paper concludes that Atomo is a general, variance-optimal framework for sparsifying stochastic gradients under atomic-decomposition constraints. Its SVD-based application can accelerate neural-network training, while joint sparsification of singular values and vectors remains future work.

7 Conclusion: Atomo applies to atomic decompositions including entry-wise and matrix SVD decompositions, and minimizes variance subject to a sparsity constraint.The framework generalizes 1-bit QSGD and TernGrad.
7 Conclusion: Sparsifying singular values of neural-network gradient matrices can produce faster training than vanilla SGD or QSGD at the same communication budget.The paper reports up to 2× speedup over QSGD and up to 3× over TernGrad.
7 Conclusion: Future work includes Fourier decompositions, joint gradient sparsification and compression, and jointly sparsifying SVD singular values and singular vectors.The current SVD procedure sparsifies only singular values.

A Proof of results

The proof establishes the variance-minimizing sampling structure for atomic sparsification using auxiliary vectors and the Cauchy-Schwarz inequality. Equality requires the sampling probabilities to be proportional to the absolute atomic coefficients.

A. Proof of results: The proof introduces auxiliary vectors α and β to analyze the feasible sampling-probability vector p.The argument assumes p satisfies the conditions in equation (4).
A. Proof of results: The Cauchy-Schwarz inequality provides the key bound used to establish the first part of Lemma 2.The proof then examines when equality can hold.
A. Proof of results: Equality holds exactly when α and β are linearly dependent, implying p_i = c|λ_i| for some constant c.This identifies the probability allocation proportional to the absolute atomic coefficients.

A.2 Proof of Lemma 4

The proof establishes Lemma 4 by constructing feasible probability vectors that preserve or reduce the objective under unbalanced atomic coefficients. It handles the two-atom case directly and extends the argument inductively to larger decompositions.

Lemma 9 supplies a feasible vector with p1 = 1 and f(p) ≤ f(q) whenever an appropriate subset satisfies the required conditions.
The base case and induction together complete the proof of Lemma 4.
For n = 2, the stated conditions hold under the assumed unbalancedness relations, proving Lemma 4 in the base case.
For n > 2, the proof considers the set of all indices except the unbalanced index and applies the lemma when the remaining coefficients are balanced.
When the remaining coefficients are unbalanced, relabeling and recursively removing indices yields a reduced feasible vector with no larger objective.

B Analysis of Atomo via the KKT Condtions

The KKT analysis characterizes the optimizer of Atomo's probability-allocation problem. After sorting coefficient magnitudes, the solution assigns probability one to the largest coefficients and proportional probabilities to the remainder.

The analysis restricts attention to feasible instances with 0 < s ≤ n, nonzero coefficients, and coefficients ordered by decreasing magnitude.
KKT conditions imply that an optimal solution has an initial block with pi = 1 and a remaining block with pi = |λi|/√µ.
The remaining probabilities are determined by the sparsity constraint after selecting the size ns of the initial probability-one block.
The optimum occurs at the first ns for which all remaining probabilities are bounded by one, equivalently when the remaining coefficients satisfy the required balance condition.
Algorithm 1 finds this optimizer by sorting the coefficients and scanning candidate values, with runtime dominated by O(n log n) sorting.
For matrix norms, the analysis compares entry-wise and singular-value decompositions, including tight bounds between ∥X∥1,1 and the nuclear norm ∥X∥∗.

D Hyperparameter optimization

The experiments tune learning rates by evaluating logarithmically spaced candidates on validation loss. This extensive tuning was feasible because the experiments used relatively small networks.

Step sizes were tuned by evaluating many logarithmically spaced candidates, such as 2^-10 through 2^0, against validation loss.
The tuning covered 8 gradient coding methods and 3 datasets.
The tuning workload was practical only because fairly small networks were used.

E Additional Experiments

Additional experiments compare spectral-Atomo with QSGD and TernGrad on runtime, communication, and convergence. The results examine ResNet-18 and VGG-11-BN across CIFAR-10 and SVHN settings.

Runtime experiments compare singular-value sparsification, QSGD, and TernGrad on ResNet-18 trained on CIFAR-10.The analysis reports computation, encoding, communication time, and gradient-message size.
Spectral-Atomo exhibits a trade-off between communication per iteration and running time as sparsity changes.
Convergence comparisons use the best QSGD and SVD performances across CIFAR-10 with ResNet-18, SVHN with ResNet-18, and CIFAR-10 with VGG-11-BN.
Table 6 reports speedups of spectral-Atomo with sparsity budget s, b-bit QSGD, and TernGrad using VGG11 on CIFAR-10 over vanilla SGD.