Personalized Federated Learning via Amortized Bayesian Meta-Learning

TL;DR

This paper introduces FedABML, a novel federated learning algorithm using Amortized Bayesian Meta-Learning to improve personalization across heterogeneous clients while providing theoretical guarantees and demonstrating superior empirical performance.

Contribution

It presents a new federated learning method leveraging hierarchical variational inference for personalized models with theoretical generalization bounds.

Findings

FedABML outperforms baseline methods in experiments.

Theoretical bounds on generalization error are established.

Hierarchical variational inference effectively captures client-specific structures.

Abstract

Federated learning is a decentralized and privacy-preserving technique that enables multiple clients to collaborate with a server to learn a global model without exposing their private data. However, the presence of statistical heterogeneity among clients poses a challenge, as the global model may struggle to perform well on each client's specific task. To address this issue, we introduce a new perspective on personalized federated learning through Amortized Bayesian Meta-Learning. Specifically, we propose a novel algorithm called \emph{FedABML}, which employs hierarchical variational inference across clients. The global prior aims to capture representations of common intrinsic structures from heterogeneous clients, which can then be transferred to their respective tasks and aid in the generation of accurate client-specific approximate posteriors through a few local updates. Our theoretical analysis provides an upper bound on the average generalization error and guarantees the generalization performance on unseen data. Finally, several empirical results are implemented to demonstrate that \emph{FedABML} outperforms several competitive baselines.