PaperlayerView exact revision on arXiv

Source-linked AI summary

Fed-MIWAE: Federated Imputation of Incomplete Data via Deep Generative Models

Irene Balelli, Aude Sportisse, Francesco Cremonesi, Pierre-Alexandre Mattei, Marco Lorenzi

arXiv:2304.08054v1stat.MLcs.LG

TL;DR

Local preprocessing in federated learning can leave missing-data handling and standardization disconnected across sites, while federated imputation methods have limited support for MAR. The paper combines federated standardization with Fed-MIWAE and reports performance comparable to centralized methods despite heterogeneous data, with multiple imputation for uncertainty quantification.

  • Problem

    Preprocessing, including missing-data handling, is often performed locally in federated learning, and existing federated imputation methods do not adequately address MAR data.

    Source: §1 · p. 2

  • Method

    The paper federates standardization and trains Fed-MIWAE, a federated MIWAE model, using decentralized standardized datasets and a classical federated optimizer.

    Source: §1 · p. 3, §2.2 · p. 4

  • Results

    Fed-MIWAE outperforms tested local imputation methods and achieves performance comparable to centralized methods under simulated MAR and highly heterogeneous patient distributions.

    Source: §1 · pp. 3–4

  • Takeaways & Limitations

    Fed-MIWAE supports multiple imputation, allowing uncertainty in federated imputations to be quantified.

    Source: §3.4 · p. 8

Abstract

from arXiv · show

Federated learning allows for the training of machine learning models on multiple decentralized local datasets without requiring explicit data exchange. However, data pre-processing, including strategies for handling missing data, remains a major bottleneck in real-world federated learning deployment, and is typically performed locally. This approach may be biased, since the subpopulations locally observed at each center may not be representative of the overall one. To address this issue, this paper first proposes a more consistent approach to data standardization through a federated model. Additionally, we propose Fed-MIWAE, a federated version of the state-of-the-art imputation method MIWAE, a deep latent variable model for missing data imputation based on variational autoencoders. MIWAE has the great advantage of being easily trainable with classical federated aggregators. Furthermore, it is able to deal with MAR (Missing At Random) data, a more challenging missing-data mechanism than MCAR (Missing Completely At Random), where the missingness of a variable can depend on the observed ones. We evaluate our method on multi-modal medical imaging data and clinical scores from a simulated federated scenario with the ADNI dataset. We compare Fed-MIWAE with respect to classical imputation methods, either performed locally or in a centralized fashion. Fed-MIWAE allows to achieve imputation accuracy comparable with the best centralized method, even when local data distributions are highly heterogeneous. In addition, thanks to the variational nature of Fed-MIWAE, our method is designed to perform multiple imputation, allowing for the quantification of the imputation uncertainty in the federated scenario.

1 Introduction and previous work

Federated learning can combine decentralized hospital data, but preprocessing is commonly done locally, while existing federated imputation methods have limited support for MAR data. This work proposes federated standardization and Fed-MIWAE, reporting performance comparable to centralized methods even with heterogeneous local data.

  • Federated learning aggregates information from decentralized datasets, but data standardization and missing-entry management are generally performed locally at each site.Reliable federated preprocessing is identified as important for real-world deployment.Source: §1 · p. 2
  • Prior imputation approaches include probabilistic, low-rank, random-forest, multiple-imputation, GAN, and VAE methods, but none were specifically designed for federated learning.Multiple imputation can account for uncertainty in imputed values.Source: §1 · pp. 2–3
  • Fed-MIWAE achieves higher generalization than tested local imputation methods and performance comparable to centralized methods under simulated MAR and heterogeneous patient distributions.The authors also report that it can outperform the best centralized model in some settings.Source: §1 · pp. 3–4
  • The paper proposes federated standardization and Fed-MIWAE, a federated version of MIWAE based on variational inference.Fed-MIWAE is designed to handle MCAR and MAR data and allow multiple imputation.Source: §1 · p. 3

2 An approach to federated data pre-processing

MIWAE uses a latent-variable model and a variational bound to impute missing data. The proposed preprocessing pipeline first standardizes data federatively, then trains Fed-MIWAE across sites and imputes locally.

2.1 Missing data importance-weighted autoencoder (MIWAE)

  • MIWAE models data from a latent Gaussian variable and optimizes a stochastic lower bound on observed likelihood because missing values make that likelihood intractable.Its imputation estimates the conditional expectation of missing values given observed values via self-normalized importance sampling.Source: §2.1 · p. 4

2.2 Proposed method

  • Federated standardization shares site-level feature means and standard deviations in one communication round so each site can standardize against global statistics.Fed-MIWAE is then trained over standardized datasets across R communication rounds using a classical federated optimizer.Source: §2.2 · p. 4
  • After training, each center uses the final global Fed-MIWAE model to impute its local missing values.The model parameters can be aggregated with schemes such as FedAvg, FedProx, or Scaffold.Source: §2.2 · p. 4

3 Application to the Alzheimer’s Disease Neuroimaging Initiative dataset (ADNI)

The ADNI evaluation simulates federated sites with MAR missingness and compares Fed-MIWAE against local and centralized approaches. Fed-MIWAE performs consistently in heterogeneous settings, while also supporting multiple imputation for uncertainty quantification.

3.1 Data

  • The study uses 311 ADNI participants from 55 centers, with cognitive scores and MRI, FDG-PET, and AV45-PET features totaling 130 dimensions per fully observed patient.The cohort includes cognitively normal individuals and patients diagnosed with Alzheimer’s disease.Source: §3.1 · pp. 4–5

3.2 Scenarios

  • The evaluation simulates 30% missingness under MCAR or MAR, using natural-split and highly heterogeneous Not-IID scenarios with separate testing data.Natural-split experiments use four simulated datasets in repeated four-fold cross-validation; Not-IID experiments train on AD-only and NC-only clients and test on a mixed client.Source: §3.2 · pp. 5–6

3.3 Benchmark

  • The benchmark compares FedProx-trained and locally standardized FedProx MIWAE with local MIWAE, centralized MIWAE, mean imputation, random forests, and chained equations using normalized MSE.MIWAE hyperparameters are held constant across training settings.Source: §3.3 · p. 6

3.4 Results

  • In MAR tests, FedProx-trained Fed-MIWAE outperforms local MIWAE, and exceeds the best centralized model on Not-IID external testing; Random Forests lead the tested benchmark approaches.Fed-MIWAE also maintains similar MSE in Not-IID and natural-split settings, while local models drop on the external test set.Source: §3.4 · pp. 6–8, Figure 1 · p. 7
  • Fed-MIWAE supports multiple imputation by drawing samples from the estimated posterior, enabling imputation-uncertainty quantification.The authors report lower variance for the conditional posterior than for p(x_missing).Source: §3.4 · p. 8

4 Conclusion

The paper proposes a federated preprocessing pipeline for MCAR and MAR data, reporting comparable performance to centralized methods and accuracy in highly heterogeneous settings. It identifies MNAR handling, stronger privacy measures, and further study of clinical preprocessing pipelines as future directions.

4 Conclusion

  • The proposed pipeline handles MCAR and MAR values in federated learning and provides highly generalizable results comparable to centralized methods.The authors report accuracy in highly heterogeneous settings, which they say should be considered in real-world federated learning.Source: §4 · p. 8
  • The authors identify handling MNAR data and improving privacy with secure aggregation or differential privacy as future work.They also call for further exploration of federated preprocessing pipelines for real clinical studies.Source: §4 · p. 8
  • Figure 2 presents test-data multiple imputation in the ADNI MAR setting and a comparison of the conditional imputation distribution with the marginal distribution.Source: Figure 2 · p. 9

Supplementary Material

The supplementary material describes the ADNI sample and simulated-center sizes, Fed-MIWAE hyperparameters, and selected evaluation settings. It includes material for both MAR and MCAR scenarios.

Supplementary Material

  • Table S1 reports demographics of the ADNI clinical sample.Source: Table S1 · p. 13
  • Table S2 gives clinical sample sizes per simulated center for the natural split.Source: Table S2 · p. 13
  • Table S3 lists the Fed-MIWAE hyperparameters and defines C as the number of clients.Source: Table S3 · p. 13
  • Table S4 concerns the ADNI testing dataset under the MAR setting and natural split.Source: Table S4 · p. 13
  • Table S5 concerns the ADNI testing dataset under the MCAR setting and not-IID split.Source: Table S5 · p. 14
  • Figure S1 describes MCAR comparisons of federated and local models for unseen test data and training data, under natural-split and not-IID scenarios.The upper row is for the unseen testing dataset and the bottom row is for local training datasets.Source: Figure S1 · p. 14
Loading 2304.08054v1…