Federated Generalised Variational Inference: A Robust Probabilistic Federated Learning Framework

2025-02-02

AI Summary

Existing federated learning methods suffer from large prediction bias, miscalibrated uncertainty estimates, and high client-side computational overhead under data non-IIDness and simultaneous misspecification of both prior and likelihood. This paper proposes FedGVI—the first probabilistic federated learning framework with theoretical robustness guarantees against *dual misspecification* (prior + likelihood). Its core contributions are: (1) integrating generalized variational inference with cavity distribution modeling to achieve robust posterior approximation; (2) extending partitioned variational inference to enable conjugate updates, drastically reducing client computation; and (3) establishing algorithmic stability via fixed-point convergence analysis. Experiments across diverse synthetic and real-world benchmarks demonstrate that FedGVI significantly improves predictive accuracy and robustness while yielding well-calibrated uncertainty quantification—outperforming state-of-the-art baselines under challenging distributional shifts and model misspecifications.

Technology Category

• Machine Learning: Calibration & Uncertainty Quantification
• Reasoning under Uncertainty: Probabilistic Programming
• Computer Vision: Adversarial Attacks & Robustness

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
• Web Mining and Content Analysis: Robustness and generalizability of Web mining methods

Abstract

We introduce FedGVI, a probabilistic Federated Learning (FL) framework that is provably robust to both prior and likelihood misspecification. FedGVI addresses limitations in both frequentist and Bayesian FL by providing unbiased predictions under model misspecification, with calibrated uncertainty quantification. Our approach generalises previous FL approaches, specifically Partitioned Variational Inference (Ashman et al., 2022), by allowing robust and conjugate updates, decreasing computational complexity at the clients. We offer theoretical analysis in terms of fixed-point convergence, optimality of the cavity distribution, and provable robustness. Additionally, we empirically demonstrate the effectiveness of FedGVI in terms of improved robustness and predictive performance on multiple synthetic and real world classification data sets.

Problem

• Federated Learning
• Non-ideal Data Handling
• Prediction Uncertainty Estimation

Innovation

• FedGVI
• Non-ideal Data Handling
• Uncertainty Estimation