Personalized Bayesian Federated Learning with Wasserstein Barycenter Aggregation

2025-05-20

AI Summary

To address the dual limitations in personalized Bayesian federated learning under non-IID clients—namely, posterior inference relying on strong parametric assumptions and server-side aggregation restricted to simple averaging—this paper proposes the first framework integrating nonparametric particle-based variational inference with Wasserstein barycenter geometric aggregation. Our method eliminates reliance on pre-specified posterior distribution families, enabling flexible local uncertainty modeling; meanwhile, Wasserstein barycenter aggregation ensures global consistency and provides, for the first time, theoretical guarantees including a lower bound on local KL divergence convergence and barycenter consistency. Experiments demonstrate that our approach consistently outperforms state-of-the-art baselines in prediction accuracy, uncertainty calibration, and convergence speed. Ablation studies further confirm its robustness to data heterogeneity and label noise.

Technology Category

• Machine Learning: Bayesian Learning
• Reasoning under Uncertainty: Probabilistic Programming
• Multiagent Systems: Adversarial Agents

Application Category

• User Modeling, Personalization and Recommendation: Federated recommendation systems and personalization
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

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.

Problem

• Handles non-i.i.d. client data in federated learning
• Overcomes restrictive parametric assumptions in Bayesian inference
• Improves server aggregation via Wasserstein barycenter geometry

Innovation

• Particle-based variational inference for local nonparametric posterior
• Wasserstein barycenter aggregation for global geometric optimization
• Theoretical convergence guarantees for both local and global performance