# Federated Learning as Variational Inference: A Scalable Expectation Propagation Approach

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

## Key facts

- Cited by 17 papers in the map (data through 14 Oct 2026).
- Builds on 38 of its 55 references that are in the map.
- First posted on arXiv 8 Feb 2023.

## Builds on (most cited first, 30 of 38)

- [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)
- [On Calibration of Modern Neural Networks](https://mlsanity.com/paper/1706.04599/) (2017)
- [Advances and Open Problems in Federated Learning](https://mlsanity.com/paper/1912.04977/) (2019)
- [Federated Optimization in Heterogeneous Networks](https://mlsanity.com/paper/1812.06127/) (2018)
- [Federated Learning: Challenges, Methods, and Future Directions](https://mlsanity.com/paper/1908.07873/) (2019)
- [On the Convergence of FedAvg on Non-IID Data](https://mlsanity.com/paper/1907.02189/) (2019)
- [Adaptive Federated Optimization](https://mlsanity.com/paper/2003.00295/) (2020)
- [Tackling the Objective Inconsistency Problem in Heterogeneous Federated Optimization](https://mlsanity.com/paper/2007.07481/) (2020)
- [Towards Federated Learning at Scale: System Design](https://mlsanity.com/paper/1902.01046/) (2019)
- [Federated Optimization: Distributed Machine Learning for On-Device Intelligence](https://mlsanity.com/paper/1610.02527/) (2016)
- [Federated Learning with Matched Averaging](https://mlsanity.com/paper/2002.06440/) (2020)
- [A Simple Baseline for Bayesian Uncertainty in Deep Learning](https://mlsanity.com/paper/1902.02476/) (2019)
- [Agnostic Federated Learning](https://mlsanity.com/paper/1902.00146/) (2019)
- [Deep Kernel Learning](https://mlsanity.com/paper/1511.02222/) (2015)
- [Bayesian Nonparametric Federated Learning of Neural Networks](https://mlsanity.com/paper/1905.12022/) (2019)
- [A Field Guide to Federated Optimization](https://mlsanity.com/paper/2107.06917/) (2021)
- [Stochastic Gradient Descent as Approximate Bayesian Inference](https://mlsanity.com/paper/1704.04289/) (2017)
- [FedML: A Research Library and Benchmark for Federated Machine Learning](https://mlsanity.com/paper/2007.13518/) (2020)
- [Can You Really Backdoor Federated Learning?](https://mlsanity.com/paper/1911.07963/) (2019)
- [Advances in Variational Inference](https://mlsanity.com/paper/1711.05597/) (2017)
- [Limitations of the Empirical Fisher Approximation for Natural Gradient Descent](https://mlsanity.com/paper/1905.12558/) (2019)
- [Expectation Propagation for approximate Bayesian inference](https://mlsanity.com/paper/1301.2294/) (2013)
- [Practical Deep Learning with Bayesian Principles](https://mlsanity.com/paper/1906.02506/) (2019)
- [Minibatch vs Local SGD for Heterogeneous Distributed Learning](https://mlsanity.com/paper/2006.04735/) (2020)
- [A Complete Recipe for Stochastic Gradient MCMC](https://mlsanity.com/paper/1506.04696/) (2015)
- [Fast and Scalable Bayesian Deep Learning by Weight-Perturbation in Adam](https://mlsanity.com/paper/1806.04854/) (2018)
- [Noisy Natural Gradient as Variational Inference](https://mlsanity.com/paper/1712.02390/) (2017)
- [Federated Learning via Posterior Averaging: A New Perspective and Practical Algorithms](https://mlsanity.com/paper/2010.05273/) (2020)
- [Personalized Federated Learning with Gaussian Processes](https://mlsanity.com/paper/2106.15482/) (2021)

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

- [Bayesian Federated Learning: A Survey](https://mlsanity.com/paper/2304.13267/) (2023)
- [FedFisher: Leveraging Fisher Information for One-Shot Federated Learning](https://mlsanity.com/paper/2403.12329/) (2024)
- [FedHB: Hierarchical Bayesian Federated Learning](https://mlsanity.com/paper/2305.04979/) (2023)
- [Leveraging Function Space Aggregation for Federated Learning at Scale](https://mlsanity.com/paper/2311.10291/) (2023)
- [FedSI: Federated Subnetwork Inference for Efficient Uncertainty Quantification](https://mlsanity.com/paper/2404.15657/) (2024)
- [Connecting Federated ADMM to Bayes](https://mlsanity.com/paper/2501.17325/) (2025)
- [FedBEns: One-Shot Federated Learning based on Bayesian Ensemble](https://mlsanity.com/paper/2503.15367/) (2025)
- [Federated ADMM from Bayesian Duality](https://mlsanity.com/paper/2506.13150/) (2025)
- [Federated Variational Inference: Towards Improved Personalization and Generalization](https://mlsanity.com/paper/2305.13672/) (2023)
- [FedPeWS: Personalized Warmup via Subnetworks for Enhanced Heterogeneous Federated Learning](https://mlsanity.com/paper/2410.03042/) (2024)
- [Federated Learning with Uncertainty and Personalization via Efficient Second-order Optimization](https://mlsanity.com/paper/2411.18385/) (2024)
- [Personalized Federated Learning via Amortized Bayesian Meta-Learning](https://mlsanity.com/paper/2307.02222/) (2023)
- [One-Shot Federated Learning with Bayesian Pseudocoresets](https://mlsanity.com/paper/2406.02177/) (2024)
- [Information-Geometric Barycenters for Bayesian Federated Learning](https://mlsanity.com/paper/2412.11646/) (2024)
- [Federated Generalised Variational Inference: A Robust Probabilistic Federated Learning Framework](https://mlsanity.com/paper/2502.00846/) (2025)
- [Personalized Bayesian Federated Learning with Wasserstein Barycenter Aggregation](https://mlsanity.com/paper/2505.14161/) (2025)
- [Cost-Free Personalization via Information-Geometric Projection in Bayesian Federated Learning](https://mlsanity.com/paper/2509.10132/) (2025)

## Abstract (arXiv)

The canonical formulation of federated learning treats it as a distributed optimization problem where the model parameters are optimized against a global loss function that decomposes across client loss functions. A recent alternative formulation instead treats federated learning as a distributed inference problem, where the goal is to infer a global posterior from partitioned client data (Al-Shedivat et al., 2021). This paper extends the inference view and describes a variational inference formulation of federated learning where the goal is to find a global variational posterior that well-approximates the true posterior. This naturally motivates an expectation propagation approach to federated learning (FedEP), where approximations to the global posterior are iteratively refined through probabilistic message-passing between the central server and the clients. We conduct an extensive empirical study across various algorithmic considerations and describe practical strategies for scaling up expectation propagation to the modern federated setting. We apply FedEP on standard federated learning benchmarks and find that it outperforms strong baselines in terms of both convergence speed and accuracy.