Venue: ICLR 2021 Conference

Reviews:

final_score: rating: 4: Ok but not good enough - rejection
review_content: **Summary**
The paper proposes OT-GNN, which incorporates optimal transport distance to message passing of GNN. The message passing is aggregated by using a Wasserstein discrepancy for a point cloud. The contrastive regularization is utilized to overcome extreme clustering of nodes of the same class. Also, in theory, the author shows that the Wasserstein kernel is universal while the "agg" kernel is not. The proposed model is tested on several molecular property prediction tasks. The OT-GNN achieves slightly better performance against existing methods.

**Pros**
1. The paper introduces a non-Euclidean metric in message-passing aggregation and develops OT-GNN based on it. The proposed OT-GNN has good performance for molecular property prediction tasks.
2. The paper provides a theoretical analysis of the universality and positiveness of the Wasserstein kernel. It shows the Wasserstein metric-induced message passing is superior to the Euclidean metric case in terms of universality. The paper also shows the $L_2$ Wasserstein kernel is not conditionally negative positive.
3. The computational complexity is analyzed for Wasserstein optimal transport in the OT-GNN. 

**Cons**
1. For Wasserstein discrepancy in (4), there seem many pairs of $H$ and $Q_i$, where $Q_i$ contains a set of free parameters. Then the OT-GNN would have too many parameters to train. Will it bring about overparameterization?
2. Using the Wasserstein metric, the computational cost will increase much. Besides, the contrastive regularization will add more training time. What is the GPU wall time for OT-GNN training for the molecule regression tasks?
3. The theoretical analysis for the properties of Wasserstein kernel seems not complicated, and it does not give a good interpretation of the learning ability, like universal approximation property or generalization of OT-GNN. In particular, the discussion of positiveness of the kernel is of no significance or use.
4. The authors also need to try more public graph datasets, such as Open Graph Benchmark, https://ogb.stanford.edu/.
5. Line 4 below (6): the last sentence seems not finished. 
review_title: Review for Optimal Transport Graph Neural Networks

final_score: rating: 5: Marginally below acceptance threshold
review_content: This paper combines OT  with parametric graph neural network. It replace the inner product between the graph embedding and the first layer weights of MLP by the Wasserstein distance between the node embeddings and some point clouds. Then the GNN, point clouds and the downstream MLP are trained in an end-to-end way. A regularization term is adopted to enforce the point clouds are not collapsed. The authors then theoretically show that the Wasserstein kernel is universal.

Pros:
1. The empirical study is thorough and the model has good empirical performance. 
2. The idea of using prototypes is interesting. That being said, I think the word "prototype" here is actually misleading. In Snell 2017, their prototypes is a representation of the data. In contrast, here Q_i is NOT the representation of H, since we are not minimizing the Wasserstein distance between them. Instead, Q_i just helps to keep some information that is useful for the task. 

Areas to improve:
1. My major concern is about computation time.  In each training step, the algorithm computes batch size*(the range of indices i) of OT problems. Will the OT part introduce a large training time overhead? Could you please report the training time of each method?
2. Since the model has significantly more parameters, e.g., in Q_i, it would be better to also compare it to a baseline model with comparable number of parameters and the node embedding info. For example, is it possible to compare it with: randomly selecting k embeddings, concatenate, and feed into MLP (adjusting k to make the parameter size comparable)?
3. Some notations are not properly defined. For example, \mathfrak{agg} first appears in section 4.1, but its formal definition is in section 4.2. Another example is indices i and j in eq (4). What is the range of i and the range of j? In other words, how is the range of i, the range of j, the dimension of q_i^j corresponds to n_pc, pc_size, pc_hidden in the appendix? 
4. Could you please provide some intuition that different variants of the proposed model performs the best for different dataset (table 1)? In other words, what variants is suitable for what kind of dataset?
review_title: Interesting Wasserstein kernel, good empirical performance

final_score: rating: 5: Marginally below acceptance threshold
review_content: This paper introduces OT-GNN, a combination of optimal transport and graph neural network, for graph-level tasks. Instead of the conventional graph embedding+MLP strategy, the authors propose to replace the readout layer in MPNN-like structures with Wasserstein discrepancy. In addition, the noise contrastive regularizer is added to the model so that the optimal transport plan is discriminative, and the new model could outperform its Euclidean counterpart. While the construction is sophisticated, the core idea behind the algorithm is not hard to follow. However, I do have some concerns regarding the presentation and experiments, which, at this stage, prevent me from recommending the paper confidently to the broader community.


Below I would like to list my main concerns or confusions regarding the paper.

The construction of prototypes is not very clear to me. I would expect the authors to provide more discussion before the experiment section regarding its structure, requirements, or formulation details.
The authors, if my understanding is correct, named the same subject to different terms as ‘prototypes’ and ‘free parameters’. This indeed creates unnecessary hurdles of understanding the paper, and I would recommend the names could be identified, or at least clearly and properly defined. 
The design of Figure 3 is confusing. What is the difference between Prototype 1-5 and Prototype 6-10? As they are designed with different shapes, I would expect them to represent different kinds of prototypes. However, unless I missed something, I didn’t find a clear statement of it. Also, I’m not sure what is the purpose of showing five real molecules? Are we supposed to compare them with the corresponding prototypes?
What is the partition function you mentioned in the 3rd line after Eq.6? Do you use different terminology from in Nickel & Kiela (2017)? 
The authors only compared their method with naive node embedding+sum aggregation models but ignored a major type of graph representation learning methods as graph pooling. This branch is developed rapidly in the recent 2 years, and many simple yet powerful methods have been proposed. I don’t think it’s fair to compare with the models that are specialized in node embedding tasks but neglect the real methods that are designed for graph embedding tasks. 
Following the last point, I would expect more recent research to be included in the Related Work section.
Some figures could be redesigned for better presentation. For example, it’s hard to read the axis information in Figure 4.

Minor typos:

(line 2, Section 2.1) “…property prediction by Yang et al. Yang et al. (2019).”
(line 1, Section 2.2) “Optimal Transport (OT) Peyré et al. (2019) is a mathematical framework…”

Given the problems stated as above, I would suggest the paper to be carefully polished before it’s ready to publish.
review_title: Review for OT-GNN

final_score: rating: 7: Good paper, accept
review_content: ## Summary

The paper introduces a novel approach to aggregate information of graph neural network node embeddings in order to support graph-level machine learning, like graph classification or graph regression. The aggregation is performed by comparing the node embeddings of a graph to learned, prototypical node embeddings via the Wasserstein distance. The vector of distances to the prototypes then serves as a concise representation that can be fed into a subsequent standard multi-layer perceptron for classification or regression. This approach is proven to be strictly more powerful than just adding up the node embeddings, which is the most common state-of-the-art. Finally, the paper evaluates the proposed aggregation function on two graph classification and two graph regression data sets, showing superior performance in all cases.

## Strengths

The paper has several strengths worth highlighting:

* The idea of using prototypical node embeddings as a method for aggregation is easy to understand and potentially interpretable. For example, one could try to construct graphs that best correspond to each prototype and thereby gather understanding what typical graphs look like based on which the model makes its decision.
* The idea also corresponds well to past research on kernels as well as distances, thus nicely integrating the strengths of graph neural nets and kernel/distance theory. The use of the Wasserstein distance to compare point clouds is particularly compelling and combines two hot topics in machine learning today, namely graph neural nets and the Wasserstein distance, which should be interesting to a broad section of the ICLR community.
* The theoretical results, especially the universal approximation result, are compelling and further make clear that the present work is strictly more powerful than the state of the art (i.e. just summing up all node embeddings)
* The experimental evaluation is performed against multiple reasonable baselines, including ablation variants of the proposed model, on several data sets, covering both regression and classification problems. Multiple repeats are performed, giving an impression of variance. Additionally, for the regression data sets, the paper shows both correlation as well as dimensionality reduction results indicating that the proposed representation is more smooth in relation to the targets compared to existing models.

## Weaknesses

That being said, I also see some issues which could be addressed to make the paper even stronger.

* Most importantly, from my perspective, the relation to existing work warrants much more discussion in two directions:
    * Currently, the discussion of prototype learning is reduced to the work of Snell (2017), whereas a long research tradition of learning prototypes for classification and regression is currently not covered. In particular, learning vector quantization approaches (refer e.g. to [Sato & Yamada; 1995][GLVQ]) have been used to learn prototypes that represent classes well and can then be used for a one-nearest neighbor classification. Even more similar to the present paper, radial basis function networks (e.g. [Chen, Cowan, and Grant, 1991][RBF]) have been used to solve regression tasks in the space of RBF kernel values to prototypes. It would be, I believe, helpful to contrast the present prototype learning approach to these past works. Additionally, these past works may give inspiration how to learn prototypes well and how to solve the issues of prototype collapse.
    * Currently, the paper treats distances as a kind of non-definite kernel. This can be done, but is unusual. An understanding of kernel and distance as opposing (but related) concepts is more common and, I believe, more instructive. However, the use of distances as features can still be justified, e.g. using the work of [Pekalska & Duin (2012)][DistanceTheory]. Similarly, it would be helpful to relate the present work to attempts that build bridges between kernels and neural nets, such as [Wilson, Hu, Salakhutdinov, and Xing (2016)][DeepKernels]
* I see several issues with how prototypes are integrated into the model.
    * It is not clear to me why the distances to the prototypes are plugged into another MLP instead of using them directly, at least for classification. In particular, prior work of Snell (2017) as well as [learning vector quantization][GLVQ] suggests to associate each prototype with a label and learn the prototypes such that a one-nearest-neighbor classification solves the classification problem. For regression, the issue is a bit more difficult, but here, as well, I would expect a single additional layer, similar to [RBF networks][RBF] and not several additional layers. Ideally, it would be beneficial to provide comparisons to these alternatives as additional ablation results in the appendix.
    * Using prototypes introduces two new hyperparameters, namely the number of prototypes and the size of points within each prototype. The paper currently gives no guidance on how to choose these hyperparameters.
    * Prototypes have the advantage of being a sparse representation of the data which lends itself to interpretation. However, currently no such interpretation is attempted. This seems to me like a missed opportunity. Showing an example graph for each prototype (i.e. a graph that is close to one prototype but far from every other prototype) could give additional insights into how the model represents the data.
    * The regularizer suggested in the paper does not seem sufficiently justified to me, yet. In particular, there are exponentially many 'bad' transport plans which I would not like to choose. How can I guarantee that my set N(T_i) is sufficiently rich to 'protect' me against all bad transport plans? Conversely, if the aim is just to prevent collapsing of the prototypes, would there not be simpler ways to prevent such a collapse, such as regularizing with the negative log of pairwise distances between points inside a prototype?
* Leaving the optimal transport plan T constant during backpropagation seems reasonable to me but the justification appears still rather hand-wavy (this applies to the cited paper of Xu (2019) as well). After inspecting the original paper of Afriat (1971), I am not ale to see how the results described there are applicable to transport plans (I don't even see an 'envelop theorem' in there). The connection may be there, but it is not obvious to me. Going back to a higher-level intuition, I think the key issue here is that even a very slight change in the points of a prototype set could dramatically change the optimal transport plan (that is just how transport plans behave, is my understanding; whenever another assignment becomes slightly better, suddenly the assignment changes in a discontinuous fashion). This speaks against leaving it constant. In the worst case, it could occur that the current transport plan suggests to move a point into one direction, but the transport plan at the next point suggests to move it back, yielding an endless cycle. One could, however, approximate the transport plan smoothly, e.g. as a mixture of several transport plans between which we interpolate continuously, and then one could make the argument that small gradient steps do not change the plan much, hence we can leave it constant. But again: Such an argument is not trivial and is currently missing. Furthermore, it would be good to clarify that the transport plan is only treated as constant for backpropagation, but that it is updated after each gradient step (that is, at least, how I interpret the approach).
* It would help to qualify Theorem 1 a little. First, Theorem 1 does not imply that the Wasserstein kernel can distinguish between all graphs; only between the graphs that are distinct in the multi-set of node embeddings, i.e. all Weisfeiler-Lehman distinguishable graphs. Second, Theorem 1 does not imply that the model as proposed is a universal approximator because the proposed model only uses a finite set of prototypes to represent the data. To make it a universal approximator, one would need infinitely many prototypes, if I am not mistaken. In practice, this is likely not an issue because few prototypes are sufficient to cover the space pretty well. It is still a limitation to how much Theorem 1 actually tells us about the proposed model.
* The experiments currently do not report empiric runtime results. These would be a valuable addition because they illustrate how much overhead prototypes introduce compared to just summing up the final layer.

There are also a few smaller issues which could be clarified:

* page 3: I wonder why the marginal distributions over points in X and points in Y are restricted to be uniform. Would the more general case not be interesting as well, where, say, prototypes receive more weight if they are more important for representing the data?
* Page 3: The vectors r_i are not introduced. Are these supposed to be the weight vectors of neurons in the MLP?
* Page 3: It is not clear to me why the negative inner product and the Euclidean distance should yield the same transport plan. I can see that for the special case of vectors with norm 1 where ||x - y||^2 = 1 - 2 * <x, y> + 1, but not in the general case.

## Judgment

Overall, I believe that this paper is strong enough to warrant acceptance. In particular, the heart of the paper, namely the proposed model architecture, the main theoretical result, and the experiments, are strong and my issues - namely related work discussions, some theoretical fine points, and ablation studies - are aspects that can be addressed during a revision without affecting the core paper too much. Furthermore, I believe that the core idea - combining prototypes, optimal transport, and graph neural nets - has great potential to make graph neural nets stronger and more interpretable, which are both crucial goals for the research community right now. As such, I am confident that the paper, as it stands, is worthy of publication and, with some additional work during revision, can become even better.

## References

* Sato, A., & Yamada, K. (1995). Generalized Learning Vector Quantization. Proceedings of NIPS 1995. [Link][GLVQ]
* de Vries, H., Memisevic, R., & Courville, A. (2016). Deep Learning Vector Quantization. Proceedings of ESANN 2016. [Link][DLVQ]
* Chen, S., Cowan, C., and Grant, P. (1991). Orthogonal least squares learning algorithm for radial basis function networks. IEEE Transactions on neural networks, 2(2). [Link][RBF]
* Duin, R., and Pekalska, E. (2012). The dissimilarity space: Bridging structural and statistical pattern recognition. Pattern Recognition Letters, 33(7). doi:[10.1016/j.patrec.2011.04.019][DistanceTheory]
* Wilson, A., Hu, Z., Salakhutdinov, R., and Xing, E. (2016). Deep Kernel Learning. Proceedings of AISTATS 2016. [Link][DeepKernels]

[GLVQ]:https://papers.nips.cc/paper/1113-generalized-learning-vector-quantization "Sato, A., & Yamada, K. (1995). Generalized Learning Vector Quantization. Proceedings of NIPS 1995."
[DLVQ]:https://www.elen.ucl.ac.be/Proceedings/esann/esannpdf/es2016-112.pdf "de Vries, H., Memisevic, R., & Courville, A. (2016). Deep Learning Vector Quantization. Proceedings of ESANN 2016."
[RBF]:https://eprints.soton.ac.uk/251135/1/00080341.pdf "Chen, S., Cowan, C., and Grant, P. (1991). Orthogonal least squares learning algorithm for radial basis function networks. IEEE Transactions on neural networks, 2(2)."
[DistanceTheory]:https://doi.org/10.1016/j.patrec.2011.04.019 "Duin, R., and Pekalska, E. (2012). The dissimilarity space: Bridging structural and statistical pattern recognition. Pattern Recognition Letters, 33(7)."
[DeepKernels]:http://proceedings.mlr.press/v51/wilson16.html "Wilson, A., Hu, Z., Salakhutdinov, R., and Xing, E. (2016). Deep Kernel Learning. Proceedings of AISTATS 2016"

## Clarifications

While I am quite confident in my review, it would help if the authors could verify that I understood some key points in the paper correctly:

* Each prototype here is a collection of points, and the number of prototypes as well as the number of points is a fixed hyperparameter, whereas the location of each point is learned, correct?
* The pipeline for classification/regression is to first embed the nodes of a graph via a graph neural net, yielding a point cloud with one point per node; second, to compare this point cloud to all prototypes using the Wasserstein distance/inner product; third to collect the resulting distances/inner products into a vector d/k with as many entries as there are prototypes and to feed this vector into a standard MLP to make the final classification/regression decision, correct?
* The optimal transport plan T is treated as constant for backprop, but is recomputed after each gradient step, correct?
* Are the limitations to Theorem 1 described above correct or did I misunderstand the Theorem?

## Typos

Beyond the points mentioned above, there are also a few typos and very minor things that do not affect my judgment but could easily be fixed to improve the paper.

* Abstract: The abbreviation 'OT' is not introduced
* page 1: 'As a result, some of the information naturally extracted by node embeddings may be lost' <-- This is correct, but could be underlined more with a citation
* page 2: 'that is universal approximator.' -> that is a universal approximator
* page 2: The abbreviation 'D-MPNN' is not introduced
* Figure 3 is currently rather busy and not very easy to interpret. Would it, perhaps, suffice to only show prototype 1 and one reference molecule?
* As far as I know, ICLR permits an appendix. As such, it would be good to move the appendix into the main paper.
review_title: A strong combination of optimal transport and graph neural nets, with some room for improvement with respect to prototypes

Meta-review: This paper proposes a very interesting approach using Wasserstein distance between graph embedded by GNN to perform prediction. The paper is well written and the experiments suggest that the method works  well in practice.

Several reviewers had some concerns about  computational complexity and model parameters that were very well answered during the discussion and with the new version of the paper but it was not enough to convince them to change their overall opinion on the paper. 

Note that a lengthy discussion with the reviewers stemming from an unclear Figure 3  was done about the need for a  contrastive regularization that raised important questions that should be addressed. In short, despite the claim by the authors that the contrastive regularization is important the experiments show very little difference in performances with or without the regularization which asks the question of its usefulness. As a matter of fact looking at Figure 3 the regularization will spread the samples in the distributions, making the prototype more similar . The argument about the sample collapse is not good enough because if the sample collapse completely during the optimization the method converges toward prototype L2 which is not the case since the proposed approach is much better than L2 even without regularization. So if there is some collapse it actually serves the discrimination and leads to a better solution. Also the Wasserstsein can reverse the collapsing (the samples  are never exactly at the same position and the gradient can be very different) if it helps for the optimization problem as is well known from the Wasserstein GAN literature. 

Due to the remaining concerns of the majority of reviewers, the AC recommends a reject but encourages the authors to resubmit the paper after taking into account the reviewers comments and the questions about the regularization .
Decision: Reject