Communication-efficient distributed SGD with Sketching

Nikita Ivkin; Daniel Rothchild; Enayat Ullah; Vladimir Braverman; Ion Stoica; Raman Arora

Source-linked AI summary

Communication-efficient distributed SGD with Sketching

Nikita Ivkin, Daniel Rothchild, Enayat Ullah, Vladimir Braverman, Ion Stoica, Raman Arora

arXiv:1903.04488v3Mar 12, 2019cs.LGcs.DCmath.OCstat.ML

TL;DR

Distributed SGD can be limited by communication, motivating a method that sends sketches instead of full gradients. The paper introduces Sketched SGD, which aggregates Count Sketches and recovers sparse updates, achieving O(log d) communication and up to 40x lower total communication without loss in model quality.

Problem
Communication can overwhelm local computation in large-scale distributed SGD, motivating lower per-worker communication as models and worker counts grow.
Method
Sketched SGD sends gradient sketches, merges them at the parameter server, recovers heavy coordinates, and requests their exact values for sparse updates.
Results
SKETCHED-SGD reduces necessary communication to O(log d), provably converges at the same rate as SGD, and reduces total communication by up to 40x without loss in model quality.
Takeaways & Limitations
The method is presented as a communication-efficient approach for scaling distributed SGD to large models and many workers.
Takeaways & Limitations
Combining SKETCHED-SGD with quantization and asynchronous updates remains future work, despite preliminary 16-bit sketch quantization results on CIFAR-10 ResNets.

Abstract

from arXiv · show

Large-scale distributed training of neural networks is often limited by network bandwidth, wherein the communication time overwhelms the local computation time. Motivated by the success of sketching methods in sub-linear/streaming algorithms, we introduce Sketched SGD, an algorithm for carrying out distributed SGD by communicating sketches instead of full gradients. We show that Sketched SGD has favorable convergence rates on several classes of functions. When considering all communication -- both of gradients and of updated model weights -- Sketched SGD reduces the amount of communication required compared to other gradient compression methods from $\mathcal{O}(d)$ or $\mathcal{O}(W)$ to $\mathcal{O}(\log d)$, where $d$ is the number of model parameters and $W$ is the number of workers participating in training. We run experiments on a transformer model, an LSTM, and a residual network, demonstrating up to a 40x reduction in total communication cost with no loss in final model performance. We also show experimentally that Sketched SGD scales to at least 256 workers without increasing communication cost or degrading model performance.

1 Introduction

Distributed synchronous SGD can become communication-bound as worker count grows, so the paper targets per-worker communication rather than increasing batch size. It introduces sketch-based gradient transmission to reduce communication while preserving convergence and model accuracy.

Communication bottleneck: Synchronous data-parallel SGD makes workers compute local gradients, aggregate them at a parameter server, and receive updated weights.The parameter server sums worker gradients into the full mini-batch gradient before sending updated weights back.
Communication bottleneck: Increasing workers yields increasingly marginal training-time reductions while overall training cost grows linearly with W.Even with optimal communication-computation interleaving, training time is bounded by the larger per-worker communication or computation time.
Motivation: Large-batch training can increase per-worker computation, but evidence indicates a maximum batch size beyond which convergence stops improving and generalization error increases.The paper therefore focuses on decreasing communication cost per worker instead.
Proposed direction: Sketched SGD sends gradient sketches of size O(log d) instead of full gradients, with claimed convergence guarantees, sub-linear dependence on d, and constant dependence on W.The paper reports high compression for large d without loss in model accuracy and expected scaling to large W.

2 Related Work

Prior distributed-SGD compression methods mainly quantize or sparsify gradients, but their communication reductions trade off against iteration complexity or lack convergence guarantees. Earlier sketching approaches also leave important distributed communication overheads or theoretical gaps.

Quantization: Quantization methods either provide only constant per-iteration savings or increase the iteration count asymptotically.The paper notes that the latter tradeoff can leave total training communication unimproved or worse.
Sparsification: Gradient sparsification can perform well empirically, but estimated mini-batch gradients may be far from the true gradients and generally lack convergence guarantees.Some methods explicitly assume conditions that control this approximation error.
Bias and convergence: Unbiased compression preserves standard SGD guarantees, while biased compression can work empirically and has only more recently received convergence analyses.The related-work discussion distinguishes theoretical guarantees from observed practical performance.
Sketching: Earlier distributed and federated sketching methods either lack provable guarantees or apply sketches only to portions of the data, leaving Ω(Wd) communication overhead.The cited approaches include sketched updates and hashing or quantization methods that improve constants in O(Wd).
Sketching: Other sketching work targets models that cannot fit in memory, whereas this paper assumes a dense model fits and uses sketches to reduce communication.The distinction is between memory reduction and communication reduction.

3 Preliminaries

The paper formulates distributed SGD and introduces Count Sketch as a way to identify large gradient coordinates using sublinear communication. Its linearity lets workers’ sketches be merged before recovering approximate heavy hitters.

SGD: SGD updates model parameters using a step size and an unbiased stochastic gradient computed from a randomly selected mini-batch.The objective is the expected loss over data points.
Assumptions: The analysis assumes bounded stochastic-gradient moments and uses smoothness and strong convexity as standard function conditions.These assumptions support the stated SGD convergence results.
Convergence background: For smooth strongly convex functions, vanilla SGD has an O(1/T) convergence rate.The paper cites prior work for this standard result.
Count Sketch: A gradient sketch has size O(1/ε log d), providing a communication representation smaller than the full d-dimensional gradient.The sketch is designed for large-coordinate recovery rather than exact recovery of every coordinate.
Count Sketch: Count Sketch approximates gradient coordinates and the full ℓ2 norm, enabling recovery of ℓ2 heavy hitters and approximation of top-k coordinates.With sufficiently small ε, heavy hitters can approximate the largest gradient coordinates.
Distributed aggregation: Because Count Sketch is linear, the parameter server can sum worker sketches to obtain a sketch of the summed gradient.The merged sketch then supports approximate top-k recovery across distributed workers.

4 Sketched SGD

Sketched SGD replaces full distributed gradients with merged Count Sketches, recovers heavy coordinates, and requests their exact values for sparse updates. The analysis establishes contraction-based convergence under stated assumptions and extends it to distributed training.

Algorithm: Each worker transmits a gradient sketch; the server sums sketches, recovers large coordinates, and requests exact values for a second-round sparse update.The update retains only selected coordinates, while omitted gradient components are accumulated locally for later iterations.
Algorithm: The recovery routine queries an approximate gradient norm, combines heavy hitters with randomly selected non-heavy coordinates, and outputs a k-sparse exact update.The selected coordinates are obtained through a second communication round.
Convergence results: Theorem 1 analyzes SKETCHED-SGD for smooth strongly convex objectives with W workers, sketch size O(k log(dT/δ)), and a decreasing step size.The theorem provides a bound on expected final suboptimality after T steps.
Analysis: The proof first establishes contraction-based convergence for single-machine sparsified SGD, then uses Count Sketch linearity and norm approximation for the distributed setting.The full proof is deferred to the appendix.
Analysis: HEAVYMIX with sketch size Θ(k log(d/δ)) is a k/d-contraction with probability at least 1−δ.This contraction property supplies the single-machine convergence ingredient for the algorithm.

5 Empirical Results

The experiments evaluate SKETCHED-SGD’s implementation and communication–accuracy tradeoffs across transformer, LSTM, and residual-network models. Results show substantial compression with comparable performance, while sketch size, recovery parameters, and worker count affect convergence and accuracy.

Training Algorithm: SKETCHED-SGD uses momentum, Count Sketch aggregation, exact-value recovery for candidate coordinates, and sparse weight updates.Its resulting per-worker communication cost is Pk + |S| + k, including the sketch, second-round values, and returned updates.
Sketching Implementation: The Count Sketch provides a provable ±ε∥g∥2 approximation for each queried gradient coordinate, with GPU-wide querying used to select large elements.The implementation queries every gradient coordinate and chooses the largest elements because efficient heavy-coordinate discovery with negative inputs is unavailable.
Large d: 90M-parameter transformer and 73M-parameter LSTM experiments on WMT 2014 measure total communication as 2d/(|S| + Pk + k).The denominator includes sketch transmission, second-round exact values, and modified parameters sent back to workers.
Large d: Extremely small sketch sizes or P values can diverge or produce very low BLEU, whereas higher compression prefers a relatively larger P.At comparable sketch-to-exact-value ratios, uncompressed SGD trains fastest, followed by top-k SGD, then 20x and 40x SKETCHED-SGD.
Large W: For 4-worker CIFAR-10 training, increasing compression decreases test accuracy, while very large sketches relative to Pk also tend to yield poor results.The observed tradeoff can sometimes be offset by training longer.

6 Discussion

SKETCHED-SGD reduces distributed-SGD communication to O(log d), while maintaining convergence and model quality. Experiments report up to 40x lower total communication and no appreciable performance loss through 256 workers.

Neither SKETCHED-SGD nor local top-k shows an appreciable performance drop through 256 workers on CIFAR10.
Local top-k communication grows toward ≈d = 6.5 × 10^6 as worker count increases, limiting its best overall compression to 2x for many workers.
O(log d) communication further reduces the O(d) or O(W) requirements of uncompressed SGD and other gradient-compression methods.
Up to 40x lower total communication is achieved without loss in model quality.
Quantizing sketches to 16 bits caused no accuracy drop in preliminary ResNet experiments, while full evaluation with quantization and asynchronous updates remains future work.
The O(log d) communication requirement is presented as having potential for workloads made difficult by growing model sizes and worker counts.

Supplementary

The supplementary material analyzes HEAVYMIX, which recovers heavy coordinates from a Count Sketch and completes a k-element update with exact-value communication. Its guarantees rely on contraction properties and a second communication round.

HEAVYMIX extracts heavy elements from a Count Sketch and partitions recovered coordinates into heavy and not-heavy sets.
The heavy set contains at most k elements, and uniformly sampled not-heavy elements fill the output to k coordinates.
A second communication round obtains exact values for the selected k coordinates used by the update.
With probability at least 1 − δ, the recovered heavy coordinates have exact values under the heavy-hitters recovery event.
The proof bounds the residual not-heavy vector using the heavy-coordinate threshold and the relation between the full and heavy-vector norms.

A.1 Proof of the main theorem

The main theorem proof transfers single-machine compressed-SGD convergence to distributed training by exploiting Count Sketch linearity and unbiased averaged stochastic gradients. The result is a high-probability convergence guarantee under smoothness, strong convexity, and contraction assumptions.

Distributed extension: Count Sketch linearity lets the merged sketch represent the sum of workers’ gradients, whose recovered top-k coordinates define the distributed update.
Distributed extension: Averaging stochastic gradients preserves unbiasedness, reducing the distributed analysis to the one-machine SKETCHED-SGD result.
Theorem conditions: Theorem 2 assumes an L-smooth, µ-strongly convex objective, sketch size O(k log(dT/δ)), and step size η_t = 1/(t+ξ).
Theorem guarantee: The theorem guarantees the output after T steps satisfies its stated error bound with probability at least 1 − δ.
Proof ingredients: The proof controls stochastic-gradient terms using unbiasedness, variance bounds, strong convexity, smoothness, and step-size restrictions.
Compression guarantee: HEAVYMIX supplies the required contraction property with sketch size Θ(k log(d/δ)) and probability at least 1 − δ.

B Auxiliary results

The auxiliary results establish convergence guarantees for Sketched SGD under smooth strongly convex, convex nonsmooth, and non-convex smooth objectives. They also connect HEAVYMIX’s contraction guarantee to these single-machine and distributed analyses.

HEAVYMIX is a k/d-contraction with sketch size Θ(k log(d/δ)), enabling the convergence analyses to apply to the sketched update.The proof strategy directly plugs this contraction property into existing sparsified-SGD results.
Theorem 3 analyzes Sketched SGD for smooth strongly convex objectives with bounded stochastic gradients and sketch size O(k log(dT/δ)).The result uses a step size η_t = 1/(t+ξ) and gives a high-probability guarantee after T steps.
Theorem 4 extends the analysis to non-convex smooth objectives under unbiased, bounded stochastic gradients and sketch size O(k log(dT/δ)).The theorem uses η = 1/√T.
Theorems 5 and 6 provide auxiliary convergence results for convex nonsmooth and smooth strongly convex objectives, respectively.The convex nonsmooth result uses averaged iterates, while the strongly convex result gives an output after T steps.
Mergeability lets the parameter server combine worker sketches and recover heavy coordinates of the aggregate gradient.Because averaging preserves the contraction guarantee, the distributed procedure reduces to single-machine Sketched SGD with an unbiased average gradient.

C Sketching

The paper builds on streaming sketching primitives, especially AMS and Count Sketch, to estimate norms and identify heavy coordinates with compact, mergeable representations. Count Sketch supplies the heavy-coordinate recovery used by Sketched SGD.

The AMS sketch estimates the ℓ2 norm of a d-dimensional update vector using random signs and repeated measurements.Independent repetitions control variance, while a median filter controls failure probability.
Count Sketch extends AMS-style sketching to identify heavy coordinates by hashing coordinates into buckets and estimating each bucket’s ℓ2 norm.When values are skewed, large-norm buckets are likely to contain heavy coordinates.
Count Sketch is mergeable: sketches of separate vectors can be combined without changing approximation guarantees or memory footprint.Its simple implementation and parallelizability also facilitate GPU acceleration.
Count Sketch finds approximate top-k coordinates with probability at least 1−δ using sublinear space.The approximation scheme uses the k-th largest coordinate as its threshold.
Count Sketch approximates heavy-coordinate values with error ±ε∥g∥2 and has a compact memory footprint.The same guarantee is stated for (α, ℓ2)-heavy coordinates.

D Model Training Details

The experiments train sketched and baseline models across translation and image-classification settings, using OpenNMT implementations for sequence models and a DAWNBench residual network for CIFAR-10.

The transformer and LSTM experiments use modified OpenNMT code and evaluate the WMT 2014 English-to-German translation task.The data are processed with a standard tokenizer and shuffled.
Transformer: The baseline transformer uses six layers, 512-dimensional recurrent and word-vector settings, SGD, momentum 0.9, and 60,000 training steps.Its command specifies a token batch size of 1024 and four-way accumulation.
Residual network: The residual network is trained on CIFAR-10 using the fastest-training DAWNBench model, batch size 512, and a learning rate scheduled from 0 to 0.4 and back to 0.Training data are augmented with padding and random cropping.
Each experimental run uses a single GPU selected from Titan X, Titan Xp, Titan V, Tesla P100, and Tesla V100 systems.

E Additional experiments

Additional experiments compare vanilla and sketched SVM and ℓ2-regularized logistic regression on MNIST, covering nonsmooth convex and smooth strongly convex optimization.

The MNIST experiments compare vanilla and sketched SVM and ℓ2-regularized logistic regression against the theoretical rates from Theorems 5 and 1.The sketch has size 280, with 40 columns and 7 rows; k = 10 and P = 10.