When Does Task-Aware Imputation Matter? An Analysis of Robustness Under Train-Test Mismatch
The paper investigates when jointly optimizing imputation with a downstream classifier improves robustness under train–test mismatches in missingness rates or mechanisms. Instead of treating imputation as fixed preprocessing, the authors cast missing values as learnable variables regularized around a reference imputation, giving both amortized and latent instantiations. Through 98 tabular classification datasets—and targeted shift experiments on 36 datasets—they show that reconstruction quality is a poor proxy for predictive utility and that task-aware corrections matter most in hard regimes such as MNAR and severe distribution shift.
Overall, the paper presents a well-motivated, large-scale empirical analysis that challenges the conventional practice of evaluating imputers via reconstruction error. The proposed framework is conceptually simple—essentially co-optimizing missing entries with a predictor anchored to a reference imputation—and is most compelling when train–test mismatch is severe, though its average-case benefits over a strong reference are modest.
The extensive benchmark across 98 datasets and multiple missingness mechanisms (MCAR, MAR, sequential MAR, and two MNAR variants) is a clear strength, and the systematic finding that reconstruction rank and predictive rank are only weakly correlated challenges standard imputation benchmarks. The reference-based parameterization—learning residual corrections around a strong baseline—stabilizes optimization and yields consistent gains when classical methods degrade, for example under large test-time missingness rate increases.
The baseline comparison is narrow: the authors compare only against mean, median, K-NN, MissForest, and TabImpute, omitting recent deep generative and task-aware competitors such as GRAPE, GEDI, and MIRACLE, which weakens claims of broad superiority in hard regimes. More importantly, the authors acknowledge that Joint-AmortizedImpute uses fixed hyperparameters across all datasets, and the main results show that when train–test shifts are small or absent the method often underperforms simple baselines and offers only marginal average improvements over its own reference initialization, raising practical questions about when to deploy it.
The central empirical claim—that reconstruction error is a poor proxy for downstream classification performance—is strongly supported by rank correlations and concrete counter-examples such as MissForest. However, the comparison to related task-aware work is largely rhetorical; the authors position their method as architecture-agnostic but do not run head-to-head experiments against graph-based end-to-end imputers or other joint-training baselines beyond simple classical methods, making it difficult to attribute gains to task-awareness itself versus the joint optimization procedure and reference anchoring.
Experimental details are thorough in the appendices: exact splits (70/15/15, seed 42), a lightweight classifier ($f_\theta(x)=W_2\sigma(W_1 x+b_1)+b_2$ with $h=32$), Gumbel-Softmax schedules, and explicit loss formulations are all provided. Yet the paper does not indicate whether code or dataset configs will be released, and the task-driven reference selection requires training and validating multiple classical imputers per dataset and missingness setting—a substantial computational step whose cost is not reported—creating a barrier to independent reproduction.
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.
Pick a starting point or write your own. Challenges run in the background, so you can keep reading while the AI investigates.
No challenges yet. Disagree with the review? Ask the AI to revisit a specific claim.