Wasserstein diffusion on graphs with missing attributes

Abstract: Many real-world graphs are attributed graphs where nodes are associated with non-topological features. While attributes can be missing anywhere in an attributed graph, most of existing node representation learning approaches do not consider such incomplete information.
      In this paper, we propose a general non-parametric framework to mitigate this problem. Starting from a decomposition of the attribute matrix, we transform node features into discrete distributions in a lower-dimensional space equipped with the Wasserstein metric. On this Wasserstein space, we propose Wasserstein graph diffusion to smooth the distributional representations of nodes with information from their local neighborhoods. This allows us to reduce the distortion caused by missing attributes and obtain integrated representations expressing information of both topology structures and attributes. We then pull the nodes back to the original space and produce corresponding point representations to facilitate various downstream tasks. To show the power of our representation method, we designed two algorithms based on it for node classification (with missing attributes) and matrix completion respectively, and demonstrate their effectiveness in experiments.

Reviews and author responses:

**Q1.For the ground metric in Eq. (3), what is the rationale for choosing as the norm (not sure which norm) of square eigenvalues?**

ANS: The metric we use to measure the distance among basis vectors is the Euclidean metric ($||X(v_i - v_j)||^2=||\lambda_i v_i-\lambda_j v_j||^2=|\lambda_i^2-\lambda_j^2|$) and for the transformed discrete distributions we use Wasserstein distance, here {\lambda_i} are the corresponding eigenvalues, $|v_i|^2 = |v_i|^2 = 1$.

**Q2.The paragraph including Eq. (6) ("Recall that the set ... notoriously difficult to solve.") does not contribute any further information as it is not used in the rest of the paper.**

ANS: This paragraph is to illustrate our motivation to use fixed-support barycenter instead of the free-support barycenter. 

**Q3.In Eq. (9), does the inverse map of X need to be consistent with the observed feature of that node?**

ANS: The inverse map of X does not need to be consistent with the observed feature. In Eq.(7), we obtain node distributional representations in discrete Wasserstein space through the Wasserstein diffusion process. Then we want to leverage the inverse map to transform such distributional representations into the original space (since we want to incorporate the information captured by “V”). Therefore, the inverse map of X is indeed node representations in the original space. Also, as we explained in the matrix completion tasks, we can add a learnable transform to the inverse map of X and a reconstruction constraint to the loss function to make the inverse map of X consistent with the observed feature.

**Q4. How to handle SVD with missing nodes? It seems that if the graph is 100% missing features, i.e. the graph with no features, the method does not work as we can not perform the SVD step.**

ANS: We did not handle SVD with missing nodes but just replaced the missing value with zero value. Our motivation is to utilize valid but incomplete feature information to learn expressive node representation, therefore we are not trying to cope with the graph with 100% missing features (it means that there is no feature information). However, for node classification, we can generate one-hot features according to node labels and take it as our input such that our method can still work when there are no attributed features.

**Q5.Is there any convergence guarantee for the proposed Wasserstein diffusion process in Eq. (7)? If there is no rigorous proof, an intuition justification is the least.**

ANS: Actually, a convergence guarantee is not needed here or in other words, convergence is unwanted in the Wasserstein diffusion process. To some extent, the Wasserstein diffusion process can be understood as a kind of aggregation function (introduced in many traditional GCNs models) in discrete distribution space. It can help convalidation and propagate node distributional information. While if the iteration number is big enough or “convergent”,  it will lead to over-smoothing.

**Q6.Some SOTA baseline methods such as GCNmf [1] or methods therein (e.g. GAIN) should be compared with.**

ANS: We refer to the experimental results of GCNmf and GAIN reported in [1] and take Citeseer as an example (in the whole missing setting):

miss_rate  $\qquad$  0.1|   $\quad$   0.2 |   $\quad$    0.3  |   $\quad$     0.4 |  $\qquad$       0.5|     $\quad$     0.6 |  $\quad$       0.7 | $\quad$        0.8 |    $\quad$     0.9

GAIN    $\qquad$    0.6947| 0.6786 | 0.6588 | 0.6396 | 0.5996 | 0.5424 | 0.4121 | 0.2531 | 0.1789

GCNmf  $\quad$   0.7044 | 0.6856 | 0.6657 | 0.6539 | 0.6344 | 0.6004 | 0.5688 | 0.5137 | 0.3986 

Wgd_mlp  0.7040 | 0.6946 | 0.6825 | 0.6634 | 0.6482 | 0.6329 | 0.6205 | 0.5864 | 0.5461



Thank you very much for your encouraging comment. We changed the format and topos as you suggested.

**Q5.It seems that the performance of the proposed method is exactly the same in both settings. (Figure 1) This is curious. Can the authors comment on what is happening in this case?**

ANS: Our model’s performance in the two settings may be similar but obviously not the same. Take the Cora dataset as an example, we list the numeral results here:

Whole: 0.8020 | 0.7945 | 0.7886 | 0.7831 | 0.7742 | 0.7592 | 0.7493 | 0.7311 | 0.7015

Partial: 0.8076 | 0.8041 | 0.7974 | 0.7890 | 0.7818 | 0.7658 | 0.7584 | 0.7272 | 0.7011  

Moreover, we find that other baselines also show similar performance in the two cases. In the paper we use figures to save the space and make the comparison more understandable.

**Q6.In multiple instances, the authors introduce experiments but don’t comment on them or explain the similarity or differences in performance between all of these experimental settings.**

ANS: Thanks for the suggestion. Compared to other models, our model finds a new way for directly modeling missing features in graphs and it belongs to a different category. In this case, that is quite common to just claim the performance in ML papers. Besides, we have comments for comparison to our ablation models (e.g. SVD-GCN), which can help explain the benefit of our model. To further clarify, we also add additional details of other baselines in the revised version.


**Q3. Because many choices are not justified or described, a number of questions are left unanswered. What is fundamentally different between a WGD layer and a GCN layer? What is the influence of the number of WGD layers? And what is changing in the representation after each WGD layer? What does the induced latent space look like? How does it relate to the complete feature information setting? Is the dimension of the latent space equal to that of the original space? Why does feature recovery using Wasserstein barycenter update trumps simple average or k-means of neighbors, what makes the metric be more adequate than Euclidean distance?**

ANS: The main difference is the WGD layers diffuse and consolidate node distributional information while the GCN layers diffuse and consolidate node attribute information. To some extent, Barycenter_Update in each WGD layer is a generalized aggregation approach in the Wasserstein space, i.e. the node aggregation in WGD considers also the pairwise distance between different components while GCN does not.

The latent space induced in each WGD layer is a Wasserstein space composed of discrete distributions with the common support points which are also shared among all the latent spaces. Therefore the shape of matrices of the latent node distributional representations is fixed and the same as the original principal components matrix U. Our method can be used in the complete feature setting, since the procedure is the same. But that is not our focus. 

As we illustrated in Q1.a, leveraging the Wasserstein barycenter update can incorporate the valid information of basis into node representations computation. If we just aggregate information of neighbors in the principal component space (a low-dimensional Euclidean space), we will ignore the ground geometry information. In the experiments, we have an ablation model SVD-GCN, which makes us compare the WGD layer with the GCN layer, and our final model is clearly better.

**Q4.The authors claim that their work “is the ﬁrst work to compute embeddings of a graph with incomplete attributes directly.” but their proposed algorithm is two-stage, first matrix completion and then neural network training. The paper seems to be more relevant to data pre-processing than representation learning itself.**

ANS: As we demonstrated in Q1.c, we are not trying to impute the incomplete matrix or pre-process data but to directly derive node representations from missing features. Our model is an end-to-end model to generate distributional representations first and node representations in the original feature space eventually. Besides, in Section 4, we show that we can naturally generalize our model to matrix completion tasks if we apply an extra learnable transformation to our computed node representations and add reconstruction regularization to the loss functions. Definitely, it is still an end-to-end method. Therefore, the statement that “the proposed algorithm is two-stage, first matrix completion and then neural network training” is incorrect (actually **the $\tilde{X}$ in Equation (9) is not a matrix completion for the original $X$, for the elements in $X$ do not remain the same.**) This neatly illustrates the flexibility and capacity of our framework which can be used to compute node representations or matrix completion respectively. 

In fact, the updated node representations in the principal component space (tilde U_k in eq.(8)) already have good expression capacity. To verify this statement, we directly take this node representations as input of MLP for node classification without pulling it back to feature matrix, and conduct experiments on Cora in the whole missing settings:

Missing rate         $ \qquad$  0.5   |    $ \quad$   0.6    |$ \qquad$        0.7| $ \quad$     0.8       |  $ \quad$    0.9  

Original model    0.7742  |   0.7592  |     0.7493  |     0.7311  |   0.7015

Ablation model   0.7684  |    0.7525 |     0.7404  |     0.7144  |   0.6968


**Q1.While the experiments show promising results, the authors do not explain the choices behind their method. Why do the authors use SVD and define the distance between transformed features in the way they did? Is there a reason for using the exponential function in the transformation other than for inducing the positiveness of probability? How does the inverse mapping come up and why does it work?**

ANS: Thanks for the suggestion. We add more explanation of the model design choices in the revised paper for clarification.  

- a. The choices of SVD and distance. At the beginning of Section 3.2,  we demonstrated our motivation for using SVD. We follow the common assumption: the feature matrix is low-ranked and is inspired by the popular approach for matrix completion: Alternating Least Square(ALS) Algorithm which uses SVD to factorize matrix into two low-ranked submatrices, i.e. X=AB, A=U(S^½), B=(S^½)V^T. ALS alternatively optimizes A and B then replaces X at each step using the most recently computed A and B till convergence to impute the incomplete matrix. However, our proposed method is not for matrix completion but node representation learning and we don’t need to optimize the low-ranked submatrix. Through matrix factorization using SVD, we obtain initial node representations in the principal components space (U) along with the corresponding basis (SV). We aim to learn expressive node representations in the principal components space (i.e. only optimize U), meanwhile we want to take the information of the basis into account. It is elegant and novel to regard the principal components vectors as generalized discrete distributions with basis vectors as the support points. The metric we use to measure the distance among basis vectors is the Euclidean metric (||X(vi-vj)||^2=||sivi-sjvj||^2=|si^2-sj^2|) and for the transformed discrete distributions we use Wasserstein distance.

- b. Exponential function. As we mentioned above, we regarded the principal components vectors as generalized discrete distributions and we introduced a reversible positive function to map them into the standard discrete distribution space. The choice of such function depends on data, in our paper we use exponential function and it can be replaced by any other appropriate reversible positive function.

- c. Inverse map. Through the Wasserstein diffusion process, we obtain distributional node representations in the discrete Wasserstein space, then we pull them back to the principal components space (log transformation and Gram-Schmidt Orthogonalization). Although we can directly apply such “componential” node representations to graph-based tasks, the valid information of basis is ignored. Therefore, we pull node representations back to the original feature space using the updated principal components matrix and basis matrix. In summary, we transform node features to principal components space, discrete Wasserstein space, and then back to the original feature space.   
 
**Q2. As for the choice of architectures, is there a reason for using 7 WGD layers versus 2 GCN layers for the baselines?**

ANS: We fine-tune the hyperparameters on one setting and then use them for all other settings. Using 7 WGD layers is the best hyperparameter choice in our model and it is also a good tradeoff between accuracy and efficiency. We provide sensitive analysis in revision version.  Using 2 GCN layers is a well-known default design choice for GCNs (as indicated in the original GCN paper); it can make a trade-off between the expressiveness and the over-smoothing. 




 **Q1.The idea in this work seems to be a combination of techniques from different previous works: Use optimal transport to complete missing data; Graph-based semi-supervised learning.**

ANS: The reviewer seems to have some misunderstandings of the paper. Although our method incorporates optimal transport technique and can be used to handle multi-graph matrix completion and many graph-based tasks including semi-supervised node classification, our method is not a combination of previous works. Actually, it is novel and quite different from most of the related works including the two works you mentioned. 
About our originality:

- a. Our method is **not trying to complete missing data** but aims to generate node representations incorporating such incompleted feature information as well as graph structure. We introduced optimal transport in our model, not for feature completion or label propagation but to consolidate the valid latent information captured from the limited observed features such that we can obtain rich characteristics of each node and generate powerful node representations.

- b. Our model can be applied to many different downstream graph-based tasks. We just tested the power of our produced node representations on the semi-supervised node classification task. And we also showed that our model can be naturally generalized to complete missing data with graph structure.

- c. Besides, our methodology has little to do with that of [1] and [2].
 [1] leverages optimal transport distance to work as its loss function with the assumption that random batches from the same dataset follow the same distribution, while the optimal transport distance is the foundation of the Wasserstein diffusion process and our loss function only depends on the corresponding task. In our diffusion process, each node is represented as an independent discrete distribution. Moreover, our problem setting is also different: [1] focus on matrix imputation without leveraging structure information while we focus on missing data with graph structure and we actually **did not complete the missing data but directly model it for node classification**.
[2] minimizes the Dirichlet energy with Wasserstein distance in the space of distribution-valued maps with prescribed distributions while in our method, we just used Wasserstein barycenter to update the consolidated distributional information.

**Q2.The "exp" function is not properly transformed function to keep the feature information. The model only tends to capture those large positive components and ignores those negative components (due to the exp function).**

ANS: As we mentioned in our paper, we need a reversible positive function to obtain discrete distributions, the exponential function is just one of the optional choices. We agree it may lead to more emphasis on large positive components, but the negative components are not ignored since the comparative gaps between the negative ones and others remain the same. The experimental results have already shown that exp is a valid choice on all our benchmark datasets. Of course there is no guarantee that it will work for any datasets, but the choice of reversible positive function does not change our architecture and it is fine to leverage any appropriate functions.



**Q3.Use better baselines for node classification. Moreover, better methods with only graph structures should be traditional methods, such as label propagation.**

ANS: The principal way to cope with missing data is using different kinds of matrix completion methods for data preprocessing, but they do not consider the graph structure in the matrix completion phase and have limited capacity when the missing rate is large. One of the main aims of this paper is to show that our method is better at node classification with missing features even without explicit matrix completion pre-processing. To prove this statement, we leverage GCN to be the basic model for node classification on processed feature matrices imputed by various matrix completion methods. The GCN is the most well-known and popular node classification model on graphs; most advanced GNNs actually have limited performance gain over it on those standard datasets. The experimental results confirmed our inference: most matrix completion methods are hard to compensate for the loss of feature information caused by missing values such that there may be no sufficient information for node classification. Nothing comes from nothing. **Using a better node classification model (more advanced GNN) does not solve the “missing feature” problem** and the comparison with different variations of GCNs demonstrates what we expect to show.
 
Thanks for your suggestion about using traditional methods as the baseline for only topological information. We **added the Label Propagation Algorithm** as baselines which once again verifies our impressive representation capacity as our method consistently achieves a huge improvement.

**Q4.To demonstrate the effectiveness of this approach, experiments should also be done over heterophilic networks.**

We respectfully disagree with the necessity of doing experiments on heterophilic networks. Our method is a general approach for dealing with missing features in graphs whose motivation has no relation to heterophilic networks. For matrix completion, we are using very common datasets. For node classification with missing features, we are also using standard benchmark datasets, which is enough to demonstrate our contribution.


We thank all the reviewers for helpful and detailed feedback. The main modifications in our revised paper are listed as follows:

1.We added more explanation about the motivation of the model design choices in section 3; 

2.We  provide sensitive analysis for hyperparameters;

3.We  add Label Propagation as another baseline for node classification;

4.We add additional details of baselines of matrix completion and correct some typos that reviewers pointed out.


This paper presents a Wasserstein diffusion based method for estimating missing node labels on attributed graphs. The algorithms use linear algebraic decompositions to represent the node labels in a low dimensional space and then uses a Wasserstien Barycenter approach to implement the diffusion before lifting back to the original feature space. None of these components is novel in this setting but the combined algorithm is interesting and the provided code is helpful and demonstrates the naturalness and ease of implementation of the proposed model. 

The empirical experiments are beyond sufficient given the space constraints and highlight the flexiblity of this formulation of the problem. The application to multigraph completion is particularly interesting - this is related to several other recent problems of interest around node embeddings and inference around families of graphs on common node sets that would be an interesting extension of this work. 

The typesetting of text in equations should be placed in something like \operatorname so that it isn't squashed and italicized. Also the final table needs to be reformatted, as it is difficult to parse currently. 

Typo second to last paragraph of page 5 low-rand -> low-rank

In the paper, the authors proposed a framework to learn representation for graphs with node attributes that are allowed to be missed.   The learning process includes two steps: (1) node features transformation using SVD which allows representing a node as a discrete distribution with support points of orthogonal basis vectors. (2) message passing and updating the node distribution using Wasserstein barycenter, i.e. one node representation is the barycenter of its neighbors and itself.  Two downstream tasks were used to demonstrate the performance of the proposed methods:  node classification and multi-graph completion.

Overall, the idea is interesting however the writing of technical ingredients is not convinced enough. There are points not clear and can be improved to get a better version of the paper.
 - For the ground metric in Eq. (3), what is the rationale for choosing as the norm (not sure which norm) of square eigenvalues?
 - The paragraph including Eq. (6) ("Recall that the set ... notoriously difficult to solve.") does not contribute any further information as it is not used in the rest of the paper.
 - In Eq. (9), do the inverse map of  X need to be consistent with the observed feature of that node?
 - How to handle SVD with missing nodes? It seems that if the graph is 100% missing features, i.e. the graph with no features, the method does not work as we can not perform the SVD step.
 - Is there any convergence guarantee for the prosed Wasserstein diffusion process in Eq. (7)? If there is no rigorous proof, an intuition justification is the least.
 - Some SOTA baseline methods such as GCNmf [1] or methods therein (e.g. GAIN) (if you think the paper is not peer-reviewed yet) should be compared with.

There are several typos and mistakes:
 - Page 3, "matrixV"=> "matrix V"
 - Page 4, "sqaure root" => "square root"
 - Page 4, "In practice, We use" => "In practice, we use"
 - Page 6, "training process with patience 100," ?!
 - Table 1, to tight table caption

[1] Taguchi, H., Liu, X., & Murata, T. (2020). Graph Convolutional Networks for Graphs Containing Missing Features. arXiv preprint arXiv:2007.04583.

Summary and contributions: Briefly summarize the paper and its contributions:
The paper considers node representation learning in the setting of incomplete node attributes. The proposed method transforms node features to a latent space endowed with the Wasserstein metric, uses the Wasserstein barycenter of nodes’ neighbors to recover missing features, and finally transforms them back to their original space using a proposed inverse mapping. The algorithm has been tested on benchmark classification tasks as well as matrix completion tasks.

Strengths: 
+ The paper appears to be a novel application of ideas from both matrix completion and optimal transport. 
+ In the node classification experiments, multiple benchmark datasets were used and multiple methods to complete missing values were considered as baselines. 

Weaknesses: 
- While the experiments show promising results, the authors do not explain the choices  behind their method. 
- Why do the authors use SVD and define the distance between transformed features in the way they did? Is there a reason for using the exponential function in the transformation other than for inducing positiveness of probability? How does the inverse mapping come up and why does it work? 
- In the classification task experimental section, it is not clear why methods, like IGMC (used in Section 4,2) paired with an MLP, weren’t tested for the classification tasks. As for the choice of architectures, is there a reason for using 7 WGD layers versus 2 GCN layers for the baselines?
- Because many choices are not justified or described, a number of questions are left unanswered. What is fundamentally different between a WGD layer and a GCN layer? What is the influence of the number of WGD layers? And what is changing in the representation after each WGD layer? What does the induced latent space look like? How does it relate to the complete feature information setting? Is the dimension of the latent space equal to that of the original space? Why does feature recovery using Wasserstein barycenter update trumps simple average or k-means of neighbors, what makes the metric be more adequate than Euclidean distance? 
- The authors claim that their work “is the ﬁrst work to compute embeddings of a graph with incomplete attributes directly.” but their proposed algorithm is two-stage, first matrix completion and then neural network training. The paper seems to be more relevant to data pre-processing than representation learning itself. 
- The authors conducted experiments in two settings, partially missing and entirely missing. It seems that the performance of the proposed method is exactly the same in both settings. (Figure 1) This is curious. Can the authors comment on what is happening in this case?
- In multiple instances, the authors introduce experiments but don’t comment on them or explain the similarity or differences in performance between all of these experimental settings.

Clarity: 
Many terms are used without properly defining them; the equivalence of p-Wasserstein distance to eq. (4) is unclear unless readers are already familiar with optimal transport; the iterative Bregman projection and Gram-Schmidt process have not been elaborated clearly. The statement “However, most of the methods, which embed nodes into a lower-dimensional Euclidean space, suﬀer from common limitations” is in reference to GCNs which are not limited to low-dimensional spaces. So it’s not clear what the authors meant by this.

Relationship to prior work: 
The state-of-the-art alternatives have not been well-addressed and compared with. A lot of work has been done with graph convolutional networks and in this literature, the issue of missing attributes is very common. There is no mention of matrix completion methods in related works, these methods were used in Section 4.2, without explaining how they are different from the proposed method, or why this set of methods was selected for comparison.

Additional feedback:
Multiple typos are present throughout the paper, these can be addressed to improve the readability, for example:  In section 3.3 “expect” should be “except” in the sentence: Note that it is similar to the message aggregation in graph neural networks, expect we do it in a Wasserstein space and introduce no parameters.
In section 4.1 “some entities of the feature matrix is missing” should be “some entities of the feature matrix are missing”. “90% attributes” should be “90% of attributes”. “Pumbed” should be “Pubmed” in “on Cora and Pumbed datasets.”

This work studies the missing data issue in the graph-based learning tasks, including node classification and matrix completion. The idea is to perform dimension reduction of the observed features. Then, allow these features to diffuse over graphs by computing barycenter update to fill in the missing features. The prediction is performed by adding MLP over the obtained features. The paper is written well and easy to follow. However, the idea reads very heuristic with some potentially issues. The experiments read weak. 

First, regarding originality, this work seems to combine some previously well-studied techniques: Use optimal transport to complete missing data; Graph-based semi-supervised learning. For example, [1] discussed using optimal transport to fill in missing data, which is not referred in this work. [2] studied label propagation over graph via optimal transport, which is also not referred in this work. The idea in this work seems to be a combination of techniques from different previous works. It is okay to combine techniques but one needs to sufficiently demonstrate the effectiveness, which have not been well done in this work. 

Second, it reads weird to use exp() followed by a softmax to obtain a discrete distribution. "exp" function is not properly transformed function to keep the feature information. For example, if some components in the features obtained by SVD have a few negative ones and a few large positive components. Of course, both the negative and positive features are important. However, if we adopt the method proposed in the work, the model only tends to capture those large positive components and ignores those negative components (due to the exp function). Later, the authors use l2-distance (a very component-wise well-balanced metric) to define W-distance. This yields a very imbalanced treatment on the features.

Third, the experimental section is not persuasive. It is widely know that GCN does not work very well for node classification tasks. There are many other models, e.g., APPNP [3], performing  good node classification over homorphilic networks. These better baselines should be used. Moverover, for only topological information, GCN without node features is not a strong baseline. Better methods with only graph structures should be traditional methods, such as label propagation. Moveover, to demonstrate the effectiveness of this approach, experiments should also be done over heterophilic networks.                



[1] Missing Data Imputation using Optimal Transport, Muzellec et al., ICML 2020
[2] Wasserstein Propagation for Semi-Supervised Learning, Solomon et al., ICML 2014.
[3] Prediction and propagation, 2020

Meta-review: The rationality of the proposed method, especially its implementation detail, is challenged by the reviewers. Additionally, the experimental part and the writing of the paper should be improved. According to the feedback of the reviewers, I don't think this work is qualified enough at its current status. 

Decision: withdrawn-rejected-submissions