# When Does Task-Aware Imputation Matter? An Analysis of Robustness Under Train-Test Mismatch

arXiv: https://arxiv.org/abs/2610.20342 · HTML: https://mlsanity.com/paper/2610.20342/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 14 of its 35 references that are in the map.
- First posted on arXiv 14 Oct 2026.

## Builds on (most cited first, 14 of 14)

- [Revisiting Deep Learning Models for Tabular Data](https://mlsanity.com/paper/2106.11959/) (2021)
- [GAIN: Missing Data Imputation using Generative Adversarial Nets](https://mlsanity.com/paper/1806.02920/) (2018)
- [Handling Incomplete Heterogeneous Data using VAEs](https://mlsanity.com/paper/1807.03653/) (2018)
- [Handling Missing Data with Graph Representation Learning](https://mlsanity.com/paper/2010.16418/) (2020)
- [Missing Data Imputation using Optimal Transport](https://mlsanity.com/paper/2002.03860/) (2020)
- [HyperImpute: Generalized Iterative Imputation with Automatic Model Selection](https://mlsanity.com/paper/2206.07769/) (2022)
- [OpenML-Python: an extensible Python API for OpenML](https://mlsanity.com/paper/1911.02490/) (2019)
- [Generating and Imputing Tabular Data via Diffusion and Flow-based Gradient-Boosted Trees](https://mlsanity.com/paper/2309.09968/) (2023)
- [MIRACLE: Causally-Aware Imputation via Learning Missing Data Mechanisms](https://mlsanity.com/paper/2111.03187/) (2021)
- [ReMasker: Imputing Tabular Data with Masked Autoencoding](https://mlsanity.com/paper/2309.13793/) (2023)
- [Classification of datasets with imputed missing values: does imputation quality matter?](https://mlsanity.com/paper/2206.08478/) (2022)
- [Domain Adaptation under Missingness Shift](https://mlsanity.com/paper/2211.02093/) (2022)
- [GEDI: A Graph-based End-to-end Data Imputation Framework](https://mlsanity.com/paper/2208.06573/) (2022)
- [missForestPredict -- Missing data imputation for prediction settings](https://mlsanity.com/paper/2407.03379/) (2024)

## Abstract (arXiv)

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.