Federated Generalized Bayesian Learning via Distributed Stein Variational Gradient Descent

Rahif Kassab, Osvaldo Simeone

Keywords: Federated Learning, Distributed Variational Inference

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

One-sentence Summary: This paper introduces a non-parametric generalized Bayesian federated learning framework based on SVGD that allows a flexible trade-off between per-iteration communication load and number of communication rounds.

Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics

Reviewed Version (pdf): https://openreview.net/references/pdf?id=EqKvcIpeL