Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Federated Learning via Meta-Variational Dropout

Insu Jeon, Minui Hong, Junhyeog Yun, Gunhee Kim

TL;DR

Federated learning struggles with non-IID client data and limited local data, which can lead to overfitting and poor global convergence. MetaVD reframes FL as amortized variational inference across clients, using a server-side hypernetwork to predict client-specific dropout rates $\alpha^m$ and a shared global parameter $\theta$, with local posteriors $q(w^m;\phi)$. A Bayesian aggregation rule combines per-client posteriors to form a MAP global parameter $\theta_*^{agg}$, weighted by an inverse relation to uncertainty, and the method employs a hierarchical prior that yields a tractable KL term. Empirical results across CIFAR-10/100, FEMNIST, CelebA, and multi-domain FL demonstrate improved accuracy and particularly enhanced uncertainty calibration for OOD clients, along with substantial model compression in the personalized layer and reduced communication costs, illustrating MetaVD's versatility and practicality for non-IID, data-limited FL scenarios.

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.

Federated Learning via Meta-Variational Dropout

TL;DR

Federated learning struggles with non-IID client data and limited local data, which can lead to overfitting and poor global convergence. MetaVD reframes FL as amortized variational inference across clients, using a server-side hypernetwork to predict client-specific dropout rates and a shared global parameter , with local posteriors . A Bayesian aggregation rule combines per-client posteriors to form a MAP global parameter , weighted by an inverse relation to uncertainty, and the method employs a hierarchical prior that yields a tractable KL term. Empirical results across CIFAR-10/100, FEMNIST, CelebA, and multi-domain FL demonstrate improved accuracy and particularly enhanced uncertainty calibration for OOD clients, along with substantial model compression in the personalized layer and reduced communication costs, illustrating MetaVD's versatility and practicality for non-IID, data-limited FL scenarios.

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.
Paper Structure (29 sections, 9 equations, 5 figures, 17 tables)

This paper contains 29 sections, 9 equations, 5 figures, 17 tables.

Table of Contents

  1. Introduction
  2. Background
  3. Approach: Meta-Variational Dropout
  4. Variational Inference for FL
  5. Meta-Variational Dropout
  6. Experiments
  7. Generalization on Non-IID Settings
  8. Ablation Study
  9. Uncertainty Calibration
  10. Client Participation
  11. Muti-domain Datasets
  12. Model Compression
  13. Conclusion
  14. The Global Posterior Decomposition in Bayesian FL
  15. The Motivation of Hierarchical Prior in MetaVD
  16. ...and 14 more sections

Figures (5)

  • Figure 1: Overview of the Meta-Variational Dropout algorithm. (a) The server's hypernetwork predicts client-specific dropout rates from client embedding, $e^m$. The global parameters, $\theta$, and dropout variables, $\alpha^m$, are sent to each $m$-th client. (b) A few-shot local adaptation is performed on each client's device in parallel; then, the updated parameters are sent back to the server. (c) The server then aggregates those parameters to update its variational parameters $\theta$, $\psi$, and $e$.
  • Figure 2: Visualization of client's data distribution in different non-IID degrees ($\dot{\alpha}=[0.5, 0.1]$).
  • Figure 3: Uncertainty calibration scores (ECE and MCE) measured on the OOD client in CIFAR-100 ($\dot{\alpha} = 0.1$). The lower, the better.
  • Figure 4: Reliability diagrams for (a) FedAvg, (b) FedAvg+MetaVD, (c) Reptile, (d) Reptile+MetaVD, (e) MAML, (f) MAML+MetaVD, (g) PerFedAvg and (h) PerFedAvg+MetaVD in CIFAR-100 ($\dot{\alpha}=0.5$).
  • Figure 5: Reliability diagrams for (a) FedAvg, (b) FedAvg+MetaVD, (c) Reptile, (d) Reptile+MetaVD, (e) MAML, (f) MAML+MetaVD, (g) PerFedAvg and (h) PerFedAvg+MetaVD in CIFAR-10 ($\dot{\alpha}=0.5$).