Message Passing Based Wireless Federated Learning via Analog Message Aggregation
Yiting Zhang, Wei Xu, An Liu, Vincent K. N. Lau · 2024
Federated learning (FL) is a promising learning paradigm that can tackle the increasingly prominent isolated data islands problem while keeping users’ data locally with privacy and security guarantee. Most of the existing works on FL are based on stochastic gradient or momentum acceleration. However, stochastic gradient iterations are known to be slow with only sub-linear order of convergence. As such, it is highly desirable to speed up the convergence of FL and reduce the number of communication rounds. In this paper, we formulate FL as a Bayesian inference problem and propose a message passing based federated learning (MP-FL) scheme to obtain the marginal posterior distribution of the DNN parameters, which not only can be used to accelerate the training process for regression/classification tasks but also quantify the model’s uncertainty. In addition, we present an analog message aggregation scheme to further reduce the communication overhead per round and the multiple access latency by exploiting the free-aggregation of signals over the wireless channel. Simulations show that our proposed method has a faster convergence speed and can significantly reduce both the total number of communication rounds and the communication overhead per round.