Federated Generalised Variational Inference: A Robust Probabilistic Federated Learning Framework

We introduce FedGVI, a probabilistic Federated Learning (FL) framework that is 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 to likelihood misspecification. Further, we empirically demonstrate the effectiveness of FedGVI in terms of improved robustness and predictive performance on multiple synthetic and real world classification data sets.

Our paper tackles the following questions: (1) Can we train a probabilistic machine learning model on different, private data sets from individuals without sharing them with other individuals or a server, (2) can we be sure that potential misspecification, such as outliers in any individual's data, does not negatively impact the model, and (3) can we train our model efficiently? These questions have been answered in isolation but not in combination. Our approach, FedGVI, allows us to answer (1), (2), and (3) simultaneously. The main result is that FedGVI successfully disregards misspecification during training that can harmfully impact the model, and that this model can be found faster than models found through existing methods that only solve (1). We demonstrate this by training a Neural Network for classifying images when some of the training data has the wrong label showing that FedGVI leads to better results than other methods. These findings have important practical implications, for instance in healthcare where patient data is sensitive and models need to be reliable.