Personalized Bayesian Federated Learning with Wasserstein Barycenter Aggregation

Feng Zhou (Renmin University of China) · Yifan Sun (University of Illinois at Urbana-Champaign) · Ting Wei (Renmin University of China) · Biao Mei (Renmin University of China) · Junliang Lyu (Peking University) · Renquan Zhang (Dalian University of Technology)

bayesian inference convergence guaranteesgeometrically meaningful approachkl divergence decreasenaive parameter averagingnon-i.i.d. client datanonparametric posterior representationparametric formsparticle-based variational inferenceparticle-based wasserstein barycenter aggregationpersonalized bayesian federated learningposterior aggregationposterior inference uncertainty calibrationwasserstein barycenter

Personalized Bayesian federated learning (PBFL) handles non-i.i.d. client data and quantifies uncertainty by combining personalization with Bayesian inference. However, current PBFL methods face two main limitations: posterior inference on clients often assumes restrictive parametric forms, and server-side posterior aggregation typically relies on naive parameter averaging. 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.