Source-linked AI summary
Handling Missing Data with Graph Representation Learning
Jiaxuan You, Xiaobai Ma, Daisy Yi Ding, Mykel Kochenderfer, Jure Leskovec
TL;DR
Missing-data learning requires either imputing absent features or predicting downstream labels directly, while existing methods face restrictive assumptions, heuristics, or scalability limitations. GRAPE represents observations and features in a bipartite graph and uses GNNs for edge-level imputation and node-level label prediction. Across nine benchmark datasets, it reports lower MAE for both tasks and generalizes to unseen observations without retraining.
Problem
Existing imputation methods rely on strong assumptions and cannot learn from downstream tasks, while label-prediction methods may use heuristics and face scalability issues.
Method
GRAPE represents observations and features as bipartite-graph nodes, observed values as attributed edges, and solves imputation and label prediction with a GNN.
Results
20% lower MAE for imputation and 10% lower MAE for prediction were reported at a 30% missing-data rate across nine UCI benchmark datasets.
Takeaways & Limitations
GRAPE provides one end-to-end graph framework for both missing-feature imputation and downstream label prediction, with robust generalization to unseen data points.
Takeaways & Limitations
Deep-learning imputers can fail to use feature values from other observations and may bias imputations through special default-value initialization.
Abstract
from arXiv · showhide
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.
1 Introduction
Missing-data methods separately address feature imputation or downstream label prediction, but existing approaches face assumptions, limited flexibility, heuristics, or scalability issues. GRAPE represents incomplete data as a bipartite graph and uses GNNs to solve both tasks, outperforming state-of-the-art methods on nine benchmarks.
- Motivation: Existing missing-data methods either estimate missing features or predict labels directly from incomplete inputs.Feature imputation estimates missing feature values from observed values, whereas label prediction targets downstream classification or regression.
- Motivation: Statistical imputation methods make strong distributional assumptions, struggle with mixed data types, and may require retraining for unseen samples.The passage specifically identifies these shortcomings for statistical and matrix-completion approaches.
- Motivation: Deep generative imputers often overlook feature values from other observations and introduce bias through special default-value initialization.These limitations are attributed to existing GAN- and autoencoder-based approaches.
- GRAPE framework: GRAPE models observations and features as bipartite-graph node types, with observed feature values represented as attributed edges.This representation turns feature imputation into edge-level prediction and label prediction into node-level prediction.
- GRAPE framework: GRAPE uses a GNN with edge embeddings, augmented node features, and edge dropout to support mixed edge information, inductive learning, and reduced overfitting.The architecture is designed for continuous and discrete edge features and end-to-end learning of both tasks.
- Results: 20% lower MAE for imputation and 10% lower MAE for prediction were achieved at a 30% missing-data rate across nine UCI benchmark datasets.The authors also report superior performance on unseen observations without retraining.
2 Related Work
Prior work includes statistical, deep-learning, and tree-based approaches for incomplete data, but these methods have limitations in flexibility, cross-observation modeling, efficiency, or scalability. GRAPE instead provides a graph-based framework covering both imputation and label prediction.
- Feature imputation: Statistical imputation methods include expectation-maximization, MICE, KNN, and matrix completion, but often rely on distributional assumptions or lack mixed-data flexibility.The related-work discussion identifies these limitations across joint modeling and matrix-completion approaches.
- Feature imputation: Deep autoencoder and GAIN imputers use only a single observation, whereas GRAPE captures interactions among multiple observations and features.GRAPE’s graph representation explicitly models these cross-observation and cross-feature interactions.
- Label prediction: Tree-based label-prediction methods use surrogate splitting heuristics, while random forests additionally suffer scalability issues from their multiple-tree structure.The passage contrasts these approaches with GRAPE’s direct graph handling of missing entries.
- GRAPE’s contribution: GRAPE’s main contribution is a graph-based framework that addresses both feature imputation and label prediction, rather than a particular GNN architecture.The authors describe it as the first graph-based solution covering both aspects of missing-data problems.
3 The GRAPE Framework
GRAPE represents incomplete feature matrices as bipartite graphs and uses GNN message passing to support imputation and label prediction. Its design combines attributed edges, augmented node features, and edge dropout for inductive generalization.
- 3.2 Missing Data Problem as a Graph Prediction Task: GRAPE represents observations and features as two node types, with observed feature values assigned to edges in a bipartite graph.
- 3.2 Missing Data Problem as a Graph Prediction Task: Missing feature imputation becomes edge-level prediction, while downstream label prediction becomes node-level prediction on the graph.
- 3.3 Learning with GRAPE: GRAPE extends GraphSAGE for bipartite message passing by incorporating edge embeddings and updating node and edge representations.The architecture uses aggregation, nonlinear updates, trainable weights, and feedforward networks for predictions.
- 3.3 Learning with GRAPE: One-hot feature-node representations and constant observation-node vectors differentiate feature semantics and support generalization to unseen data points.The formulation also supports transferring knowledge from an external dataset with the same feature set.
- 3.3 Learning with GRAPE: Edge dropout masks observed edges during training to prevent identity mappings and improve generalization to unseen edge values.The full graph is used at test time after edges are randomly dropped according to rdrop during training.
4 Experiments
Experiments on nine UCI datasets evaluate GRAPE for imputation, label prediction, varying missing ratios, unseen observations, and end-to-end training. GRAPE generally outperforms baselines and retains performance without retraining on unseen observations.
- Experimental setup: Nine UCI datasets spanning multiple domains are evaluated after randomly removing values from fully observed data matrices.The smallest dataset has 314 observations and 6 features, while the largest has over 45,000 observations and 9 features.
- Feature imputation: GRAPE achieves 20% lower MAE than KNN for feature imputation at a 30% missing rate.Figure 2 reports averaged results over 5 trials, normalized by Mean imputation.
- Label prediction: GRAPE achieves 10% lower MAE than the best baselines for label prediction across the evaluated datasets.The label-prediction setup uses a 70/30% training/test split of labels; most baselines first impute and then apply linear regression.
- Robustness against missingness: Across missing ratios of 0.1, 0.3, 0.5, and 0.7, GRAPE reduces imputation MAE by -8%, 20%, 20%, and 17%, respectively.For prediction, the corresponding reductions are -15%, 10%, 10%, and 4%.
- Generalization: On unseen observations, GRAPE yields 21% lower MAE than MICE without retraining.The model conditions on the entire new dataset, allowing it to capture distributional changes.
- Ablation study: End-to-end training reduces averaged prediction MAE by 19% over a pipeline that first imputes data and then performs linear regression, excluding NAVAL and WINE.The comparison uses the same prediction model as GRAPE’s linear output layer.
5 Conclusion
GRAPE frames missing-data handling as graph-based edge- and node-level prediction tasks trained end-to-end with a GNN. It improves both tasks, generalizes across missingness ratios, and performs well on unseen data points.
- GRAPE formulates feature imputation and label prediction as edge-level and node-level prediction tasks on a graph.
- GRAPE adapts GNN structures to handle continuous edge values while training both tasks end-to-end.
- GRAPE shows significant improvement on both tasks against state-of-the-art imputation approaches across nine standard UCI datasets.
- GRAPE generalizes robustly to unseen data points and different data missing ratios.
Broader Impact
Missing-data quality affects dataset reliability and downstream-task success. GRAPE offers a graph-based perspective for analyzing and handling these problems.
- The quality of imputed data influences dataset reliability and the success of downstream tasks.
- GRAPE provides a new perspective for analyzing and handling missing-data problems through graph representations.
A Additional Details on Baseline Implementation
The experiments use established implementations and specified parameter settings for several feature-imputation baselines. These settings include neighbor count, matrix rank, iteration count, and shrinkage selection.
- Mean, KNN, MICE, SVD, and Spectral baselines use implementations from the fancyimpute package.
- KNN uses 50 nearest neighbors, while SVD sets its rank to m −1, where m is the number of features.
- MICE uses a maximum of 3 iterations, and Spectral uses the default heuristic for its shrinkage value.
B Running Time Comparison
The running-time comparison measures test-time feature-imputation cost across methods. Timing is reported in seconds using method-specific execution units and averaged over five trials in the appendix.
- For Mean, KNN, MICE, SAC, and Spectral, timing covers one function call imputing the entire dataset.
- For GAIN and GRAPE, timing covers one forward pass of the network.
- Appendix B averages running times over five trials using the same setting as Section 4.2.
- Table 2 reports test-time running clock time in seconds for different feature-imputation methods.
C Comparisons with Additional Baselines
GRAPE is compared with missMDA and MIWAE for feature imputation, achieving the smallest imputation error across all datasets under the reported UCI evaluation setting.
- GRAPE is compared with missMDA, a statistical multiple-imputation method, and MIWAE, a deep generative model.
- GRAPE yields the smallest imputation error on all datasets compared with missMDA and MIWAE.
- The UCI evaluation reports averaged MAE for feature imputation at a data missing level of 0.3.
D Experiments on Larger Datasets
Experiments on larger datasets evaluate GRAPE against GNN-based matrix-completion methods using RMSE. GRAPE outperforms GC-MC and performs close to IGMC across the tested datasets.
- The larger-dataset experiments use Flixster, Douban, and YahooMusic feature-imputation datasets with thousands of features.Flixster has 2956 features, Douban 3000, and YahooMusic 1363.
- GRAPE is compared with the GNN-based approaches GC-MC and IGMC on these datasets.
- GRAPE performs better than GC-MC and close to IGMC on all tested datasets.The reported metric is RMSE test performance.