Federated Learning as Variational Inference: A Scalable Expectation Propagation Approach

TL;DR

This paper presents a novel federated learning framework based on variational inference and expectation propagation, enabling scalable, efficient, and accurate inference of global models across distributed data sources.

Contribution

It introduces FedEP, a new federated learning algorithm using variational inference and message passing, extending the inference perspective of federated learning.

Findings

FedEP outperforms baseline methods in convergence speed.

FedEP achieves higher accuracy on standard benchmarks.

The approach scales effectively to large federated settings.

Abstract

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.