An end-to-end distribution- and task-aware parallel imputation method for improving target task performance

Abstract

Most machine-learning imputation techniques treat missing values independently of the downstream task, resulting in suboptimal predictive performance. While some recent methods jointly train an imputer with a target predictor, they fail to produce diverse and context-sensitive imputations and suffer from training inefficiencies. Additionally, in many methods, missing values are ignored during training, and the imputers are trained only on observed data. To overcome these limitations, we propose an imputation method, named end-to-end task-aware parallel Imputation with class-wise prototypes (ETPI). It captures the class-conditional distributions of input data using a few proxies. For the missing entries in any sample, ETPI generates class-aware pseudo-targets based on the predictor loss and the estimated class-conditional distribution. By exposing the imputer to both observed data and optimal pseudo-targets during training, ETPI effectively leverages training information to fit the imputer model and aligns imputation with the objectives of the target task. Extensive experiments on classification and regression tasks show that ETPI outperforms other state-of-the-art methods. It also maintains high imputation quality even with limited training data or high missing rates, mainly due to the high quality of the generated pseudo-targets and the integration of imputation and prediction into a single end-to-end pipeline.