Communication-Efficient Distributed Bayesian Federated Learning Over Arbitrary Graphs

Abstract

This paper investigates a fully distributed federated learning (FL) problem, in which each device is restricted to only utilize its local dataset and the information received from its adjacent devices that are defined in a communication graph to update the local model weights for minimizing the global loss function. To incorporate the communication graph constraint into the joint posterior distribution, we exploit the fact that the model weights on each device is a function of its local likelihood and local prior and then, the connectivity between adjacent devices is modeled by a Dirac Delta distribution. In this way, the joint distribution can be factorized naturally by a factor graph. Based on the Dirac Delta-based factor graph, we propose a novel distributed approximate Bayesian inference algorithm that combines loopy belief propagation (LBP) and variational Bayesian inference (VBI) for distributed FL. Specifically, VBI is used to approximate the non-Gaussian marginal posterior as a Gaussian distribution in local training process and then, the global training process resembles Gaussian LBP where only the mean and variance are passed among adjacent devices. Furthermore, we propose a new damping factor design according to the communication graph topology to mitigate the potential divergence and achieve consensus convergence. Simulation results verify that the proposed solution achieves faster convergence speed with better performance than baselines.