Federated Generalized Bayesian Learning via Distributed Stein Variational Gradient Descent

This paper introduces Distributed Stein Variational Gradient Descent (DSVGD), a non-parametric generalized Bayesian inference framework for federated learning that enables a flexible trade-off between per-iteration communication load and performance. DSVGD maintains a number of non-random and interacting particles at a central server that represent the current iterate of the model global posterior. The particles are iteratively downloaded and updated by one of the agents by minimizing the local free energy with the end goal of minimizing the global free energy. By using a sufficiently large number of particles, DSVGD is shown to outperform benchmark frequentist and Bayesian federated learning strategies, also scheduling a single device per iteration, in terms of accuracy, number of communication rounds, and scalability with respect to the number of agents, while also providing well-calibrated, and hence trustworthy, predictions.