Source-linked AI summary
When Does Task-Aware Imputation Matter? An Analysis of Robustness Under Train-Test Mismatch
Maria V. Salvador, Ana F. Sequeira, Jaime S. Cardoso
TL;DR
Imputation is often separated from prediction, despite missingness potentially differing between training and deployment. This paper jointly optimizes missing values with the predictive model and studies amortized and latent strategies. Task-aware imputation is most beneficial in challenging regimes, while reconstruction quality is only weakly correlated with downstream performance.
Problem
Most imputation methods treat missing data as preprocessing without accounting for when or how missingness occurs, including train–test mismatches.
Source: §1 · pp. 1–2Method
The paper unifies task-aware imputation by treating missing values as learnable variables jointly optimized with the predictive model.
Source: §2 · p. 3Results
Task-aware imputation benefits challenging regimes, particularly distribution shift and MNAR settings where classical methods fail to recover complete-data performance; reconstruction quality correlates weakly with downstream performance (Spearman ρ ≈0.16–0.32).
Source: §6 · p. 9Takeaways & Limitations
Imputation should be treated as task-dependent rather than as preprocessing focused solely on reconstruction.
Source: §6 · p. 9Experiments cover tabular classification, and the method does not explicitly model the missingness mechanism.
Source: §6 · p. 10
Abstract
from arXiv · showhide
Missing data is a pervasive challenge in machine learning, often degrading predictive performance and complicating model deployment. Most existing approaches treat imputation as a preprocessing step independent of the downstream task, a separation that becomes problematic under train-test mismatches in missingness mechanisms or rates. We introduce a unified framework that treats missing entries as variables in a task-conditioned optimization problem. We instantiate this framework in two settings: an amortized approach for scenarios with missing data at both training and test time, and a latent optimization approach for settings where only training data is incomplete. The evaluation is conducted on 98 tabular datasets, and robustness under train-test mismatches is assessed on 36 datasets. Results reveal that reconstruction error constitutes a poor proxy for predictive utility (Spearman rho approximately 0.16 to 0.32), and that task-aware imputation becomes critical in scenarios where classical methods fail to recover complete-data performance, achieving win rates of up to 80% in challenging train-test mismatch scenarios.
1 Introduction
Missingness can undermine predictive models, while its timing and mechanism create distinct challenges. The work asks when task-aware imputation improves prediction, especially under train–test mismatch.
1 Introduction
- Missing values can bias estimates, reduce statistical power, and degrade performance, with implications for predictive-model validity, reliability, and fairness.Source: §1 · p. 1
- MCAR, MAR, and MNAR differ in whether missingness depends on data, observed values, or unobserved values, respectively.These mechanisms interact with when missingness occurs in the learning pipeline.Source: §1 · p. 1
- Most imputation methods treat missing data as static preprocessing without accounting for when or how missingness occurs.Missingness at training only primarily affects representation learning, while missingness at test time requires robustness to distribution shift.Source: §1 · pp. 1–2
- The paper investigates when task-aware imputation improves predictive performance, particularly when missingness persists at test time or differs between training and deployment.Source: §1 · p. 2
- Task-aware imputation is not universally beneficial, but becomes critical in challenging regimes such as distribution shift and intrinsically hard missingness settings.The paper also introduces a joint optimization formulation and reports weak correlation between reconstruction metrics and downstream performance.Source: §1 · p. 2
2 Related Work
Prior work spans classical, generative, and task-aware imputation, alongside research on missingness shifts. The paper unifies task-aware formulations and studies train–test mismatch directly.
2 Related Work
- Handling missing data includes imputation, omission, and methods that operate directly on incomplete data; this work focuses on imputation.Source: §2 · p. 2
- Classical and iterative methods, low-rank completion, pretrained models, and deep tabular methods provide varied approaches to imputation.The passage notes trade-offs including limited modeling of complex dependencies, computational expense, and sensitivity to design choices.Source: §2 · p. 2
- Deep generative methods often prioritize data fidelity, while other work models MNAR mechanisms or aligns observed and missing-data distributions.Examples include MIRACLE and optimal transport-based approaches.Source: §2 · p. 2
- Task-aware methods range from clustering-oriented imputation to end-to-end prediction frameworks, sometimes with specialized graph-based designs.TARImpute, OptImpute, GRAPE [33], and GEDI [34] illustrate these approaches.Source: §2 · p. 2
- Missingness shifts can invalidate standard assumptions and prompt adaptation, while mask-shift work studies removing unstable mask correlations [3, 4].Source: §2 · p. 2
- Earlier work identifies robustness to missingness shifts and misalignment between reconstruction quality and predictive performance as key challenges.Two-stage pipelines separate imputation from prediction, and RMSE and MAE may not capture discrepancies in the underlying data distribution.Source: §2 · p. 2
- The proposed formulation jointly trains missing values with the predictive model and investigates train–test mismatches in missingness mechanisms and rates.It encompasses latent and amortized strategies without relying on specialized architectures or bilevel optimization.Source: §2 · p. 3
3 Joint-AmortizedImpute: Task-Aware Imputation via Joint Optimization
The method treats unobserved values as learnable variables jointly optimized with the predictive model. A reference imputation and penalty regularize this optimization, which supports amortized and per-sample latent variants.
3 Joint-AmortizedImpute: Task-Aware Imputation via Joint Optimization
- The dataset consists of n samples with feature vectors x_i, some of whose components may be missing.The indicator m_i marks observed components with 1 and missing components with 0.Source: §3 · p. 3
- The method jointly optimizes the predictive model and unobserved input components under a task-oriented objective.Observed and unobserved components are distinguished, and the objective minimizes expected predictive loss.Source: §3 · p. 3
- Missing values are represented as deviations from a reference imputation, which anchors optimization in a plausible region and reduces the effective search space.A quadratic regularizer can encourage imputations to remain close to the reference.Source: §3 · pp. 3–4
- A unified imputation operator parameterized by ψ yields a common objective for different task-aware imputation strategies.The operator specifies how imputations are computed.Source: §3 · p. 4
- Joint-AmortizedImpute uses shared parameters to predict imputations in one forward pass, while Joint-LatentImpute directly optimizes per-sample latent variables.The amortized variant is described as enabling generalization to unseen data and reducing computational cost; the latent objective sets λrec = 0.Source: §3 · p. 4
- The amortized objective includes penalty and reconstruction-loss terms with nonnegative weights.The implementation details for both variants are provided in the appendices.Source: §3 · p. 4
4 Experimental Setup
The experiments cover classification with missingness at inference and training-only missingness, using benchmarks, baselines, and recovery-oriented metrics. Robustness tests vary missingness mechanisms and rates across training and test data.
4 Experimental Setup
- Joint-AmortizedImpute is evaluated with missingness at inference, including train–test shifts in missing rates or mechanisms and same-mechanism stability scenarios.The complementary Joint-LatentImpute setting tests incomplete training data with clean deployment data.Source: §4 · p. 4
- The stability evaluation spans six missingness mechanisms and five missing rates, while the training-only setting uses the same mechanism and rate grid.Source: §4 · pp. 4–5
- Baselines include simple statistical imputations, K-NN, MissForest, TabImpute, and a complete-data performance upper bound.The reference imputation is selected per dataset and missingness setting by validation performance among candidate imputers.Source: §4 · p. 5
- Imputation quality uses RMSE for numerical features and categorical error, while predictive utility uses accuracy or F1 and gaps from complete-data performance.Recovery ratio compares accuracy under missing data with the corresponding complete-data accuracy; scenarios are categorized into four hardness levels.Source: §4 · p. 5
- The evaluation covers 98 OpenML classification datasets overall and 36 datasets for distribution-shift experiments, with identical data splits across methods.The shift subset excludes datasets with imbalance ratio at or above 10; numerical features are standardized using training-set observed values.Source: §4 · p. 5
- Methods are compared with Friedman and Nemenyi tests, and reconstruction quality versus downstream performance is analyzed using Spearman rank correlation.Critical difference diagrams show average ranks; tests use α = 0.05.Source: §4 · pp. 5–6
5 Results
Task-aware imputation performs best in challenging settings, especially when test missingness substantially exceeds training missingness, but its advantage is not universal. Reconstruction quality is only weakly associated with predictive performance, while both proposed approaches recover complete-data performance well on average.
5.1.1 Train/test mismatch
- For 0.1 →0.5 and 0.3 →0.7 missing-rate shifts, Joint-AmortizedImpute gains exceed 0.10 on average and win in over 80% of cases, while small shifts favor the best classical baseline.Across scenarios, its recovery ratio is 0.8181 versus 0.69–0.70 for classical baselines.Source: §5.1.1 · p. 6, Figure 1 · p. 6, Table 1 · p. 6
- Under mechanism shifts, Joint-AmortizedImpute attains a 0.8381 recovery ratio versus K-NN at 0.7220 and outperforms classical methods on MNAR tests across training mechanisms.Its results are weaker under MAR, while MNAR-trained models generalize well to MCAR.Source: §5.1.1 · pp. 6–7, Figure 2 · p. 7, Table 2 · p. 7
5.1.2 Difficult scenarios not due to train/test mismatch
- In hard and severe scenarios, Joint-AmortizedImpute ranks lowest for accuracy and F1, while Joint-LatentImpute is top-performing on both metrics.Joint-AmortizedImpute is not significantly different from Column Mean for accuracy but is significantly better than the remaining methods for F1; Joint-LatentImpute’s results are significant at all p < 10^-6.Source: §5.1.2 · p. 7, Figure 3 · p. 8, Figure 4 · p. 8
5.2 Reconstruction error & its relationship to downstream performance
- Reconstruction quality is only weakly correlated with downstream accuracy: Spearman ρ = 0.16 for Joint-LatentImpute and ρ = 0.32 for Joint-AmortizedImpute.MissForest ranks second for numerical reconstruction error, yet mean and K-NN imputation often perform better in downstream classification.Source: §5.2 · p. 8, Table 3 · p. 8
5.3 Overall stability assessment
- Joint-AmortizedImpute achieves the highest average recovery ratio, with only marginal changes relative to the selected reference and no significant difference from top-ranked classical methods.For accuracy, Nemenyi p-values are 0.38 and 0.96 versus Column Mean and Column Median; for F1, p = 0.11 versus Column Mean.Source: §5.3 · pp. 8–9, Table 4 · p. 9, Table 5 · p. 9
- Joint-LatentImpute has the highest average recovery across scenarios, improves over the selected reference in approximately 45% of cases, and degrades by more than 5% in 2% and more than 10% in 0.3%.The reported results characterize large performance drops as rare.Source: §5.3 · p. 9, Table 6 · p. 9, Table 7 · p. 9
6 Conclusions, Limitations and Future Work
Task-aware imputation benefits predictive performance mainly in challenging regimes, while reconstruction quality is only weakly correlated with downstream performance. The paper also notes methodological and scope limitations and potential fairness risks.
Conclusions
- Task-aware imputation primarily benefits challenging regimes, particularly distribution shift and MNAR settings where classical methods fail to recover complete-data performance.Reconstruction quality is only weakly correlated with downstream performance (Spearman ρ ≈0.16–0.32).Source: §6 · p. 9
Limitations and Future Work
- Joint-AmortizedImpute uses fixed hyperparameters, compares against a small set of baselines, and is evaluated only on tabular classification without explicitly modeling missingness mechanisms.The authors identify dataset-specific tuning, broader benchmarking, other task settings, and mechanism-aware components as future directions.Source: §6 · p. 10
Broader impacts
- Task-aware imputation may support better decisionmaking but may also amplify observed-data biases, risking unfair or misleading outcomes in sensitive domains.The authors call for fairness-aware imputation and safeguards.Source: §6 · p. 10
A Experimental setup
Experiments used Python 3.10.19 and standard machine-learning and data-analysis libraries, with GPU hardware.
Experimental setup
- Experiments used Python 3.10.19, PyTorch, scikit-learn, NumPy, pandas, and the tabimpute library; accuracy and F1-score were computed with scikit-learn.Tabular foundation models were incorporated through tabimpute.Source: A Experimental setup · p. 12
- Experiments ran on a machine equipped with an NVIDIA GPU with CUDA support.Source: A Experimental setup · p. 12
A.1 Datasets and data preprocessing
The main evaluation uses 98 OpenML classification datasets without missing values, enabling controlled missingness generation with complete-data ground truth. A separate 36-dataset set supports distribution-shift experiments.
Datasets and data preprocessing
- The evaluation is restricted to OpenML classification datasets without missing values, allowing controlled missingness generation with complete-data ground truth.Source: §A.1 · p. 12
- 98 datasets span domains, dataset sizes, feature dimensionalities, categorical-feature proportions, and binary and multiclass classification tasks.Tables 8 and 9 summarize dataset characteristics and manually annotated domains.Source: §A.1 · p. 12, Table 8 · p. 13, Table 9 · p. 13
- The distribution-shift experiments use 36 datasets selected for stable evaluation under cross-mechanism and varying missing-rate scenarios while maintaining domain diversity.Tables 10 and 11 report their characteristics and domains.Source: §A.1 · p. 15, Table 10 · p. 15, Table 11 · p. 15
A.2 Datasets with reported win-rate results
The listed datasets include examples from several domains, but these passages do not report win-rate values or comparisons.
Datasets with reported win-rate results
- The list includes compass, electrical-grid-stability, and covertype datasets.Source: §A.2 · p. 15
- Additional listed datasets include cpu_small, diabetes risk prediction, compas-two-years, and balance-scale.Source: §A.2 · p. 16
A.3 Data splits and missingness
The appendix describes the benchmark setup and two task-aware approaches: Joint-LatentImpute optimizes missing entries with the classifier, while Joint-AmortizedImpute learns an imputer for use with the classifier.
A.3 Data splits and missingness
- Training, validation, and test sets use stratified 70/15/15 splits with seed 42, and missingness is generated independently after splitting; additional mechanisms model sequential dependence, latent factors, and lower-tail censoring.The structured mechanisms are MAR Sequential, MNAR latent factor, and MNAR Polarization.Source: §A.3 · p. 16
A.4 Benchmark imputation methods
- Five imputers—mean, median, K-NN, MissForest, and TabImpute—are benchmarked, with the highest validation-accuracy method selected for each dataset and missingness configuration.Imputation quality is assessed on missing entries against ground truth.Source: §A.4 · p. 17
A.5 Classifier
- A single-hidden-layer ReLU classifier with hidden dimension 32 is kept lightweight so performance differences primarily reflect imputation strategy rather than model capacity..Source: §A.5 · p. 17
A.6 Joint-LatentImpute
- Joint-LatentImpute initializes missing values from the best baseline, then learns numerical offsets and categorical logits while jointly optimizing imputations and classifier parameters.Regularization keeps imputations close to the reference solution.Source: §A.6 · p. 18
A.7 Joint-AmortizedImpute
- Joint-AmortizedImpute uses a learned residual correction to a reference imputation, preserves observed entries, and trains with artificial masking alongside classifier prediction.Its model parameter count depends on input dimensionality and architecture, not the number of training samples.Source: §A.7 · pp. 19–20
B.1 Overall stability
Across datasets, both joint methods show robust accuracy relative to classical and complete-data baselines as missingness increases. Dataset characteristics explain only a limited portion of performance differences, and the approaches incur different computational costs.
B.1 Overall stability
- For Joint-AmortizedImpute, dimensionality, dataset size, missingness mechanism, class count, and imbalance are significant factors with small effects, while missing rate is not significant.For Joint-LatentImpute, categorical-feature count, dataset size, and missingness mechanism are significant; imbalance, missing rate, and class count are not, and effects are small.Source: §B.1 · p. 21, Table 12 · p. 21, Table 13 · p. 22
- Joint-AmortizedImpute shows few extreme test-accuracy deviations from the best classical method and the smallest accuracy gap to complete data as missingness rises.Figure 6 shows differences from the best classical method; Figure 7 shows the gap to the complete-data baseline by missing rate.Source: §B.1 · p. 21, Figure 6 · p. 21, Figure 7 · p. 21
- Joint-LatentImpute stays close to the best classical method with fewer extreme degradations, and has the smallest accuracy gap to complete data as missingness increases.Figure 8 compares accuracy with the best classical method; Figure 9 plots the gap to complete data by missing rate.Source: §B.1 · pp. 21–22, Figure 8 · p. 22, Figure 9 · p. 22
B.2 Computational cost
- Joint-AmortizedImpute takes a mean 39.3 seconds to train but imputes at inference with one forward pass; Joint-LatentImpute averages 31.7 seconds to impute training data.The latter timing reflects end-to-end optimization; iterative baseline timings include fitting where applicable.Source: §B.2 · p. 22, Table 14 · p. 23