Personalizing Low-Rank Bayesian Neural Networks Via Federated Learning

TL;DR

LR-BPFL is proposed, a novel BPFL method that learns a global deterministic model along with personalized low-rank Bayesian corrections, and incorporates an adaptive rank selection mechanism to tailor the local model to each client's inherent uncertainty level.

Abstract

To support real-world decision-making, it is crucial for models to be well-calibrated, i.e., to assign reliable confidence estimates to their predictions. Uncertainty quantification is particularly important in personalized federated learning (PFL), as participating clients typically have small local datasets, making it difficult to unambiguously determine optimal model parameters. Bayesian PFL (BPFL) methods can potentially enhance calibration, but they often come with considerable computational and memory requirements due to the need to track the variances of all the individual model parameters. Furthermore, different clients may exhibit heterogeneous uncertainty levels owing to varying local dataset sizes and distributions. To address these challenges, we propose LR-BPFL, a novel BPFL method that learns a global deterministic model along with personalized low-rank Bayesian corrections. To tailor the local model to each client's inherent uncertainty level, LR-BPFL incorporates an adaptive rank selection mechanism. We evaluate LR-BPFL across a variety of datasets, demonstrating its advantages in terms of calibration, accuracy, as well as computational and memory requirements.

Figures (7)

• Figure 1: Average reliability diagrams across clients on the CIFAR-10 dataset for FedAvg mcmahan2017communication, a standard FL algorithm; Per-FedAvg fallah2020personalized, a benchmark PFL scheme; pFedBayes zhang2022personalized, a state-of-the-art BPFL algorithm; and the proposed LR-BPFL. (See Section \ref{['sec:uc']} for details.)
• Figure 2: Top: In the proposed LR-BPFL scheme, each client uploads its updated shared model to the server and then downloads the deterministic shared model from the server, while retaining and updating the Bayesian low-rank corrections locally. Bottom: During local training, after receiving the shared model from the server, clients perform local adaptation of the low-rank corrections before proceeding to update the shared model.
• Figure 3: Reliability diagrams of the worst-calibrated client on CIFAR-10 with 50 clients. The diagonal line represents perfect calibration. Each plot also displays the ECE and MCE values.
• Figure 4: Calibration performance for new clients as a function of the parameter $\alpha$, which dictates the degree to which the local data distributions of the test clients differ from those of the training clients. As $\alpha$ moves away from 0.1, the distributions of new clients diverge more from those of the training clients.
• Figure 5: Comparison of reliability diagrams for LR-BPFL with and without the ARS module.