Handling Missing Data with Graph Representation Learning

TL;DR

GRAPE introduces a graph neural network framework for missing data handling, enabling effective feature imputation and label prediction by modeling data as a bipartite graph, outperforming existing methods on benchmark datasets.

Contribution

This paper presents GRAPE, a novel graph-based framework that unifies feature imputation and label prediction for missing data using bipartite graph modeling and GNNs.

Findings

20% lower mean absolute error for imputation

10% lower error for label prediction

Outperforms state-of-the-art methods on nine datasets

Abstract

Machine learning with missing data has been approached in two different ways, including feature imputation where missing feature values are estimated based on observed values, and label prediction where downstream labels are learned directly from incomplete data. However, existing imputation models tend to have strong prior assumptions and cannot learn from downstream tasks, while models targeting label prediction often involve heuristics and can encounter scalability issues. Here we propose GRAPE, a graph-based framework for feature imputation as well as label prediction. GRAPE tackles the missing data problem using a graph representation, where the observations and features are viewed as two types of nodes in a bipartite graph, and the observed feature values as edges. Under the GRAPE framework, the feature imputation is formulated as an edge-level prediction task and the label prediction as a node-level prediction task. These tasks are then solved with Graph Neural Networks. Experimental results on nine benchmark datasets show that GRAPE yields 20% lower mean absolute error for imputation tasks and 10% lower for label prediction tasks, compared with existing state-of-the-art methods.