Scalable Missing Data Imputation With Graph Neural Networks

Missing features in tabular and graph structured data are common: a company may not want to disclose all of their accounting, and users online do not always engage in social platforms in the same way as their peers. Recently, models such as the GRAPE architecture have achieved state-of-the-art results in the task of feature imputation. We present an extension of GRAPE that performs mini-batch learning on datasets which do not fit in the GPU. Moreover, we add preprocessing and post-processing steps that allow the model to be used with graph structured data. We experimentally show the behaviour of the model on an academic citation network under different regimes of missing data. We observe that the model performance starts decreasing when we have less than 1% of observed edges. We additionally perform an ablation study of key elements of the model, such as its capacity, the batch size and the number of layers.