Information-Geometric Barycenters for Bayesian Federated Learning

TL;DR

This paper introduces an information-geometric framework for federated learning, interpreting model aggregation as finding barycenters of local posteriors, and proposes BA-BFL, an algorithm with strong convergence and performance in non-i.i.d. settings.

Contribution

It unifies federated learning and Bayesian inference through a geometric perspective and proposes BA-BFL, a new algorithm with theoretical guarantees and practical advantages.

Findings

BA-BFL achieves performance comparable to state-of-the-art methods.

The geometric interpretation provides new insights into model aggregation.

BA-BFL retains convergence properties in non-convex, non-i.i.d. scenarios.

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.