Federated Learning as Variational Inference: A Scalable Expectation Propagation Approach
This work addresses scalability and performance in federated learning for distributed data settings, presenting a novel inference-based method that is incremental over prior inference formulations.
The paper tackles federated learning by framing it as a variational inference problem, proposing FedEP, an expectation propagation approach that iteratively refines global posterior approximations through message-passing. It shows that FedEP outperforms strong baselines in convergence speed and accuracy on standard benchmarks.
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.