Federated Learning via Meta-Variational Dropout

TL;DR

This paper introduces MetaVD, a Bayesian meta-learning approach for federated learning that personalizes models, reduces overfitting, and improves communication efficiency in non-IID data settings.

Contribution

MetaVD predicts client-specific dropout rates using a shared hypernetwork, enhancing personalization and robustness in federated learning with limited non-IID data.

Findings

MetaVD achieves high classification accuracy and uncertainty calibration.

MetaVD effectively handles out-of-distribution clients.

MetaVD reduces communication costs by compressing local models.

Abstract

Federated Learning (FL) aims to train a global inference model from remotely distributed clients, gaining popularity due to its benefit of improving data privacy. However, traditional FL often faces challenges in practical applications, including model overfitting and divergent local models due to limited and non-IID data among clients. To address these issues, we introduce a novel Bayesian meta-learning approach called meta-variational dropout (MetaVD). MetaVD learns to predict client-dependent dropout rates via a shared hypernetwork, enabling effective model personalization of FL algorithms in limited non-IID data settings. We also emphasize the posterior adaptation view of meta-learning and the posterior aggregation view of Bayesian FL via the conditional dropout posterior. We conducted extensive experiments on various sparse and non-IID FL datasets. MetaVD demonstrated excellent classification accuracy and uncertainty calibration performance, especially for out-of-distribution (OOD) clients. MetaVD compresses the local model parameters needed for each client, mitigating model overfitting and reducing communication costs. Code is available at https://github.com/insujeon/MetaVD.