GEDI: A Graph-based End-to-end Data Imputation Framework

Katrina Chen, Xiuqin Liang, Zheng Ma, Zhibin Zhang

Abstract: Data imputation is an effective way to handle missing data, which is common in practical applications. In this study, we propose and test a novel data imputation process that achieves two important goals: 1) preserving the row-wise similarities among observations and column-wise contextual relationships among features in the feature matrix. 2) tailoring the imputation process to some specific downstream label prediction task. The proposed imputation process uses Transformer and graph structure learning to iteratively refine the contextual relationships among features and similarities among observations. Moreover, it implements a meta-learning framework to select features that are influential to the downstream prediction task of interest. We conduct experiments on real-world datasets, and show that the proposed method consistently improves imputation and label prediction performance over a variety of benchmark methods.