Stochastic gradient compression for federated learning over wireless network

Abstract

As a mature distributed machine learning paradigm, federated learning enables wireless edge devices to collaboratively train a shared AI-model by stochastic gradient descent (SGD). However, devices need to upload high-dimensional stochastic gradients to edge server in training, which cause severe communication bottleneck. To address this problem, we compress the communication by sparsifying and quantizing the stochastic gradients of edge devices. We first derive a closed form of the communication compression in terms of sparsification and quantization factors. Then, the convergence rate of this communication-compressed system is analyzed and several insights are obtained. Finally, we formulate and deal with the quantization resource allocation problem for the goal of minimizing the convergence upper bound, under the constraint of multiple-access channel capacity. Simulations show that the proposed scheme outperforms the benchmarks.