A Bayesian Approach for Personalized Federated Learning in Heterogeneous Settings

Disha Makhija, Joydeep Ghosh, Nhat Ho

Primary Area: general machine learning (i.e., none of the above)

Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.

Keywords: Federated Learning, Heterogeneous Settings, Bayesian Learning, Privacy-aware

Submission Guidelines: I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2024/AuthorGuide.

TL;DR: A privacy-aware algorithm for federated learning in heterogeneous settings.

Abstract: In several practical applications of federated learning (FL), the clients are highly heterogeneous in terms of both their data and compute resources, and therefore enforcing the same model architecture for each client is very limiting. The need for uncertainty quantification is also often particularly amplified for clients that have limited local data. This paper presents a unified FL framework based on training customized local Bayesian models that can simultaneously address both these constraints. A Bayesian framework provides a natural way of incorporating supervision in the form of prior distributions. We use priors in the functional (output) space of the networks to facilitate collaboration across heterogeneous clients via an unlabelled auxiliary dataset. We further present a differentially private version of the algorithm along with formal differential privacy guarantees that apply to general settings without any assumptions on the learning algorithm. Experiments on standard FL datasets demonstrate that our approach outperforms strong baselines in both homogeneous and heterogeneous settings and under strict privacy constraints, while also providing characterizations of model uncertainties.

Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors' identity.

No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.

Submission Number: 6848