[Paper Review] Federated Learning via Posterior Averaging: A New Perspective and Practical Algorithms

Maruan Al-Shedivat, Jennifer Gillenwater+2 more|arXiv (Cornell University)|Oct 11, 2020

Privacy-Preserving Technologies in DataComputer Science47 references38 citations

TL;DR

This paper reframes federated learning as global posterior inference by aggregating local posteriors, and introduces FEDPA, a practical algorithm that generalizes FEDAVG and achieves state-of-the-art results on several benchmarks.

ABSTRACT

Federated learning is typically approached as an optimization problem, where the goal is to minimize a global loss function by distributing computation across client devices that possess local data and specify different parts of the global objective. We present an alternative perspective and formulate federated learning as a posterior inference problem, where the goal is to infer a global posterior distribution by having client devices each infer the posterior of their local data. While exact inference is often intractable, this perspective provides a principled way to search for global optima in federated settings. Further, starting with the analysis of federated quadratic objectives, we develop a computation- and communication-efficient approximate posterior inference algorithm -- federated posterior averaging (FedPA). Our algorithm uses MCMC for approximate inference of local posteriors on the clients and efficiently communicates their statistics to the server, where the latter uses them to refine a global estimate of the posterior mode. Finally, we show that FedPA generalizes federated averaging (FedAvg), can similarly benefit from adaptive optimizers, and yields state-of-the-art results on four realistic and challenging benchmarks, converging faster, to better optima.

Motivation & Objective

• Propose a probabilistic perspective on federated learning by viewing the global model as a posterior mode obtained from local posteriors.
• Develop a computation- and communication-efficient FEDPA algorithm that works with stateless clients.
• Show how FEDPA generalizes FEDAVG and explain convergence/Bias-Variance trade-offs under local posterior estimation.
• Demonstrate state-of-the-art performance on realistic FL benchmarks across vision and language tasks.

Proposed method

• Formulate FL as a global posterior inference problem where the global posterior decomposes into a product of local posteriors (proposition 1).
• Derive a practical algorithm (FEDPA) that uses local posterior sampling on clients and server-side optimization to find the global posterior mode (equations 3 and 4).
• Implement local posterior inference on clients via SG-MCMC (IASG) to obtain estimates of local means and covariances (µi, Σi).
• Compute client deltas using a linear-time, memory-efficient scheme (∆i = Σi^{-1}(θ − µi)) with a shrinkage covariance estimator to enable O(d) cost per round (theorem 3).
• Show that FEDAVG is a special case of FEDPA with identity covariance, explaining observed stagnation and bias in local-update schemes.

Experimental results

Research questions

• RQ1 Can federated learning be effectively framed as global posterior inference rather than pure optimization?
• RQ2 How can we design a practically efficient FEDPA algorithm with linear cost in model dimensionality for stateless clients?
• RQ3 What is the impact of using local posterior information (means and covariances) on convergence, bias, and final optima compared to FEDAVG?
• RQ4 Do posterior-based methods achieve faster convergence and better optima on realistic FL benchmarks?

Key findings

• FEDPA converges faster and to better optima than FEDAVG on multiple benchmarks by leveraging local posterior information.
• FEDAVG with many local steps corresponds to FEDPA with identity covariances, which introduces bias and stagnation; FEDPA mitigates this with covariance-informed updates.
• Using IASG-based posterior sampling on clients and shrinkage covariance estimators enables O(d) communication and O(d) server computation per round, scalable to high-dimensional models.
• On EMNIST-62 and CIFAR-100, FEDPA-ME attains higher final accuracy and fewer rounds to reach target accuracy than strong baselines.
• On StackOverflow tasks, FEDPA-ME improves macro-F1 and micro-F1, and achieves competitive precision/recall metrics relative to baselines.