[Paper Review] PowerSGD: Practical Low-Rank Gradient Compression for Distributed Optimization

Thijs Vogels, Sai Praneeth Karimireddy+1 more|arXiv (Cornell University)|May 31, 2019

Advanced Data Compression TechniquesComputer Science35 references96 citations

TL;DR

PowerSGD introduces a low-rank gradient compressor based on power iteration with error feedback and all-reduce aggregation, achieving SGD-like accuracy with substantial communication speedups in distributed training.

ABSTRACT

We study gradient compression methods to alleviate the communication bottleneck in data-parallel distributed optimization. Despite the significant attention received, current compression schemes either do not scale well or fail to achieve the target test accuracy. We propose a new low-rank gradient compressor based on power iteration that can i) compress gradients rapidly, ii) efficiently aggregate the compressed gradients using all-reduce, and iii) achieve test performance on par with SGD. The proposed algorithm is the only method evaluated that achieves consistent wall-clock speedups when benchmarked against regular SGD with an optimized communication backend. We demonstrate reduced training times for convolutional networks as well as LSTMs on common datasets. Our code is available at https://github.com/epfml/powersgd.

Motivation & Objective

• Motivate and address the communication bottleneck in data-parallel distributed optimization for deep learning.
• Develop a scalable, linear gradient compressor that supports all-reduce aggregation.
• Ensure convergence and maintain test accuracy through error feedback for biased compression.
• Demonstrate wall-clock speedups over full-precision SGD on CNNs and LSTMs across datasets and hardware.
• Provide a practical, adaptable method with open-source implementation.

Proposed method

• Introduce a rank-r gradient compressor that approximates each gradient matrix M as PQ^T using one step of subspace/power iteration with warm-starting from the previous step.
• Leverage linearity to enable all-reduce based aggregation of compressed gradients across W workers.
• Apply error feedback with post-compression momentum to enable convergence with biased compression.
• Decompose gradients per layer into matrices and compress each matrix independently, using P in R^{n×r} and Q in R^{m×r} for r small (1–4).
• Decompress as PQ^T and perform distributed updates with momentum SGD (EF-SGD with Momentum).
• Empirically compare against SGD and other compressors (Signum, Spectral Atomo) on CIFAR-10/ResNet18 and Wikitext-2/LSTM, reporting wall-clock time, data sent, and accuracy.

Experimental results

Research questions

• RQ1 Can a low-rank gradient compressor based on power iteration achieve substantial communication reductions without sacrificing test accuracy?
• RQ2 Does error feedback enable convergence and robustness of biased compression schemes in non-convex deep learning settings?
• RQ3 How does warm-starting the power iteration affect approximation quality and training accuracy over time?
• RQ4 What are the trade-offs between rank, communication overhead, and end-to-end training time across CNNs and RNN/LSTM models?
• RQ5 How does PowerSGD scale with the number of workers and different communication backends (NCCL vs GLOO)?

Key findings

• PowerSGD achieves wall-clock speedups over regular SGD in a 16-GPU setting using optimized NCCL backends, while compressing gradients more than 120×.
• On CIFAR-10 with ResNet-18, Rank-2 PowerSGD reaches test accuracy 94.4% with 8 MB communicated per epoch, compared to SGD’s 94.3% with 1023 MB.
• For language modeling on Wikitext-2 with an LSTM, Rank-4 PowerSGD attains similar perplexity to SGD while reducing data sent per epoch by ~99% and end-to-end training time by ~55%.
• PowerSGD with error feedback outperforms unbiased rank-based compression in test accuracy (e.g., biased Rank-2 reaches 94.4% vs unbiased ~75–76% under similar settings).
• Warm-starting the power iteration closes the gap to the best rank-r approximation, maintaining or improving test accuracy while avoiding SVD costs.
• The method scales favorably with the number of workers via all-reduce, outperforming alternatives (Signum) on both speed and scalability on optimized backends.