Partitioned Variational Inference: A unified framework encompassing federated and continual learning

Abstract

Variational inference (VI) has become the method of choice for fitting many modern probabilistic models. However, practitioners are faced with a fragmented literature that offers a bewildering array of algorithmic options. First, the variational family. Second, the granularity of the updates e.g. whether the updates are local to each data point and employ message passing or global. Third, the method of optimization (bespoke or blackbox, closed-form or stochastic updates, etc.). This paper presents a new framework, termed Partitioned Variational Inference (PVI), that explicitly acknowledges these algorithmic dimensions of VI, unifies disparate literature, and provides guidance on usage. Crucially, the proposed PVI framework allows us to identify new ways of performing VI that are ideally suited to challenging learning scenarios including federated learning (where distributed computing is leveraged to process non-centralized data) and continual learning (where new data and tasks arrive over time and must be accommodated quickly). We showcase these new capabilities by developing communication-efficient federated training of Bayesian neural networks and continual learning for Gaussian process models with private pseudo-points. The new methods significantly outperform the state-of-the-art, whilst being almost as straightforward to implement as standard VI.

Toward Understanding Catastrophic Forgetting in Continual Learning

“…The Kullback-Leibler term in the variational lower bound of VCL naturally regularizes the approximate posterior toward the prior. Improved training procedures have also been developed for this type of approximate Bayesian continual learning through the use of natural gradients [10,47], fixed-point updates [52], and local approximation [6]. More expressive classes of variational distributions were also considered, including channel factorized Gaussian [22], multiplicative normalizing flow [22], or structured Laplace approximations [33].…”

Compressed Particle-Based Federated Bayesian Learning and Unlearning

“…In Bayesian federated learning, the goal is to obtain a variational distribution q(θ) on the model parameter space that minimizes the global free energy (see, e.g., [1], [7], [14])…”

Communication Efficiency in Federated Learning: Achievements and Challenges

“…In another study [56], authors introduce a Partitioned Variational Inference (PVI) for probabilistic models that work well over federated data. They train a Bayesian Neural Network (BNN) over an FL environment that is allowed for both synchronous or asynchronous model updates across many machines, their proposed approach along with the integration of other methods could allow a more communication-efficient training of BNN on non-iid federated data.…”

“…Ability to handle different phases of the model training well [56] a Partitioned Variational Inference (PVI)…”