Soft-Sign Stochastic Gradient Descent Algorithm for Wireless Federated Learning

Abstract

Federated learning over wireless networks requires aggregating locally computed gradients at a server where the mobile devices send statistically distinct gradient information over heterogenous communication links. This paper proposes a Bayesian approach for wireless federated learning referred to as soft-sign stochastic gradient descent (soft-signSGD). The idea of soft-signSGD is to aggregate the one-bit quantized local gradients at the server by jointly exploiting i) the prior distributions of the local gradients, ii) the gradient quantizer function, and iii) channel distributions. This aggregation method is optimal in the sense of minimizing the mean-squared error (MSE) under a simplified Gaussian prior assumption on the local gradient. From simulations, we demonstrate that soft-signSGD considerably outperforms the conventional sign stochastic gradient descent algorithm when training and testing neural networks using the MNIST dataset and the CIFAR-10 dataset over heterogeneous wireless networks.