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.
Figures & tables
Figure 1: Comparison of imputation methods under varying missingness probabilities. The recovery ratio (left) shows that Joint-AmortizedImpute achieves the best overall performance, while the CD diagram (right) confirms statistically significant ranking differences.
Train → Test
Mean ( ↑ )
Win ( ↑ )
Degradation ( ↓ )
0.1 → 0.3
-0.0151
25.00%
47.22%
0.1 → 0.5
0.1240
74.07%
21.30%
0.1 → 0.7
0.0873
66.67%
28.21%
0.3 → 0.1
-0.0154
26.39%
43.06%
0.3 → 0.5
0.1428
81.94%
13.89%
0.3 → 0.7
0.1063
82.05%
12.82%
Table 1: Performance of Joint imputation relative to the best classical method under train–test missingness shifts. We report mean accuracy difference, win rate, and degradation rate across scenarios.
Figure 2: Comparison of imputation methods.
Train → Test
Mean ( ↑ )
Win ( ↑ )
Degradation ( ↓ )
MAR → MCAR
-0.0249
23.53%
55.29%
MAR → MNAR
0.1536
77.65%
18.82%
MCAR → MAR
-0.0134
24.32%
51.35%
MCAR → MNAR
0.1783
79.73%
13.51%
MNAR → MAR
-0.0267
23.46%
55.56%
MNAR → MCAR
0.1352
80.00%
15.00%
Table 2: Performance of Joint imputation relative to the best classical method under train–test missingness mechanism shifts. We report mean accuracy difference (Joint minus Best Classical), win rate, and degradation rate across scenarios.
Figure 3: Joint-AmortizedImpute: CD diagrams (scenario-level) for hard and severe scenarios. Lower ranks indicate better performance.
Figure 4: Joint-LatentImpute: CD diagrams (scenario-level) for hard and severe scenarios.
Joint-LatentImpute
Joint-AmortizedImpute
Method
RMSE Rank ( ↓ )
Cat. Err. Rank ( ↓ )
RMSE Rank ( ↓ )
Cat. Err. Rank ( ↓ )
Joint (Ours)
2.96
1.8
2.80
2.01
Column Mean
3.11
1.6
2.62
1.77
K-NN
3.39
1.6
3.15
–
Column Median
4.64
1.6
3.85
–
MissForest
3.02
2.0
4.82
3.12
Table 3: Imputation quality ranking (lower is better).
Figure 5: Overall: CD diagrams comparing imputation methods based on performance rankings.
Table 9
Appendix figures & tables7 assets
Supplementary material from the paper’s appendix.
Appendix
Table 10Table 11Figure 12
Factor
Sig.
η2
Imbalance
Yes
0.005
Missing Rate
No
0.0005
Dataset Size
Yes
0.05
#Categorical
Yes
0.017
Mechanism
Yes
0.015
#Classes
Yes
0.01
Appendix
Table 12: Factor impact on performance differences between Joint-AmortizedImpute and the best classical method.
Figure 14
Factor
Sig.
η2
Imbalance
No
0.002
Missing Rate
No
0.0002
Dataset Size
Yes
0.03
#Categorical
Yes
0.04
Mechanism
Yes
0.02
#Classes
No
0.003
Appendix
Table 13: Factor impact on performance differences between Joint-LatentImpute and the best classical method.
Method
mean ( ↓ )
median ( ↓ )
std ( ↓ )
Column Mean
0.0009
0.000
0.002
Column Median
0.0009
0.000
0.003
K-NN
0.692
0.195
1.4
MissForest
12.02
8.51
11.23
TabImpute
21.94
7.15
57.47
Joint-LatentImpute (Ours)
31.68
12.66
48.75
Appendix
Table 14: Average time (in seconds) required to obtain a fully imputed training dataset across all datasets. For iterative methods (MissForest and TabImpute), this includes model fitting.