Personalized Bayesian Federated Learning with Wasserstein Barycenter Aggregation

TL;DR

This paper introduces FedWBA, a novel personalized Bayesian federated learning method that uses particle-based variational inference and Wasserstein barycenter aggregation to improve local inference, global aggregation, and convergence guarantees.

Contribution

FedWBA advances PBFL by integrating nonparametric particle-based inference and Wasserstein barycenter aggregation, overcoming parametric assumptions and naive averaging limitations.

Findings

FedWBA outperforms baselines in prediction accuracy.

It improves uncertainty calibration.

It demonstrates faster convergence and robustness.

Abstract

Personalized Bayesian federated learning (PBFL) handles non-i.i.d. client data and quantifies uncertainty by combining personalization with Bayesian inference. However, existing PBFL methods face two limitations: restrictive parametric assumptions in client posterior inference and naive parameter averaging for server aggregation. To overcome these issues, we propose FedWBA, a novel PBFL method that enhances both local inference and global aggregation. At the client level, we use particle-based variational inference for nonparametric posterior representation. At the server level, we introduce particle-based Wasserstein barycenter aggregation, offering a more geometrically meaningful approach. Theoretically, we provide local and global convergence guarantees for FedWBA. Locally, we prove a KL divergence decrease lower bound per iteration for variational inference convergence. Globally, we show that the Wasserstein barycenter converges to the true parameter as the client data size increases. Empirically, experiments show that FedWBA outperforms baselines in prediction accuracy, uncertainty calibration, and convergence rate, with ablation studies confirming its robustness.