Information-Geometric Barycenters for Bayesian Federated Learning

2024-12-16

AI Summary

In federated learning, inconsistency between Bayesian inference and model aggregation undermines uncertainty quantification and model calibration. Method: We reformulate global aggregation from an information-geometric perspective, generalizing conventional weighted averaging to a divergence-weighted barycenter optimization over local posterior distributions. Specifically, we unify Bayesian federated learning aggregation as a geometric centroid problem under Bregman divergence and propose the Bregman-Aggregation Bayesian Federated Learning (BA-BFL) algorithm. Contribution/Results: BA-BFL retains FedAvg-level convergence guarantees under non-convex, non-IID settings while ensuring theoretical interpretability and empirical robustness. Experiments demonstrate significantly improved uncertainty quantification accuracy and model calibration over state-of-the-art methods. Moreover, our framework provides a unified information-geometric interpretation for diverse mainstream aggregation strategies, bridging Bayesian inference and geometric aggregation in federated learning.

Technology Category

• Machine Learning: Distributed Machine Learning & Federated Learning
• Multiagent Systems: Multiagent Learning
• Reasoning under Uncertainty: Other Foundations of Reasoning under Uncertainty

Application Category

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

Abstract

Federated learning (FL) is a widely used and impactful distributed optimization framework that achieves consensus through averaging locally trained models. While effective, this approach may not align well with Bayesian inference, where the model space has the structure of a distribution space. Taking an information-geometric perspective, we reinterpret FL aggregation as the problem of finding the barycenter of local posteriors using a prespecified divergence metric, minimizing the average discrepancy across clients. This perspective provides a unifying framework that generalizes many existing methods and offers crisp insights into their theoretical underpinnings. We then propose BA-BFL, an algorithm that retains the convergence properties of Federated Averaging in non-convex settings. In non-independent and identically distributed scenarios, we conduct extensive comparisons with statistical aggregation techniques, showing that BA-BFL achieves performance comparable to state-of-the-art methods while offering a geometric interpretation of the aggregation phase. Additionally, we extend our analysis to Hybrid Bayesian Deep Learning, exploring the impact of Bayesian layers on uncertainty quantification and model calibration.

Problem

• Aligning federated learning with Bayesian inference in distribution spaces
• Finding barycenters of local posteriors using divergence metrics
• Improving uncertainty quantification in hybrid Bayesian deep learning

Innovation

• Information-geometric barycenters for Bayesian FL
• BA-BFL algorithm generalizes Federated Averaging
• Hybrid Bayesian Deep Learning for uncertainty quantification