Personalizing Low-Rank Bayesian Neural Networks Via Federated Learning

TL;DR

This paper introduces LR-BPFL, a novel federated learning approach that personalizes Bayesian neural networks with low-rank corrections, improving calibration and efficiency across heterogeneous client datasets.

Contribution

LR-BPFL is the first method to combine low-rank Bayesian corrections with adaptive rank selection for personalized federated learning.

Findings

Enhanced calibration and accuracy demonstrated across multiple datasets.

Reduced computational and memory requirements compared to full Bayesian methods.

Effective personalization of uncertainty levels for diverse client data.

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.