High-Dimensional Stochastic Gradient Quantization for Communication-Efficient Edge Learning

Abstract

Edge machine learning involves the deployment of machine learning algorithms at the network edge so as to leverage massive mobile data and distributed computation resources. Many edge learning frameworks (e.g., federated learning) have been developed based on distributed gradient descent. Based on the approach, stochastic gradients are computed at edge devices and then transmitted to an edge server for aggregation for updating a global AI model. Since each gradient is typically high-dimensional (with millions to billions of coefficients), communication overhead may become a bottleneck for edge learning. In this work, we propose a novel gradient compression scheme to reduce the aforementioned overhead. Specifically, in the proposed scheme, the norm of the stochastic gradient is quantized using a uniform quantizer while the normalized stochastic gradient is decomposed into block gradients. A Grassmannian codebook is applied to quantizing each normalized block gradients. Their quantized versions are assembled using a so-called hinge vector, which is quantized using another Grassmannian codebook. Furthermore, a practical bit-allocation strategy is developed. By simulations, we show that similar learning performance can be achieved with substantially lower communication overhead as compared to the one-bit scalar quantization schemes used in the state-of-the-art design, namely signed SGD.