Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Connecting Federated ADMM to Bayes

Siddharth Swaroop, Mohammad Emtiyaz Khan, Finale Doshi-Velez

TL;DR

The paper addresses bridging ADMM-based federated optimization and VB-based Bayesian federated learning by showing that ADMM dual variables align with VB site parameters under isotropic Gaussian assumptions. It then derives FedLap, FedLap-Cov, and FedLap-Func as VB-derived ADMM variants that incorporate covariance learning and function-space information. Empirical results across multiple datasets demonstrate improved convergence and performance relative to standard baselines such as FedAvg, FedProx, and FedDyn, with further gains from covariance and function-space extensions. This work provides a principled integration of duality with uncertainty in federated learning, enabling new algorithms for non-convex settings.

Abstract

We provide new connections between two distinct federated learning approaches based on (i) ADMM and (ii) Variational Bayes (VB), and propose new variants by combining their complementary strengths. Specifically, we show that the dual variables in ADMM naturally emerge through the 'site' parameters used in VB with isotropic Gaussian covariances. Using this, we derive two versions of ADMM from VB that use flexible covariances and functional regularisation, respectively. Through numerical experiments, we validate the improvements obtained in performance. The work shows connection between two fields that are believed to be fundamentally different and combines them to improve federated learning.

Connecting Federated ADMM to Bayes

TL;DR

The paper addresses bridging ADMM-based federated optimization and VB-based Bayesian federated learning by showing that ADMM dual variables align with VB site parameters under isotropic Gaussian assumptions. It then derives FedLap, FedLap-Cov, and FedLap-Func as VB-derived ADMM variants that incorporate covariance learning and function-space information. Empirical results across multiple datasets demonstrate improved convergence and performance relative to standard baselines such as FedAvg, FedProx, and FedDyn, with further gains from covariance and function-space extensions. This work provides a principled integration of duality with uncertainty in federated learning, enabling new algorithms for non-convex settings.

Abstract

We provide new connections between two distinct federated learning approaches based on (i) ADMM and (ii) Variational Bayes (VB), and propose new variants by combining their complementary strengths. Specifically, we show that the dual variables in ADMM naturally emerge through the 'site' parameters used in VB with isotropic Gaussian covariances. Using this, we derive two versions of ADMM from VB that use flexible covariances and functional regularisation, respectively. Through numerical experiments, we validate the improvements obtained in performance. The work shows connection between two fields that are believed to be fundamentally different and combines them to improve federated learning.

Paper Structure

This paper contains 25 sections, 31 equations, 1 figure, 5 tables.

Table of Contents

  1. Introduction
  2. Federated Learning with ADMM and Bayes
  3. Connecting Variational Bayes to ADMM
  4. Connecting Partitioned Variational Inference (PVI) to ADMM
  5. FedLap: A Laplace Version of PVI
  6. FedLap-Cov: A New ADMM Variant with Flexible Covariances
  7. FedLap-Func: Improving by including function-space information
  8. Experiments
  9. Conclusions and Future Work
  10. Making FedLap More Similar to the ADMM update
  11. Derivation of FedLap-Cov Updates
  12. Derivation of FedLap-Func equations
  13. Derivation of \ref{['eq:FedLapFunc_global_interm']}
  14. Deriving the FedLap-Func algorithm
  15. Benefits of FedLap over FedDyn and FedADMM
  16. ...and 10 more sections

Figures (1)

  • Figure 1: FedDyn's performance is sensitive to the additional weight-decay hyperparameter that it has compared to FedLap. When set an order of magnitude too small or too large, performance drops significantly. When there is no weight-decay, performance gets worse over rounds. Results are mean and standard deviation over three seeds on homogeneous MNIST (10 clients).