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

arXiv:2610.20342v1 Announce Type: new Abstract: 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.

Related

• When Does Self-Supervised Pretraining Help Tabular Models? A Study of Label Scarcity and Missing Data
• Integrated Imputation-Classification for Supervised Learning with Missing Data
• Missingness-aware Data Imputation via AI-powered Bayesian Generative Modeling
• Probabilistic indirect models for undrained shear strength: addressing significant data missing and variability with advanced imputation and machine learning techniques

Source: arXiv cs.LG | 2026-10-15