Task-aware imputation improves predictions under missing data differences
When Does Task-Aware Imputation Matter? An Analysis of Robustness Under Train-Test Mismatch
Machine Learning
Summary
Missing data can cause problems for machine learning models, especially when the pattern of missing information changes between training and testing. The authors show that filling in missing values while keeping the prediction task in mind works better than just guessing the missing data separately. They tested their approach on many datasets and found that better data guessing does not always mean better predictions, but their task-aware method helps when usual techniques fail under different missing data conditions.
What this means in practice
- •For data science teams: Improve model accuracy when deployed data has different missing value patterns than training data by integrating task-aware imputation methods.
- •For healthcare analytics teams: Enhance predictions on patient records with missing information that changes over time by using task-conditioned missing data imputation.
Authors
Maria V. Salvador, Ana F. Sequeira, Jaime S. Cardoso
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.