Source-linked AI summary
Distributed Learning with Compressed Gradient Differences
Konstantin Mishchenko, Eduard Gorbunov, Martin Takáč, Peter Richtárik
TL;DR
Communication bottlenecks motivate compressed distributed learning, but prior compressed-SGD work left questions about block quantization and worker data availability. DIANA is a distributed compression method, and the paper reports asymptotic expected linear convergence to the exact optimum for strongly convex problems with heterogeneous local functions.
Problem
Prior compressed-SGD work left open questions about block quantization and assumed all data were available to each worker, despite communication being a bottleneck.
Source: §1 · p. 2Method
DIANA is a distributed optimization method with compression whose iterations use unbiased local-gradient estimators assumed to have bounded variance.
Source: §3 · p. 4, §7 · p. 12Results
DIANA converges linearly in expectation to the exact optimum asymptotically for strongly convex problems with heterogeneous local functions and full local gradients.
Source: §7 · p. 12Takeaways & Limitations
For large vectors with differently scaled coordinates, blockwise quantization can prevent undersampling smaller values, and theory predicts superlinear acceleration from quantizing blocks or layers.
Source: §6 · p. 10DIANA's stated setup assumes unbiased local-gradient estimators with bounded variance.
Source: §3 · p. 4
Abstract
from arXiv · showhide
Training large machine learning models requires a distributed computing approach, with communication of the model updates being the bottleneck. For this reason, several methods based on the compression (e.g., sparsification and/or quantization) of updates were recently proposed, including QSGD (Alistarh et al., 2017), TernGrad (Wen et al., 2017), SignSGD (Bernstein et al., 2018), and DQGD (Khirirat et al., 2018). However, none of these methods are able to learn the gradients, which renders them incapable of converging to the true optimum in the batch mode. In this work we propose a new distributed learning method -- DIANA -- which resolves this issue via compression of gradient differences. We perform a theoretical analysis in the strongly convex and nonconvex settings and show that our rates are superior to existing rates. We also provide theory to support non-smooth regularizers study the difference between quantization schemes. Our analysis of block-quantization and differences between $\ell_2$ and $\ell_{\infty}$ quantization closes the gaps in theory and practice. Finally, by applying our analysis technique to TernGrad, we establish the first convergence rate for this method.
1. Introduction
The introduction frames distributed machine learning as a scaling and communication challenge, reviews approaches to reduce communication, and states the regularized learning problem and its assumptions.
Communication as the bottleneck
- Distributed training aggregates updates computed by workers, but communication can be slower than stochastic updates and prevent systems from scaling linearly as resources grow.The introduction motivates methods that balance computation and communication while remaining versatile and efficient.Source: §1 · pp. 1–2
Motivation
- The desired methods must address federated learning’s communication and local-data concerns, proximable nonsmooth regularization, and deep learning’s momentum and nonconvex convergence needs.The introduction also identifies sparse solutions as a possible requirement in distributed optimization.Source: §1 · p. 2
Communication as the bottleneck
- Prior communication strategies include block descent in Hydra and extra local computation in CoCoA variants [15], DANE, DiSCO [18], DANCE [9], and AIDE [22].Hydra’s block size controls the amount of data communicated; local solvers can instead balance computation with communication between rounds.Source: §1 · p. 2
Update compression via randomized sparsification and/or quantization
- Compressed-SGD research includes QSGD and related methods [10] [12], but QSGD left questions about block quantization and assumed each worker could access all data.Concurrent Terngrad analysis considered stochastic updates positively correlated in expectation with the solution direction and used ℓ∞ quantization, unlike QSGD’s ℓ2 quantization.Source: §1 · p. 2
The problem
- The learning problem uses each worker’s expected loss over its data distribution and a proper closed convex regularizer, without assuming similarity among worker distributions.The local loss is defined as fi(x) := Eζ∼Diϕ(x, ζ).Source: §1 · p. 2
2. Contributions
DIANA compresses gradient differences using learned vectors that converge to each worker’s gradient at the optimum. The paper analyzes its rates for strongly convex and nonconvex problems and develops convergence theory for TernGrad and 1-bit QSGD.
DIANA
- DIANA compresses gradient differences rather than gradients, using learned vectors h_i that converge to each worker’s optimum gradient; it also adds a momentum modification.The method uses a special update rule for h_i to ensure convergence to ∇f_i(x*).Source: §2 · p. 3
Rates in the strongly convex and nonconvex cases
- For smooth strongly convex objectives with arbitrary closed convex regularizers, DIANA has an iteration-complexity guarantee for reaching an optimum-centered neighborhood; for smooth nonconvex objectives without regularization, it bounds expected squared gradient norm.The strongly convex neighborhood’s size depends on the stepsize and variance.Source: §2 · p. 3
New theory for Terngrad and QSGD
- The paper derives TernGrad’s first convergence rate and a tighter analysis of 1-bit QSGD under less restrictive assumptions for specified strongly convex and nonconvex settings.It also studies momentum versions for smooth nonconvex objectives with a constant regularizer when all worker objectives are identical.Source: §2 · p. 3
3. The Algorithm
DIANA combines unbiased stochastic gradients with stored local estimates, communicating quantized gradient differences rather than raw gradients and taking a proximal momentum step. The section defines p-quantization and describes its communication-cost and convergence-analysis scope.
The Algorithm
- Each DIANA iteration uses an unbiased estimator of each node’s local gradient, with bounded stochastic-gradient variance assumed.The stochastic-gradient assumption is used throughout the method description.Source: §3 · p. 4
Quantization
- p-quantization randomly encodes gradient differences using signed, coordinatewise Bernoulli sampling scaled by the difference’s p-norm.A coordinate is sampled with probability equal to its absolute value divided by the p-norm; the resulting entries lie in {−t, 0, t}.Source: §3 · pp. 4–5
Communication cost
- Elias coding communicates quantized vectors, and Theorem 3.3 bounds expected communication cost using the quantized vector’s sparsity.The section gives an upper bound on expected communication cost; its communication expressions increase with p.Source: §3 · p. 5
Proximal step
- For a closed convex regularizer R, DIANA’s proximal step uses prox_{γR}, a nonexpansive operator.The proximal operator is defined for γ > 0.Source: §3 · p. 5
DIANA
- DIANA stores local memories h_i, quantizes each stochastic-gradient difference, and forms an unbiased full-gradient estimator for a proximal momentum update.The update direction is v_k = βv_{k−1} + ĝ_k, with 0 ≤ β ≤ 1; block quantization is also allowed.Source: §3 · p. 5
Relation to QSGD and TernGrad
- With α = 0 and h0 = 0, DIANA reduces to 1-bit QSGD for p = 2 or TernGrad for p = ∞; its iteration bounds improve with p.The bounds cover smooth strongly convex objectives with closed convex regularizers, smooth nonconvex objectives with constant regularizers, and the momentum version; the best iteration complexity is for ℓ∞ quantization. The analysis establishes a TernGrad rate where the original work provided none, and extends beyond the bounded-gradient assumption used for 1-bit QSGD by assuming bounded stochastic-gradient variance.Source: §3 · pp. 5–6
4. Theory: Strongly Convex Case
In the strongly convex case, DIANA has a geometric accuracy guarantee under the theorem’s conditions and its compressed gradient differences eventually learn the gradient at the optimum.
4. Theory: Strongly Convex Case
- Lemma 4.3 states that α_p increases with p and decreases with dimension d, with corresponding relations across dimensions for α_1 and α_2.Source: §4 · p. 6
- The Lyapunov function tracks both squared distance to the solution and accumulated estimation errors, and both terms approach zero during the algorithm.Source: §4 · p. 7
- Under L-smoothness and strong convexity, Theorem 4.4 gives a geometric accuracy guarantee for the Lyapunov function with admissible stepsizes, block sizes, and parameter c.The theorem’s stated accuracy condition bounds the expected Lyapunov value after a number of iterations depending on the initial value and target accuracy.Source: §4 · p. 7
- As the Lyapunov function vanishes, each worker’s compressed gradient-difference memory converges to the gradient at the optimum, distinguishing DIANA from prior literature.Source: §4 · p. 7
- Choosing γ at the minimum in (15) gives leading iteration-complexity term 1/(γμ) = max{1/α, (μ+L)(1+cα)/(2μ)}.Source: §4 · p. 7
Corollary
The corollary identifies p = +∞ as optimal for the leading strongly convex complexity term, shows when DIANA matches gradient descent, and gives block-quantization and decreasing-stepsize results.
4.5. Let
- The leading term 1/γµ = max {2/αp, (κ + 1) (1/2 − 1/n + 1/nαp)} decreases with p, making p = +∞ optimal.The stated conditions (14) and (15) are satisfied.Source: §4.5 · p. 7
- When d = O(min{κ^2, n^2}), DIANA with p ≥ 2 achieves O(κ log(1/ϵ)), matching gradient descent.Source: §4.5 · p. 8
- Quantization increases iteration complexity by about d/n, but quantizing blocks of size n^2 reduces this increase to 1.With about 100 workers and blocks of size about 10,000, the algorithm works as fast as SGD while communicating bits instead of floats; block quantization also avoids zeroing updates in smaller-scale layers.Source: §4.5 · p. 8
4.1. Decreasing stepsizes
- With decreasing step sizes, DIANA obtains an O(1/k) convergence rate under smoothness, strong convexity, and bounded gradient noise.Source: §4.1 · p. 8
5. Theory: Nonconvex Case
The nonconvex analysis relies on bounded data dissimilarity because the solution may be nonunique. Under additional conditions, it gives an O(1/K) complexity matching SGD.
5. Theory: Nonconvex Case
- Bounded data dissimilarity is necessary in the nonconvex case because nonuniqueness prevents bounding gradient differences using the gradient at the optimum.Assumption 5.1 posits a constant ζ ≥ 0 for the dissimilarity bound.Source: §5 · p. 9
- ζ = 0 when all workers’ functions differ only by additive constants; a more general bound adds an O(∥∇f(x)∥²) term but is omitted for complexity.The paper says this extension would overcomplicate the theory without providing more insight.Source: §5 · p. 9
- Theorem 5.2 assumes constant R, an L-smooth objective, positive constant stepsizes and parameter, and random selection of an iterate from the preceding iterates.The theorem also imposes additional constraints on the stepsizes and parameter.Source: §5 · p. 9
- O(1/K) complexity, matching SGD, results when the corollary’s first term leads and 1/n = Ω(αp); sufficiently large minibatches and blocks of size O(n²) support these conditions.The passage identifies sufficiently large minibatches for the first condition and quantization blocks of size O(n²) for the latter.Source: §5 · p. 9
6. Implementation and Experiments
The experiments examine blockwise quantization and compare DIANA with compressed-gradient methods and MPI reduction; on CIFAR-10, DIANA and SGD significantly outperform the other tested methods.
Block quantization
- Following, the authors recommend quantizing large vectors blockwise to avoid undersampling small coordinates; their theory predicts superlinear acceleration from block- or layer-wise quantization.This is motivated by coordinates having different scales, as in deep learning.Source: §6 · p. 10
Convex experiments
- In convex experiments, optimal α was usually around min_i 1/√d_i, with the minimum taken over blocks and d_i denoting block sizes.Source: §6 · p. 10
Logistic regression
- For logistic regression, DIANA with momentum works best, and the optimal block size for ℓ∞ quantization is significantly larger than for ℓ2.Source: §6 · p. 10
DIANA vs MPI
- With 2-bit DIANA versus 32-bit MPI reduction on a Cray Aries network, unusual scaling affected speed beyond 64 processes for AlexNet and 32 for VGG, as also documented in.Figure 2 compares images per second across GPU/MPI-process counts.Source: §6 · p. 10, Figure 2 · p. 11
Train and test accuracy on Cifar10
- DIANA and SGD significantly outperform QSGD and TernGrad on CIFAR-10 in the best runs across tested hyperparameters.Figure 3 shows training and test accuracy after three epochs with eight workers.Source: §6 · pp. 10–11, Figure 3 · p. 11
7. Discussion … Appendix C. Proof of Theorem 3.3
The discussion highlights DIANA’s asymptotic convergence to the exact optimum for strongly convex problems with heterogeneous workers using full local gradients; the appendices provide supporting identities, block quantization results, and a proof of Theorem 3.3.
7. Discussion
- DIANA converges asymptotically in expectation to the exact optimum on strongly convex problems with full local gradients, even when workers’ functions are arbitrarily heterogeneous.The discussion notes extensions to unbiased compression and variance reduction [7], accelerated methods [14], and decentralized methods.Source: §7 · p. 12
Appendix A. Basic Identities and Inequalities
- Appendix A collects basic tools: smoothness and strong-convexity relations, a norm identity for convex combinations, and variance decomposition.Source: Appendix A · p. 14
Appendix B. Block p-quantization
- Appendix B defines block p-quantization, allows dependent quantization across blocks, and recovers full quantization when there is one block.The blocks partition the initial vector into components of sizes d_l.Source: Appendix B · p. 14
- The block p-quantization estimator is unbiased, and Lemma B.2 gives a formula for its variance.The variance statement is expressed using Ψ_l(x) = ||x^(l)||_1||x^(l)||_p − ||x^(l)||_2.Source: Appendix B · pp. 14–15
Appendix C. Proof of Theorem 3.3
- The proof represents each quantized coordinate as ∥∆∥p sign(∆(j))ξ(j), where ξ(j) is Bernoulli with parameter |∆(j)|/∥∆∥p, and establishes the first claim using a norm bound at q = 1.The proof of the second claim uses concavity and Jensen’s inequality in one step.Source: Appendix C · pp. 15–16
Appendix D. Proof of Lemma 4.3 … Appendix I. Strongly convex case: decreasing stepsize
These appendices establish quantization and convergence results, characterize how dimension affects quantization and node choices, and analyze strongly convex convergence with decreasing stepsizes.
Appendix D. Proof of Lemma 4.3
- The proof establishes that α_p(d) increases with p and decreases with dimension d.Source: Appendix D · p. 16
Appendix E. Strongly Convex Case: Optimal Number of Nodes
- The optimal number of compute nodes n(d) increases with dimension d, supporting more nodes for larger models.Source: Appendix E · p. 17
Appendix F. Quantization Lemmas
- The quantization lemmas establish identities for DIANA iterations and bounds on the first two moments of the updated memory variables.Source: Appendix F · pp. 18–19
Appendix G. Proof of Theorem 4.4
- The proof of Theorem 4.4 derives and unrolls a Lyapunov-function recursion to establish its convergence bound.Source: Appendix G · p. 24
Appendix H. Proof of Corollary 4.5
- For iteration complexity, p = +∞ is optimal because the leading term decreases with p.Source: Appendix H · p. 24
Appendix I. Strongly convex case: decreasing stepsize
- The decreasing-stepsize analysis treats three regimes and includes a regime with θ = Θ(L) for achieving E[V_k] ≤ ε.The theorem assumes smooth, strongly convex f with bounded gradient noise and uses γ_k = 2/(μk + θ).Source: Appendix I · p. 26
Appendix J. Nonconvex analysis … L.4. Momentum version
The appendices establish nonconvex convergence guarantees for DIANA and its momentum variant, then analyze TernGrad as a DIANA special case and derive a convergence rate for it. They also provide technical lemmas and a momentum guarantee for the α = 0 setting.
Appendix J. Nonconvex analysis
- With constant R, smooth f, and the stated stepsize conditions, Appendix J derives a nonconvex guarantee for the expected gradient norm.The result samples x_k uniformly from the iterates; the supplied excerpts omit the complete accuracy expression.Source: Appendix J · pp. 27, 30
Appendix K. Momentum version of DIANA
- The proof introduces virtual iterates whose recursion is used to establish convergence by writing a recursion on f(z_k), and samples x_k uniformly from {x_0, ..., x_{k−1}}.Source: Appendix K · p. 30
L.1. Convergence Rate of TernGrad
- The appendix closes the missing convergence-rate analysis for TernGrad, generalizing beyond the original ℓ∞ sampling choice.It identifies TernGrad with DIANA when h_i^0 = 0, α = 0, and p = ∞.Source: §L.1 · p. 37
Appendix L. Analysis of DIANA with α = 0 and h0
- The Appendix L analysis treats TernGrad as DIANA with zero initial memories and α = 0, making the DIANA inequalities applicable to its iterates.The original TernGrad choice p = ∞ is a special case; the analysis also permits other sampling choices.Source: §L.1 · p. 37
L.2. Technical lemmas
- For constant R and constant stepsize, Lemma L.4 supplies an iterate bound under smoothness, with σ² defined as the average component variance.This technical result supports the subsequent nonconvex analysis.Source: §L.2 · pp. 38–39
L.3. Nonconvex analysis
- Theorem L.5 assumes constant R, Assumption 5.1, L-smooth f, and γ ≤ nαp L((n−1)αp+1), with xk chosen randomly from {x0, . . . , xk−1}.Source: §L.3 · p. 39
L.4. Momentum version
- For α = 0 and h_i = 0, the momentum variant provides an accuracy guarantee after k iterations under its stated parameter constraints.Its proof uses virtual iterates z_k with recursion z_{k+1} = z_k − γ/(1−β) ĝ_k to analyze f(z_k).Source: §L.4 · pp. 41, 43
L.5. Strongly convex analysis … M.4. Performance of GPU
The appendices analyze strongly convex constant- and decreasing-stepsize regimes and report experiments on quantization, communication, and GPU batch sizes; on Rosenbrock, DIANA vastly outperforms QSGD and TernGrad.
L.5. Strongly convex analysis
- Theorem L.9 analyzes TernGrad for µ-strongly convex, L-smooth functions using a constant stepsize.The theorem's displayed conclusion and stepsize condition are incomplete in the supplied passage.Source: §L.5 · pp. 43–44
L.6. Decreasing stepsize
- Theorem L.10 considers decreasing stepsizes for an L-smooth, µ-strongly convex objective with bounded gradient noise.Its stepsize is specified using an iteration-dependent term and a parameter θ.Source: §L.6 · p. 46
M.1. Performance of DIANA, QSGD and Terngrad on the Rosenbrock function
- DIANA vastly outperforms QSGD and TernGrad on the two-worker nonconvex Rosenbrock experiment.The experiment used deterministic gradients and a 1-bit QSGD configuration, which coincides with QGD in this setting.Source: §M.1 · p. 47, Figure M1 · p. 47
M.2. Logistic regression; M.2.1. What α is better to choose?; M.2.3. DIANA vs QSGD vs TernGrad vs DQGD
- On mushrooms logistic regression, α need only be small for either ℓ2 or ℓ∞ quantization; DIANA with momentum was also compared with QSGD, TernGrad, and DQGD.The logistic-regression experiments use ℓ2 and ℓ1 penalties, with the ℓ1 coefficient adjusted for a solution with about 20% nonzero values.Source: §M.2.1 · p. 47, §M.2 · p. 47, §M.2.3 · p. 48
M.2.2. What is the optimal block-size?
- For ℓ∞ quantization, full quantization performed best; for ℓ2, the preferred block size depended on regularization, with about 25 optimal when regularization was large.The parameter vector had dimension 112; small block sizes were preferable when regularization was small.Source: §M.2.2 · p. 48, Figure M2 · p. 49
M.3. MPI - broadcast, reduce and gather
- MPI Gather significantly outperformed point-to-point communication, while repeated shorter-vector gathers corrected the unexpected scaling of 2-bit Gather.The Multi-Gather approach retained favorable scaling as the number of MPI processes increased.Source: §M.3 · p. 48, Figure M3 · p. 50
M.4. Performance of GPU
- Small batch sizes led to low GPU utilization, while larger batches did not necessarily reduce training time.On VGG-A, GPU memory limited gradient-evaluation batches to 128 and function-evaluation batches to 256.Source: Figure M6 · p. 53, §M.4 · p. 53
M.5. Diana vs. TenGrad, SGD and QSGD · M.6. Computational Cost
On MNIST and CIFAR-10, DIANA is compared with SGD, QSGD, and TernGrad; the computational-cost section compares sparse DIANA weight updates with reduce operations across MPI process counts.
M.5. Diana vs. TenGrad, SGD and QSGD
- With DIANA’s 2-bit updates, scaling was anomalous beyond 64 MPI processes for AlexNet and 32 for VGG-A; Multi-Gather improved Gather communication.Figure M7 compares images per second for DIANA against 32-bit-float reduce; experiments used a Cray Aries network, 4 MPI processes per node, and one P100 GPU per process.Source: §M.5 · p. 54, Figure M7 · p. 54
- The MNIST and CIFAR-10 experiments used 8 workers for 3 epochs, learning rates {0.1, 0.2, 0.05}, and quantization bucket sizes {32, 128, 512}.QSGD used 2, 4, or 8 quantization levels; DIANA used α ∈ {0, 1.0/√quantization bucket sizes} with initial h = 0.Source: §M.5 · p. 54
- The experiments also varied momentum, DIANA’s ℓ2 and ℓ∞ norms, and used softmax cross entropy on MNIST-Convex, MNIST-DNN, and CIFAR10-DNN.Momentum values were {0, 0.95, 0.99}; MNIST-Convex is a DNN with no hidden layer, while the other two are convolutional networks.Source: §M.5 · p. 54
- On MNIST-Convex and MNIST-DNN, SGD and DIANA with momentum dominate the other algorithms; on CIFAR10-DNN, both significantly outperform the alternatives.Figure M8 shows the best runs over the tested hyperparameters for training and testing accuracy.Source: §M.5 · p. 55, Figure M8 · p. 56
- For MNIST-DNN, the quantized gradients appear to become sparser as training progresses.Figure M9 tracks quantized-gradient sparsity for DIANA, QSGD, and TernGrad across three problems.Source: §M.5 · p. 55, Figure M9 · p. 57
M.6. Computational Cost
- Figure M10 compares the time to update weights after reduce with the time for sparse DIANA updates across 4–128 MPI processes at 10% sparsity.Source: Figure M10 · p. 58