Communication-efficient Subspace Methods for High-dimensional Federated Learning

Zai Shi, Atilla Eryilmaz

Abstract: As an emerging technique to employ machine learning processes within an edge computing infrastructure, federated learning (FL) has aroused great interests in both industry and academia. In this paper, we consider a potential challenge of FL in a wireless setup, whereby uplink communication from edge devices to the central server has limited capacity. This is particularly important for machine learning tasks (such as training deep neural networks) in FL with extremely high-dimensional domains that can substantially increase the communication burden. To tackle this challenge, we first propose a basic method called Subspace Stochastic Gradient Descent for Federated Learning (FL-SSGD) to introduce the idea of subspace methods. Through theoretical analysis, we show that by choosing appropriate subspace matrices in FL-SSGD, we can reduce uplink communication costs compared to classical FedAvg method. To improve FL-SSGD, we then propose another method called Subspace Stochastic Variance Reduced Gradient for Federated Learning (FL-SSVRG) that has a faster convergence rate with less assumptions on objective functions. By conducting experiments of a nonconvex machine learning problem in two FL setups, we demonstrate the advantages of our methods compared to other communication-efficient methods.