Classification of datasets with imputed missing values: does imputation quality matter?

1 Introduction

• The paper considers missingness arising from incomplete or inadequate collection, dataset corruption, and differences in recording across cohorts or sources.These sources can produce MCAR, MAR, or MNAR patterns.Source: §1 · p. 1
• Incomplete samples can reduce statistical power or introduce bias when missingness relates to the outcome, motivating imputation before classification.Many classification methods require complete data; the common two-stage approach imputes missing values and then classifies the completed dataset.Source: §1 · p. 1
• Downstream classification performance is often prioritized over checking whether imputed data reflects the underlying feature distribution, despite the risk of spurious feature importance.Common sample-wise metrics such as MSE, MAE, and square deviance can be optimal even when the imputed distribution differs from the true distribution.Source: §1 · p. 2
• The study evaluates two-stage imputation-and-classification configurations and introduces sliced Wasserstein-inspired discrepancy scores for overall distribution reconstruction.It also compares standard imputation-quality metrics and examines their relationship with downstream classification performance.Source: §1 · p. 2
• Popular discrepancy measures are uncorrelated with downstream model performance, whereas distributional discrepancy measures show a correlation.The authors also report that high-performing classifiers trained on poorly imputed data can assign spurious feature importances.Source: §1 · p. 2

2.1 Datasets

• Four datasets comprise three clinical datasets and one synthetic dataset; two complete datasets permit controlled induction of missingness, while two retain natural missingness.The Simulated dataset has 1,000 samples and 25 informative features.Source: §2.1 · p. 2

2.2 Dataset partitioning

• Development and holdout cohorts are separated before development data is split into training and validation cohorts to limit overfitting.For the complete MIMIC-III and Simulated datasets, the study induces 25% or 50% missingness in development and holdout data, yielding four rate combinations.Source: §2.2 · p. 3

2.3 Imputation methods

• The benchmark compares mean imputation, MICE, MissForest, GAIN, and MIWAE, with multiple-imputation predictions pooled after classification.The datasets include numerical features and, in two datasets, categorical and ordinal features; categorical variables are one-hot encoded and ordinal variables use integer codes.Source: §2.3 · p. 3

2.4 Classification methods

• The classifiers take completed samples as input and predict binary class labels using logistic regression, Random Forest, XGBoost, NGBoost, or a neural network.The classifier list and input/output description are stated in the methods passages.Source: §2.4 · pp. 3–4

2.5 Hyperparameter optimisation

• An exhaustive grid search selects classifier hyperparameters by average performance across five validation folds for each configuration.The search spans datasets, missingness levels, imputation methods, holdout and validation sets, and multiple imputations, requiring millions of experiments.Source: §2.5 · p. 4

2.6 ANOVA analysis

• A binomial logistic-regression-based generalized linear model performs multi-factor ANOVA on holdout AUC, including factor interactions before backward elimination at the 1% level.Results are pooled across the three clinical datasets for an overall real-clinical-data analysis, and each dataset is also assessed individually.Source: §2.6 · p. 4

2.7 Measuring the quality of imputation

• Imputation quality is evaluated with sample-wise RMSE, MAE, and R2; feature-wise KL, KS, and 2-Wasserstein measures; and a proposed whole-distribution class.The new class repeatedly projects data onto random one-dimensional directions and compares sliced Wasserstein distances; the method also examines relative distance change and repeated-imputation variance. Feature-wise marginals alone can miss differences in the joint distribution.Source: §2.7 · pp. 4–5

2.8 Downstream effects of imputation quality

• For downstream analysis, the study relates discrepancy scores to classifier performance and compares well-performing models trained on poorly versus well-imputed data using Shapley-value feature importances.The feature-importance analysis uses the Simulated dataset, whose features are equally important by design, and measures importance symmetry with skewness.Source: §2.8 · p. 6

3 Results

• No imputation method leads across real-world datasets: MIWAE leads on NHSX COVID-19, while MICE usually leads on MIMIC-III but is worst on Breast Cancer.Mean imputation is worst on MIMIC-III but competitive with MIWAE on NHSX COVID-19.Source: §3 · pp. 6–7
• Increasing test missingness reduces classification performance, and the complete-data classifier usually outperforms imputed-data classifiers on MIMIC-III and Simulated.The Neural Network on MIMIC-III is the exception; test missingness affects performance more than train missingness at fixed rates.Source: §3 · p. 6, Figure 3 · p. 7
• In pooled ANOVA, test-set missingness explains most performance deviance; classifier choice is primary for NHSX COVID-19 and Breast Cancer.For Breast Cancer, classifier choice is the only significant factor at the 1% level.Source: §3 · pp. 7–8, Figure 4 · p. 8

Comparing imputation quality

• Sample-wise scores favor MissForest on MIMIC-III, whereas feature-wise metrics favor MICE and rank mean worst; higher test missingness raises feature-wise distances.For sample-wise scores, MICE generally has the lowest R2 in both datasets.Source: Comparing imputation quality · pp. 8–9, Figure 5 · p. 9, Figure 6 · p. 9
• MICE has the best sliced-Wasserstein discrepancies overall in both datasets; MIWAE is competitive at 25% test missingness on Simulated data.MICE also induces a smaller distance ratio than other methods, which rises substantially with test missingness.Source: Figure 7 · p. 10, Comparing imputation quality · p. 10
• MICE imputations remain below the 10^-7 Wasserstein-distance threshold across holdout and validation sets, unlike methods with highly variable distances.Large distance variances can arise from random projections or stochasticity in the imputation algorithm.Source: Figure 8 · p. 10, Comparing imputation quality · p. 10
• Sample-wise discrepancies correlate with downstream performance only in the Simulated dataset; feature-wise and sliced-Wasserstein discrepancies correlate negatively in both datasets.In MIMIC-III, sample-wise scores show no clear correlation with classifier performance.Source: Figure 9 · p. 11, Comparing imputation quality · p. 11
• For high-performing classifiers, MICE yields less-skewed Shapley values than mean or GAIN on most evaluated Simulated-dataset features.The comparisons favor MICE on 21/25 features for Random Forest, 19/25 for NGBoost, and 15/25 for XGBoost.Source: Comparing imputation quality · pp. 11–12

4 Discussion

• No imputation or classifier consistently maximized downstream performance across the considered datasets; test-set missingness most influenced performance.The authors conclude that heuristically identifying an optimal method combination is unlikely.Source: §4 · p. 12
• Sample-wise scores do not establish distribution reconstruction, and feature-wise marginals alone are insufficient; sliced Wasserstein-based scores assess reconstruction across all features.The paper notes that MSE is statistically justified for imputing Gaussian data, an assumption that often does not hold in practice.Source: §4 · pp. 12–13
• MICE best recreated the overall distribution in both evaluable datasets across missingness rates, with the smallest relative distance increase and minimal variance.The authors report that higher test missingness substantially increased imputation-induced distance, while higher training missingness had a minimal effect.Source: §4 · p. 13
• GAIN performed well on sample-wise scores but poorly on feature-wise and sliced Wasserstein-based discrepancies; GAIN and MIWAE sometimes showed mode collapse or local minima across imputations.The authors recommend multiple imputation to reduce the risk of poor-quality results from neural-network methods.Source: §4 · p. 13
• Better imputation quality did not necessarily yield better classification performance, yet poor imputations compromised feature importances and could mislead conclusions about patient outcomes.The authors suggest, rather than establish, that noisy imputations might regularize a classifier.Source: §4 · p. 13
• Categorical imputation remains unresolved: one-hot values need valid binary and mutually exclusive configurations, which methods do not enforce; the released codebase supports benchmarking.The codebase includes preprocessing, partitioning, and analysis pipelines.Source: §4 · pp. 13–14

Comparison between MSE and Wasserstein distance as discrepancy scores

• MICE reflects the simulated data distribution better than an MSE-optimised imputation, despite having a poorer MSE.The example illustrates why MSE alone can be insufficient for assessing imputation quality.Source: Comparison between MSE and Wasserstein… · p. 18, Figure 11 · p. 18

Dataset Description and Preprocessing Details

• The MIMIC-III analysis uses first ICU admissions from patients over 15, with at least three ICU days, seven clinical variables, and 30-day survival as the outcome.The processed dataset contains 7,214 randomly selected patients; the supplementary material also details Breast Cancer preprocessing and categorical-variable handling.Source: Dataset Description and Preprocessing… · pp. 18–19

Descriptions of the imputation methods

• The study describes mean imputation, MissForest, MICE, GAIN, and MIWAE, including their mechanisms and handling of categorical or ordinal variables.For example, MICE iteratively models each variable in turn, while MIWAE uses a deep latent variable model and permits multiple imputation through sampling importance resampling.Source: Descriptions of the imputation methods · pp. 19–20

Descriptions of the classification methods

• The benchmark tunes logistic regression, random forest, XGBoost, NGBoost, and neural-network classifiers; five-fold development-set cross-validation selects hyperparameters by mean AUC.Each dataset is split into three holdout sets and complementary development sets, with the latter divided into five validation cohorts.Source: Descriptions of the classification… · pp. 20–21, Data partitioning and hyperparameter… · p. 21

Description of the sample-wise and feature-wise discrepancy measures

• The discrepancy measures include sample-wise, feature-wise, and sliced-Wasserstein classes, with development-set means and standard deviations used to normalize data.The supplementary material describes implementations including RMSE, MAE, R2, KL divergence, KS, and two-Wasserstein distance.Source: Description of the sample-wise and… · p. 22

Outlier Analysis Details

• GAIN produced Wasserstein distances above 10^-7 in 80% of imputations and above 10^-6 in 40%, while MIWAE showed less stochasticity.The authors recommend multiple imputation runs for stochastic methods; they note that deep networks can find local minima and GANs can undergo mode collapse.Source: Outlier Analysis Details · p. 22, Figure 13 · p. 23

Interpretability

• For interpretability comparisons, the authors select each classifier’s best-performing configurations, then compare models with the smallest and largest sliced-Wasserstein distance ratios using Shapley values.The supplementary material also reports Accuracy, Brier Score, Precision, Sensitivity, and Specificity for the Simulated and MIMIC-III datasets at 25% and 50% development and test missingness.Source: Interpretability · p. 22, Additional performance metrics for the… · p. 24, Figure 15 · pp. 24–26

MIMIC-III

• Figure 18 presents sample-wise statistics for MIMIC-III at the train and test missingness rates considered.The caption does not state the direction or size of differences between settings.Source: Figure 18 · p. 29, Supplementary Figures for the Sliced… · p. 38

Simulated dataset

• Figure 19 presents sample-wise statistics for the Simulated dataset at the train and test missingness rates considered.The caption does not state the direction or size of differences between settings.Source: Figure 19 · p. 30

25% train, 50% test

• Figures 20 and 24 cover 25% train missingness with 50% test missingness.The captions identify the missingness rates but do not report an outcome.Source: Figure 20 · p. 31, Figure 24 · p. 35

50% train, 25% test

• Figures 21 and 25 cover 50% train missingness with 25% test missingness.The captions identify the missingness rates but do not report an outcome.Source: Figure 21 · p. 32, Figure 25 · p. 36

50% train, 50% test

• Figures 22 and 26 cover 50% train missingness with 50% test missingness.The captions identify the missingness rates but do not report an outcome.Source: Figure 22 · p. 33, Figure 26 · p. 37

25% train, 25% test

• Figure 23 covers 25% train missingness with 25% test missingness.The caption identifies the missingness rates but does not report an outcome.Source: Figure 23 · p. 34

MIMIC-III

• Figure 27 reports class C sliced Wasserstein discrepancies for MIMIC-III, with training missingness at 25% and test missingness at 25% and 50%; original values and logarithms are shown for clarity.Source: Figure 27 · p. 38
• Figure 28 covers MIMIC-III class C discrepancies with 50% train missingness and 25% or 50% test missingness.The original values and logarithms are shown for clarity.Source: Supplementary Figures for the Sliced… · p. 38, Figure 28 · p. 39

Simulated dataset

• Figure 29 reports class C discrepancies in sliced Wasserstein distance for simulated data, with training missingness fixed at 25% and test missingness at 25% and 50%.The original values and logarithms are shown for clarity.Source: Figure 29 · p. 40
• Figure 30 covers Simulated-data class C discrepancies with 50% train missingness and 25% or 50% test missingness.The original values and logarithms are shown for clarity.Source: Figure 30 · p. 41, C: Sliced Wasserstein discrepancy for… · p. 41

• Figure 31 plots class A, B, and C discrepancy metrics against downstream AUC for the Simulated dataset, with separate trend lines for 25% and 50% test missingness.Source: Figure 31 · p. 42
• Figure 32 shows the imputed-to-original Wasserstein-distance ratio for MIMIC-III across train and test missingness rates.Source: Figure 32 · p. 43
• Figure 33 shows the imputed-to-original Wasserstein-distance ratio for the Simulated dataset across train and test missingness rates.Source: Figure 33 · p. 43

• Table 2 lists parameter coefficients for the Breast Cancer dataset.Source: Table 2 · p. 44
• Table 3 lists parameter coefficients for the NHSX COVID-19 dataset.Source: Table 3 · p. 44
• Table 4 lists parameter coefficients for the MIMIC-III dataset.Source: Table 4 · p. 44

• Table 5 lists parameter coefficients for the Simulated dataset.Source: Table 5 · p. 45
• Table 6 gives the model configurations used in the Random Forest interpretability analysis.Source: Table 6 · p. 46
• Table 7 gives the model configurations used in the XGBoost interpretability analysis.Source: Table 7 · p. 46
• Table 8 gives the model configurations used in the NGBoost interpretability analysis.Source: Table 8 · p. 46

• Figure 34 is a correlation heatmap of all discrepancy metrics considered for the Simulated dataset.Source: Figure 34 · p. 47