[Paper Review] Handling Missing Data with Graph Representation Learning

Jiaxuan You, Xiaobai Ma+3 more|arXiv (Cornell University)|Oct 30, 2020

Advanced Graph Neural NetworksComputer Science64 references94 citations

TL;DR

Grape uses a bipartite graph representation and Graph Neural Networks to perform both feature imputation and label prediction for missing data, achieving notable MAE improvements on nine UCI 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.

Motivation & Objective

• Motivate missing data problems and limitations of prior imputation and label-prediction methods.
• Propose a general graph-based framework (Grape) that handles both tasks end-to-end.
• Show that a bipartite graph with observations and features enables effective imputation and prediction.
• Demonstrate robustness, generalization to unseen data, and scalability across multiple datasets.

Proposed method

• Construct a bipartite graph with observations and features as node types and observed values as edges.
• Formulate feature imputation as edge-level prediction and label prediction as node-level prediction on the graph.
• Develop a Graph Neural Network architecture inspired by GraphSAGE that uses edge embeddings and augmented node features for bipartite graphs.
• Incorporate edge dropout during training to reduce overfitting and improve generalization to unseen edges.
• Use end-to-end training to optimize both O_edge (imputation) and O_node (prediction) components.

Experimental results

Research questions

• RQ1 Can a graph-based representation improve imputation quality for missing features compared to state-of-the-art methods?
• RQ2 Can end-to-end GNN-based learning simultaneously improve feature imputation and downstream label prediction under missing data?
• RQ3 Does edge dropout and augmented bipartite features enhance generalization to unseen data and observations?
• RQ4 Is Grape scalable to datasets with many features and capable of transferring knowledge across datasets?
• RQ5 How does Grape perform across varying missing data rates on benchmark datasets?

Key findings

• Grape achieves about 20% lower MAE for feature imputation and 10% lower MAE for label prediction on nine UCI datasets at 30% missing data.
• Edge dropout reduces MAE by about 33% on average for imputation tasks.
• End-to-end training improves downstream prediction MAE by about 19% in most cases.
• Grape generalizes well to unseen observations without retraining, showing ~21% MAE improvement over baselines.
• Grape maintains robust performance across different missing data levels (ratios) and supports scalable bipartite graph learning.