Partitioned Variational Inference: A Framework for Probabilistic Federated Learning

Abstract

The proliferation of computing devices has brought about an opportunity to deploy machine learning models on new problem domains using previously inaccessible data. Traditional algorithms for training such models often require data to be stored on a single machine with compute performed by a single node, making them unsuitable for decentralised training on multiple devices. This deficiency has motivated the development of federated learning algorithms, which allow multiple data owners to train collaboratively and use a shared model whilst keeping local data private. However, many of these algorithms focus on obtaining point estimates of model parameters, rather than probabilistic estimates capable of capturing model uncertainty, which is essential in many applications. Variational inference (VI) has become the method of choice for fitting many modern probabilistic models. In this paper we introduce partitioned variational inference (PVI), a general framework for performing VI in the federated setting. We develop new supporting theory for PVI, demonstrating a number of properties that make it an attractive choice for practitioners; use PVI to unify a wealth of fragmented, yet related literature; and provide empirical results that showcase the effectiveness of PVI in a variety of federated settings.

Federated Learning as Variational Inference: A Scalable Expectation Propagation Approach

“…son et al, 2016) to personalized FL. Finally, some prior works also consider applying EP to federated learning (Corinzia et al, 2019;Kassab & Simeone, 2022;Ashman et al, 2022), but mostly on relatively small-scale tasks. In this work, we instead discuss and empirically study various algorithmic considerations to scale up expectation propagation to contemporary benchmarks.…”

“…However, exact posterior inference is in general intractable for even modestly-sized models and datasets and requires approx-However, scaling up classic expectation propagation to the modern federated setting is challenging due to the high dimensionality of model parameters and the large number of clients. Indeed, while there is some existing work on expectation propagation-based federated learning (Corinzia et al, 2019;Kassab & Simeone, 2022;Ashman et al, 2022), they typically focus on small models (fewer than 100K parameters) and few clients (at most 100 clients). In this paper we conduct an extensive empirical study across various algorithmic considerations to scale up expectation propagation to contemporary benchmarks (e.g., models with many millions of parameters and datasets with hundreds of thousands of clients).…”

Bayesian data fusion for distributed learning

“…In Federated Learning (FL), deterministic strategies like FedAvg and enhance convergence over heterogeneous datasets. As outlined in Chapter 2, Section 2.5.4, we employ Laplace approximation [51] for posterior estimation of local parameters before aggregation, a methodology paralleled by approaches utilizing Variational Inference [76,77] and other Bayesian frameworks like Variational Federated Learning [78,79,80], MCMC [81], and Gaussian Process [48].…”