Handling Missing Data with Graph Representation Learning

Introduction

Issues with learning from incomplete data arise in many domains including computational biology, clinical studies, survey research, finance, and economics . The missing data problem has previously been approached in two different ways: feature imputation and label prediction. Feature imputation involves estimating missing feature values based on observed values , and label prediction aims to directly accomplish a downstream task, such as classification or regression, with the missing values present in the input data .

Statistical methods for feature imputation often provide useful theoretical properties but exhibit notable shortcomings: (1) they tend to make strong assumptions about the data distribution; (2) they lack the flexibility for handling mixed data types that include both continuous and categorical variables; (3) matrix completion based approaches cannot generalize to unseen samples and require retraining when the model encounters new data samples . When it comes to models for label prediction, existing approaches such as tree-based methods rely on heuristics and tend to have scalability issues. For instance, one of the most popular procedures called surrogate splitting does not scale well, because each time an original splitting variable is missing for some observation it needs to rank all other variables as surrogate candidates and select the best alternative.

Recent advances in deep learning have enabled new approaches to handle missing data. Existing imputation approaches often use deep generative models, such as Generative Adversarial Networks (GANs) or autoencoders , to reconstruct missing values. While these models are flexible, they have several limitations: (1) when imputing missing feature values for a given observation, these models fail to make full use of feature values from other observations; (2) they tend to make biased assumptions about the missing values by initializing them with special default values.

Here, we propose Grape Project website with data and code: http://snap.stanford.edu/grape, a general framework for feature imputation and label prediction in the presence of missing data. Our key innovation is to formulate the problem using a graph representation, where we construct a bipartite graph with observations and features as two types of nodes, and the observed feature values as attributed edges between the observation and feature nodes (Figure 1). Under this graph representation, the feature imputation can then be naturally formulated as an edge-level prediction task, and the label prediction as a node-level prediction task.

Grape solves both tasks via Graph Neural Networks (GNNs). Specifically, Grape adopts a GNN architecture inspired by the GraphSAGE model , while having three innovations in its design: (1) since the edges in the graph are constructed based on the data matrix and have rich attribute information, we introduce edge embeddings during message passing and incorporate both discrete and continuous edge features in the message computation; (2) we design augmented node features to initialize observation and feature nodes, which provides greater representation power and maintains inductive learning capabilities; (3) to overcome the common issue of overfitting in the missing data problem, we employ an edge dropout technique that greatly boosts the performance of Grape.

We compare Grape with the state-of-the-art feature imputation and label prediction algorithms on 9 benchmark datasets from the UCI Machine Learning Repository . In particular, Grape yields 20% lower mean absolute error (MAE) for the imputation tasks and 10% lower MAE for the prediction tasks at the 30% data missing rate. Finally, we demonstrate Grape’s strong generalization ability by showing its superior performance on unseen observations without the need for retraining.

Overall, our approach has several important benefits: (1) by creating a bipartite graph structure we create connections between different features (via observations) and similarly between the observations (via features); (2) GNN elegantly harnesses this structure by learning to propagate and borrow information from other features/observations in a graph localized way; (3) GNN allows us to model both feature imputation as well as label prediction in an end-to-end fashion, which as we show in experiments leads to strong performance improvements.

Related Work

Feature imputation. Successful statistical approaches for imputation include joint modeling with Expectation-Maximization , multivariate imputation by chained equations (MICE) , $k$-nearest neighbors (KNN) , and matrix completion . However, joint modeling tends to make assumptions about the data distribution through a parametric density function; joint modeling and matrix completion lack the flexibility to handle data of mixed modalities; MICE and KNN cannot accomplish imputation while adapting to downstream tasks.

Recently, deep learning models have also been used to tackle the feature imputation problem . However, these models have important limitations. Denoising autoencoder (DAE) models and GAIN only use a single observation as input to impute the missing features. In contrast, Grape explicitly captures the complex interactions between multiple observations and features. GNN-based approaches have also been proposed in the context of matrix completion . However, they often make the assumption of finite, known-range values in their model design, which limits their applicability to imputation problems with continuous values. In contrast, Grape can handle both continuous and discrete feature values.

Label prediction with the presence of missing data. Various models have been adapted for label prediction with the presence of missing data, including tree-based approaches , probabilistic modeling , logistic regression , support vector machines , deep learning-based models , and many others . Specifically, decision tree is a classical statistical approach that can handle missing values for the label prediction task . With the surrogate splitting procedure, decision tree uses a single surrogate variable to replace the original splitting variable with missing values, which is effective but inefficient, and has been shown to be inferior to the “impute and then predict” procedure . Random forests further suffer from the scalability issues as they consist of multiple decision trees . In contrast, Grape handles the missing feature entries naturally with the graph representation without any additional heuristics. The computation of Grape is efficient and easily parallelizable with modern deep learning frameworks.

Overall discussion. In Grape implementation, we adopt several successful GNN design principles. Concretely, our core architecture is inspired by GraphSAGE ; we apply GraphSAGE to bipartite graphs following G2SAT ; we use edge dropout in ; we use one-hot auxiliary node features which has been used in ; we follow the GNN design guidelines in to select hyperparameters. Moreover, matrix completion tasks have been formulated as bipartite graphs and solved via GNNs in ; however, they only consider the feature imputation task with discrete feature values. We emphasize that our main contribution is not the particular GNN model but the graph-based framework for the general missing data problem. Grape is the first graph-based solution to both feature imputation and label prediction aspects of the missing data problem.

The Grape Framework

2 Missing Data Problem as a Graph Prediction Task

The key insight of this paper is to represent the feature matrix with missing values as a bipartite graph. Then the feature imputation problem and the label prediction problem can naturally be formulated as node prediction and edge prediction tasks (Figure 1).

Feature matrix as a bipartite graph. The feature matrix $\mathbf{D}$ and the mask $\mathbf{M}$ can be represented as an undirected bipartite graph $\mathcal{G}=(\mathcal{V},\mathcal{E})$, where $\mathcal{V}$ is the node set that consists of two types of nodes $\mathcal{V}=\mathcal{V}_{D}\cup\mathcal{V}_{F}$, $\mathcal{V}_{D}=\{u_{1},...,u_{n}\}$ and $\mathcal{V}_{F}=\{v_{1},\ldots,v_{m}\}$, $\mathcal{E}$ is the edge set where edges only exist between nodes in different partitions: $\mathcal{E}=\{(u_{i},v_{j},\mathbf{e}_{u_{i}v_{j}})\mid u_{i}\in\mathcal{V}_{D},v_{j}\in\mathcal{V}_{F},\mathbf{M}_{ij}=1\}$, where the edge feature, $\mathbf{e}_{u_{i}v_{j}}$, takes the value of the corresponding feature $\mathbf{e}_{u_{i}v_{j}}=\mathbf{D}_{ij}$. If $\mathbf{D}_{ij}$ is a discrete variable then it is transformed to a one-hot vector then assigned to $\mathbf{e}_{u_{i}v_{j}}$. To simplify the notation $\mathbf{e}_{u_{i}v_{j}}$, we use $\mathbf{e}_{ij}$ in the context of feature matrix $\mathbf{D}$, and $\mathbf{e}_{uv}$ in the context of graph $\mathcal{G}$.

Feature imputation as edge-level prediction. Using the definitions above, imputing missing features can be represented as learning the edge value prediction mapping: $\hat{\mathbf{D}}_{ij}=\hat{\mathbf{e}}_{ij}=f_{ij}(\mathcal{G})$ by minimizing the difference between $\hat{\mathbf{D}}_{ij}$ and $\mathbf{D}_{ij},\forall\mathbf{M}_{ij}=0$. When imputing discrete attributes, we use cross entropy loss. When imputing continuous values, we use MSE loss.

Label prediction as node-level prediction. Predicting downstream node labels can be represented as learning the mapping: $\hat{\mathbf{Y}}_{i}=g_{i}(\mathcal{G})$ by minimizing the difference between $\hat{\mathbf{Y}}_{i}$ and $\mathbf{Y}_{i},\forall\mathbf{V}_{i}=0$.

3 Learning with Grape

Grape adopts a GNN architecture inspired by GraphSAGE , which is a variant of GNNs that has been shown to have strong inductive learning capabilities across different graphs. We extend GraphSAGE to a bipartite graph setting by adding multiple important components that ensure its successful application to the missing data problem.

Grape GNN architecture. Given that our bipartite graph $\mathcal{G}$ has important information on its edges, we modify GraphSAGE architecture by introducing edge embeddings. At each GNN layer $l$, the message passing function takes the concatenation of the embedding of the source node $\mathbf{h}_{v}^{(l-1)}$ and the edge embedding $\mathbf{e}_{uv}^{(l-1)}$ as the input:

where $\textsc{Agg}_{l}$ is the aggregation function, $\sigma$ is the non-linearity, $\mathbf{P}^{(l)}$ is the trainable weight, $\mathcal{N}$ is the node neighborhood function. Node embedding $\mathbf{h}_{v}^{(l)}$ is then updated using:

where $\mathbf{Q}^{(l)}$ is the trainable weight, we additionally update the edge embedding $\mathbf{e}_{uv}^{(l)}$ by:

where $\mathbf{W}^{(l)}$ is the trainable weight. To make edge level predictions at the $L$-th layer:

The node-level prediction is made using the imputed dataset $\hat{\mathbf{D}}$:

where $\mathbf{O}_{edge}$ and $\mathbf{O}_{node}$ are feedforward neural networks.

Augmented node features for bipartite message passing. Based on our definition, nodes in $\mathcal{V}_{D}$ and $\mathcal{V}_{F}$ do not naturally come with features. The straightforward approach would be to augment nodes with constant features. However, such formulation would make Grape hard to differentiate messages from different feature nodes in $\mathcal{V}_{F}$. In real-world applications, different features can represent drastically different semantics or modalities. For example in the Boston Housing dataset from UCI , some features are categorical such as if the house is by the Charles River, while others are continuous such as the size of the house.

Instead, we propose to use $m$-dimensional one-hot node features for each node in $\mathcal{V}_{F}$ ($m=|\mathcal{V}_{F}|$), while using $m$-dimensional We make data nodes and feature nodes to have the same feature dimension for the ease of implementation. constant vectors as node feature for data nodes in $\mathcal{V}_{F}$:

Such a formulation leads to a better representational power to differentiate feature nodes with different underlying semantics or modalities. Additionally, the formulation has the capability of generalizing the trained Grape to completely unseen data points in the given dataset. Furthermore, it allows us to transfer knowledge from an external dataset with the same set of features to the dataset of interest, which is particularly useful when the external dataset provides rich information on the interaction between observations and features (as captured by Grape). For example, as a real-world application in biomedicine, gene expression data can be used to predict disease types and frequently contain missing values. If we aim to impute missing values in a gene expression dataset of a small cohort of lung cancer patients, public datasets, e.g., the Cancer Genome Atlas Program (TCGA) can be first leveraged to train Grape, where rich interactions between patients and features are learned. Then, the trained Grape can be applied to our smaller dataset of interest to accomplish imputation.

Improved model generalization with edge dropout. When doing feature imputation, a naive way of training Grape is to directly feed $\mathcal{G}=(\mathcal{V};\mathcal{E})$ as the input. However, since all the observed edge values are used as the input, an identity mapping $\hat{\mathbf{D}}_{ij}=\mathbf{e}^{(0)}_{ij}$ is enough to minimize the training loss; therefore, Grape trained under this setting easily overfits the training set. To force the model to generalize to unseen edge values, we randomly mask out edges $\mathcal{E}$ with dropout rate $r_{drop}$:

Experiments

Datasets. We conduct experiments on 9 datasets from the UCI Machine Learning Repository . The datasets come from different domains including civil engineering (concrete, energy), biology (protein), thermal dynamics (naval), etc. The smallest dataset (yacht) has 314 observations and 6 features, while the largest dataset (protein) has over 45,000 observations and 9 features. The datasets are fully observed; therefore, we introduce missing values by randomly removing values in the data matrix. The attribute values are scaled to $$ with a MinMax scaler .

Baseline models. We compare our model against five commonly used imputation methods. We also compare with a state-of-the-art deep learning based imputation model as well as a decision tree based label prediction model. More details on the baseline models are provided in the Appendix.

Mean imputation (Mean): The method imputes the missing $\mathbf{D}_{ij}$ with the mean of all the samples with observed values in dimension $j$.

K-nearest neighbors (KNN): The method imputes the missing value $\mathbf{D}_{ij}$ using the KNNs that have observed values in dimension $j$ with weights based on the Euclidean distance to sample $i$.

Multivariate imputation by chained equations (MICE): The method runs multiple regression where each missing value is modeled conditioned on the observed non-missing values.

Iterative SVD (SVD) : The method imputes missing values based on matrix completion with iterative low-rank SVD decomposition.

Spectral regularization algorithm (Spectral) : This matrix completion model uses the nuclear norm as a regularizer and imputes missing values with iterative soft-thresholded SVD.

GAIN , state-of-the-art deep imputation model with generative adversarial training .

Decision tree (Tree) , a commonly used statistical method that can handle missing values for label prediction. We consider this baseline only for the label prediction task. Random forest is not included due to the lack of a public implementation that can handle missing data without imputation.

Grape configurations. For all experiments, we train Grape for 20,000 epochs using the Adam optimizer with a learning rate at 0.001. For all feature imputation tasks, we use a 3-layer GNN with 64 hidden units and ReLU activation. The $\textsc{AGG}_{l}$ is implemented as a mean pooling function $\textsc{Mean}(\cdot)$ and $\mathbf{O}_{edge}$ as a multi-layer perceptron (MLP) with 64 hidden units. For label prediction tasks, we use two GNN layers with 16 hidden units. $\mathbf{O}_{edge}$ and $\mathbf{O}_{node}$ are implemented as linear layers. The edge dropout rate is set to $r_{drop}=0.3$. For all experiments, we run 5 trials with different random seeds and report the mean and standard deviation of the results.

2 Feature Imputation

Results. As shown in Figure 2, Grape has the lowest MAE on all datasets and its average error is 20% lower compared with the best baseline (KNN). Since there are significant differences between the characteristics of different datasets, statistical methods often need to adjust its hyper-parameters accordingly, such as the cluster number in KNN, the rank in SVD, and the sparsity in Spectral. On the contrary, Grape is able to adjust its trainable parameters adaptively through loss backpropagation and learn different observation-feature relations for different datasets. Compared with GAIN, which uses an MLP as the generative model, the GNN used in Grape is able to explicitly model the information propagation process for predicting missing feature values.

3 Label Prediction

Results. As is shown in Figure 2, on all datasets except naval and wine, Grape has the best performance. On wine dataset, all methods have comparable performance. The fact that the performance of all methods are close to the Mean method indicates that the relation between the labels and observations in wine is relatively simple. For the dataset naval, the imputation errors of all models are very small (both relative to Mean and on absolute value). In this case, a linear regression on the imputed data is enough for label prediction. Across all datasets, Grape yields 10% lower MAE compared with best baselines. The improvement of Grape could be explained by two reasons: first, the better handling of missing data with Grape where the known information and the missing values are naturally embedded in the graph; and second, the end-to-end training.

4 Robustness against Different Data Missing Levels

Setup. To examine the robustness of Grape with respect to the missing level of the data matrix. We conduct the same experiments as in Sections 4.2 and 4.3 with different missing levels of $r_{miss}\in\{0.1,0.3,0.5,0.7\}$.

Results. The curves in Figure 3 demonstrate the performance change of all methods as the missing ratio increases. Grape yields -8%, 20%, 20%, and 17% lower MAE on imputation tasks, and -15%, 10%, 10%, and 4% lower MAE on prediction tasks across all datasets over missing ratios of 0.1, 0.3, 0.5, and 0.7, respectively. In missing ratio of 0.1, the only baseline that behaves better than Grape is KNN. As in this case, the known information is adequate for the nearest-neighbor method to make good predictions. As the missing ratio increases, the prediction becomes harder and the Grape’s ability to coherently combine all known information becomes more important.

5 Generalization on New Observations

Results. As shown in Figure 4, Grape yields 21% lower MAE compared with best baselines (MICE) without being retrained, indicating that our model generalizes seamlessly to unseen observations. Statistical methods have difficulties transferring the knowledge in the training data to new data. While GAIN is able to encode such information in the generator network, it lacks the ability to adapt to observations coming from a different distribution. However, by using a GNN, Grape is able to make predictions conditioning on the entire new datasets, and thus capture the distributional changes.

6 Ablation Study

Edge dropout. We test the influence of the edge dropout on the performance of Grape. We repeat the experiments in Section 4.2 for Grape with no edge dropout and the comparison results are shown in Table 1. The edge dropout reduces the test MAE by 33% on average, which verifies our assumption that using edge dropout could help the model learn to predict unseen edge values.

Aggregation function. We further investigate how the aggregation function ($\textsc{Sum}(\cdot)$, $\textsc{Max}(\cdot)$, $\textsc{Mean}(\cdot)$) of GNN affects Grape’s performance. While $\textsc{Sum}(\cdot)$ is theoretically most expressive, in our setting the degree of a specific node is determined by the number of missing values which is random and unrelated to the missing data task; in contrast, the $\textsc{Mean}(\cdot)$ and $\textsc{Max}(\cdot)$ aggregators are not affected by this inherent randomness of node degree, therefore they perform better.

End-to-end downstream regression. To show the benefits of using end-to-end training in label prediction, we repeat the experiments in Section 4.3 by first using Grape to impute the missing data and then perform linear regression on the imputed dataset for node labels (which is the same prediction model as the linear layer used by Grape). The results are shown in Table 1. The end-to-end training gets 19% less averaged MAE over all datasets except naval and wine. The reason for the two exceptions is similar as described in Section 4.3.

7 Further Discussions

Scalability. In our paper, we use UCI datasets as they are widely-used datasets for benchmarking imputation methods, with both discrete and continuous features. Grape can easily scale to datasets with thousands of features. We provide additional results on larger-scale benchmarks, including Flixster (2956 features), Douban (3000 features), and Yahoo (1363 features) in the Appendix. Grape can be modified to scale to even larger datasets. We can use scalable GNN implementations which have been successfully applied to graphs with billions of edges ; when the number of features is prohibitively large, we can use a trainable embedding matrix to replace one-hot node features.

Applicability of Grape. In the paper, we adopt the most common evaluation regime used in missing data papers, i.e., features are missing completely at random. Grape can be easily applied to other missing data regimes where feature are not missing at random, since Grape is fully data-driven.

More intuitions on why Grape works. When a feature matrix does not have missing values, to make downstream label predictions, a reasonable solution will be directly feeding the feature matrix into an MLP. As is discussed in , an MLP can in fact be viewed as a GNN over a complete graph, where the message function is matrix multiplication. Under this interpretation, Grape extends a simple MLP by allowing it to operate on sparse graphs (i.e., feature matrix with missing values), enabling it for missing feature imputation tasks, and adopting a more complex message computation as we have outlined in Algorithm 1.

Conclusion

In this work, we propose Grape, a framework to coherently understand and solve missing data problems using graphs. By formulating the feature imputation and label prediction tasks as edge-level and node-level predictions on the graph, we are able to train a Graph Neural Network to solve the tasks end-to-end. We further propose to adapt existing GNN structures to handle continuous edge values. Our model shows significant improvement in both tasks compared against state-of-the-art imputation approaches on nine standard UCI datasets. It also generalizes robustly to unseen data points and different data missing ratios. We hope our work will open up new directions on handling missing data problems with graphs.

Broader Impact

The problem of missing data arises in almost all practical statistical analyses. The quality of the imputed data influences the reliability of the dataset itself as well as the success of the downstream tasks. Our research provides a new point of view for analysing and handling missing data problems with graph representations. There are many benefits to using this framework. First, different from many existing imputation methods which rely on good heuristics to ensure the performance , Grape formulates the problem in a natural way without the need of handcrafted features and heuristics. This makes our method ready to use for datasets coming from different domains. Second, similar to convolutional neural networks , Grape is suitable to serve as a pre-processing module to be connected with downstream task-specific modules. Grape could either be pre-trained and fixed or concurrently learned with downstream modules. Third, Grape is general and flexible. There is little limitation on the architecture of the graph neural network as well as the imputation ($\mathbf{O}_{edge}$) and prediction ($\mathbf{O}_{node}$) module. Therefore, researchers can easily plug in domain-specific neural architectures, e.g., BERT , to the design of Grape. Overall, we see exciting opportunities for Grape to help researchers handle missing data and thus boost their research.

Acknowledgments

We gratefully acknowledge the support of DARPA under Nos. FA865018C7880 (ASED), N660011924033 (MCS); ARO under Nos. W911NF-16-1-0342 (MURI), W911NF-16-1-0171 (DURIP); NSF under Nos. OAC-1835598 (CINES), OAC-1934578 (HDR), CCF-1918940 (Expeditions), IIS-2030477 (RAPID); Stanford Data Science Initiative, Wu Tsai Neurosciences Institute, Chan Zuckerberg Biohub, Amazon, Boeing, JPMorgan Chase, Docomo, Hitachi, JD.com, KDDI, NVIDIA, Dell. J. L. is a Chan Zuckerberg Biohub investigator.

References

Appendix A Additional Details on Baseline Implementation

For imputation baselines including Mean, KNN, MICE, SVD, and Spectral, we use the implementation provided in the fancyimpute package https://github.com/iskandr/fancyimpute. For KNN, we use 50 nearest neighbors. For SVD, we set the rank equal to $m-1$, where $m$ is the number of features. For MICE, we set the maximum iteration number to 3. For Spectral, we found the default heuristic for shrinkage value works the best. For a detailed explanation of the meaning of the parameters, we refer readers to the documentation of fancyimpute package. The hyper-parameter values are chosen by comparing the average imputation performance over all datasets. For GAIN, we use the source code released by the authors. All the hyper-parameters are the same as in the source code https://github.com/jsyoon0823/GAIN. We use the rpart R package for the implementation of the decision tree method.

Appendix B Running Time Comparison

Here we report the running clock time for feature imputation of different methods at test time. For Mean, KNN, MICE, SAC, and Spectral, this means the running time of one function call for imputing the entire dataset. For GAIN and Grape, this means one forward pass of the network. Table 2 shows the averaged running time over 5 different trials with the same setting as described in Section 4.2.

Appendix C Comparisons with Additional Baselines

We additionally provide the comparison results of our method with two other state-of-the-art baselines: missMDA , a statistical multiple imputation approach, and MIWAE , a deep generative model. We adapt the same setting as in Section 4.1 and the results are shown in Table 3. Grape yields the smallest imputation error on all datasets compared with the two other baselines.

Appendix D Experiments on Larger Datasets

To test the scalability of Grape, we perform additional feature imputation tests on the Flixter, Douban, and YahooMusic detests with preprocessed subsets and splits provided by . The Flixster dataset has 2341 observations and 2956 features. The Douban dataset has 3000 observations and 3000 features. The YahooMusic dataset has 1357 observations and 1363 features. These datasets only have discrete values. We compare Grape with two GNN-based approaches, GC-MC and IGMC . The results are shown in Table 4, where the results of GC-MC and IGMC are provided by . On all datasets, Grape shows a reasonable performance which is better than GC-MC and close to IGMC. Notice that the two baselines are specially designed for discrete matrix completion, where Grape is applicable to both continuous and discrete feature values and is general for both feature imputation and label prediction tasks.