Downlink Compression Improves TopK Sparsification

TL;DR

This paper demonstrates that applying gradient compression in both directions during distributed training can reduce communication costs and potentially improve convergence, challenging previous beliefs that extra compression harms model performance.

Contribution

The authors extend the convergence analysis of topK SGD to bidirectional compression and empirically show its benefits over unidirectional compression in distributed neural network training.

Findings

Bidirectional topK SGD reduces communication overhead.

Bidirectional compression can improve convergence bounds.

Models trained with bidirectional topK SGD perform as well as unidirectional methods.

Abstract

Training large neural networks is time consuming. To speed up the process, distributed training is often used. One of the largest bottlenecks in distributed training is communicating gradients across different nodes. Different gradient compression techniques have been proposed to alleviate the communication bottleneck, including topK gradient sparsification, which truncates the gradient to the largest K components before sending it to other nodes. While some authors have investigated topK gradient sparsification in the parameter-server framework by applying topK compression in both the worker-to-server (uplink) and server-to-worker (downlink) direction, the currently accepted belief says that adding extra compression degrades the convergence of the model. We demonstrate, on the contrary, that adding downlink compression can potentially improve the performance of topK sparsification: not only does it reduce the amount of communication per step, but also, counter-intuitively, can improve the upper bound in the convergence analysis. To show this, we revisit non-convex convergence analysis of topK stochastic gradient descent (SGD) and extend it from the unidirectional to the bidirectional setting. We also remove a restriction of the previous analysis that requires unrealistically large values of K. We experimentally evaluate bidirectional topK SGD against unidirectional topK SGD and show that models trained with bidirectional topK SGD will perform as well as models trained with unidirectional topK SGD while yielding significant communication benefits for large numbers of workers.