# Partitioned Variational Inference: A Framework for Probabilistic Federated Learning

arXiv: https://arxiv.org/abs/2202.12275 · HTML: https://mlsanity.com/paper/2202.12275/

## Key facts

- Cited by 13 papers in the map (data through 14 Oct 2026).
- Builds on 26 of its 68 references that are in the map.
- First posted on arXiv 24 Feb 2022.

## Builds on (most cited first, 26 of 26)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Communication-Efficient Learning of Deep Networks from Decentralized Data](https://mlsanity.com/paper/1602.05629/) (2016)
- [Variational Inference with Normalizing Flows](https://mlsanity.com/paper/1505.05770/) (2015)
- [Gaussian Processes for Big Data](https://mlsanity.com/paper/1309.6835/) (2013)
- [Variational Continual Learning](https://mlsanity.com/paper/1710.10628/) (2017)
- [Multitask learning and benchmarking with clinical time series data](https://mlsanity.com/paper/1703.07771/) (2017)
- [Advances in Variational Inference](https://mlsanity.com/paper/1711.05597/) (2017)
- [Expectation Propagation for approximate Bayesian inference](https://mlsanity.com/paper/1301.2294/) (2013)
- [Automatic Differentiation Variational Inference](https://mlsanity.com/paper/1603.00788/) (2016)
- [Conjugate-Computation Variational Inference : Converting Variational Inference in Non-Conjugate Models to Inferences in Conjugate Models](https://mlsanity.com/paper/1703.04265/) (2017)
- [Generalized Variational Continual Learning](https://mlsanity.com/paper/2011.12328/) (2020)
- [The Information Geometry of Mirror Descent](https://mlsanity.com/paper/1310.7780/) (2013)
- [Streaming Variational Bayes](https://mlsanity.com/paper/1307.6769/) (2013)
- [Streaming Sparse Gaussian Process Approximations](https://mlsanity.com/paper/1705.07131/) (2017)
- [Natural Gradients in Practice: Non-Conjugate Variational Inference in Gaussian Process Models](https://mlsanity.com/paper/1803.09151/) (2018)
- [Fast yet Simple Natural-Gradient Descent for Variational Inference in Complex Models](https://mlsanity.com/paper/1807.04489/) (2018)
- [Deep Probabilistic Programming](https://mlsanity.com/paper/1701.03757/) (2017)
- [Stochastic Expectation Propagation](https://mlsanity.com/paper/1506.04132/) (2015)
- [Overpruning in Variational Bayesian Neural Networks](https://mlsanity.com/paper/1801.06230/) (2018)
- [Faster Stochastic Variational Inference using Proximal-Gradient Methods with General Divergence Functions](https://mlsanity.com/paper/1511.00146/) (2015)
- [Distributed Bayesian Learning with Stochastic Natural-gradient Expectation Propagation and the Posterior Server](https://mlsanity.com/paper/1512.09327/) (2015)
- [Differentially Private Variational Inference for Non-conjugate Models](https://mlsanity.com/paper/1610.08749/) (2016)
- [Improving and Understanding Variational Continual Learning](https://mlsanity.com/paper/1905.02099/) (2019)
- [A trust-region method for stochastic variational inference with applications to streaming data](https://mlsanity.com/paper/1505.07649/) (2015)
- [Differentially Private Federated Variational Inference](https://mlsanity.com/paper/1911.10563/) (2019)
- [Incremental Variational Inference for Latent Dirichlet Allocation](https://mlsanity.com/paper/1507.05016/) (2015)

## Built on by (most cited first, 13 of 13)

- [Federated Learning as Variational Inference: A Scalable Expectation Propagation Approach](https://mlsanity.com/paper/2302.04228/) (2023)
- [FedHB: Hierarchical Bayesian Federated Learning](https://mlsanity.com/paper/2305.04979/) (2023)
- [Knowledge Adaptation as Posterior Correction](https://mlsanity.com/paper/2506.14262/) (2025)
- [Federated Variational Inference Methods for Structured Latent Variable Models](https://mlsanity.com/paper/2302.03314/) (2023)
- [Connecting Federated ADMM to Bayes](https://mlsanity.com/paper/2501.17325/) (2025)
- [Federated ADMM from Bayesian Duality](https://mlsanity.com/paper/2506.13150/) (2025)
- [Differentially private partitioned variational inference](https://mlsanity.com/paper/2209.11595/) (2022)
- [One-Shot Federated Learning with Bayesian Pseudocoresets](https://mlsanity.com/paper/2406.02177/) (2024)
- [Federated Generalised Variational Inference: A Robust Probabilistic Federated Learning Framework](https://mlsanity.com/paper/2502.00846/) (2025)
- [Compressed Bayesian Federated Learning for Reliable Passive Radio Sensing in Industrial IoT](https://mlsanity.com/paper/2405.05855/) (2024)
- [Bayesian Federated Learning for Continual Training](https://mlsanity.com/paper/2504.15328/) (2025)
- [Decentralized Federated Learning of Probabilistic Generative Classifiers](https://mlsanity.com/paper/2507.17285/) (2025)
- [Gaussian Mean Field Variational Inference can Overestimate Predictive Variance](https://mlsanity.com/paper/2606.25745/) (2026)

## Abstract (arXiv)

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.