Optimal Transport Graph Neural Networks
Authors: Gary Bécigneul
Octavian-Eugen Ganea
Benson Chen
Regina Barzilay
Tommi S. Jaakkola
Venue:  (2021)
Keywords: graph neural networks
optimal transport
molecular representations
molecular property prediction

Abstract: Current graph neural network (GNN) architectures naively average or sum node embeddings into an aggregated graph representation---potentially losing structural or semantic information. We here introduce OT-GNN, a model that computes graph embeddings using parametric prototypes that highlight key facets of different graph aspects. Towards this goal, we are (to our knowledge) the first to successfully combine optimal transport with parametric graph models. Graph representations are obtained from Wasserstein distances between the set of GNN node embeddings and "prototype" point clouds as free parameters. We theoretically prove that, unlike traditional sum aggregation, our function class on point clouds satisfies a fundamental universal approximation theorem. Empirically, we address an inherent collapse optimization issue by proposing a noise contrastive regularizer to steer the model towards truly exploiting the optimal transport geometry. Finally, we consistently report better generalization performance on several molecular property prediction tasks, while exhibiting smoother graph representations.

Under review as a conference paper at ICLR 2021
OPTIMAL TRANSPORT GRAPH NEURAL NETWORKS
Anonymous authors
Paper under double-blind review
ABSTRACT
Current graph neural network (GNN) architectures naively average or sum node embeddings into aggregated graph representation-ppoeenialll losing structural or semantic information. We here introduce OT-GNN, model that computes graph embeddings using parametric prototypes that highlight key facets of different graph aspects. Towards this goal, we are (to our knowledge) the first to successfully combine optimal transport (OT) with parametric graph models. Graph represen- tations are obtained from Wasserstein distances between the set of GNN node embeddings and "prototype" point clouds as free parameters. We theoretically prove that, unlike traditional sum aggregation, our function class point clouds satisfies a fundamental universal approximation theorem. Empirically, we address inherent collapse optimization issue by proposing noise contrastive regularizer to steer the model towards truly exploiting the optimal transport geometry. Finally, we consistently report better generalization performance several molecular property prediction tasks, while exhibiting smoother graph representations.
INTRODUCTION
Recently, there has been considerable interest in developing learning algorithms for structured data such as graphs. For example, molecular property prediction has many applications in chemistry and drug discovery (Yang al. 2019 Vamathevan et al.) 2019). Historically, graphs were decomposed into features such as molecular fingerprints, or turned into non-parametric graph kernels (Vish- wanathan et al. 2010 Shervashidze et al., 2011). More recently, learned representations via graph neural networks (GNNs) have achieved state-of- the-art on graph prediction tasks (Duvenaud et al.
2015 Defferrard et al. 2016 Kipf & Welling 2017 Yang al. 2019).
Despite these successes, graph neural networks are often underutilized in whole graph prediction tasks such as molecule property prediction. Specifically, while GNNs produce node embeddings for each atom in the molecule, these are typically aggregated via simple operations such sum or average, turning the molecule into a single vector prior to classification or regression. As a result, some of the information naturally extracted by node embeddings may be lost.
Departing from this simple aggregation step, Togninalli et al. (2019) recently proposed kernel function over graphs by directly comparing non-parametric node embeddings as point clouds through optimal transport (Wasserstein distance). Their non paraamettrii model yields better empirical per- formance over popular graph kernels, but this idea hasn't been extended to the more challenging parametric case where optimization difficulties have to be reconciled with the combinatorial aspects of optimal transport solvers.
Motivated by these observations and drawing inspiration from prior work on prototype learning (appendix F). we introduce a new class of GNNs where the key representational step consists of comparing each input graph to set of abstract prototypes (fig. I These prototypes play the role of basis functions: they are stored as point clouds as if they were encoded from actual real graphs. Each input graph is first encoded into a set node embeddings using any existing GNN architecture. We then compare this resulting embedding point cloud to the prototype embedding sets. Formally, the distance between two point clouds is measured by their optimal transport Wasserstein distances. be understood as keys that highlight property values
The prototypes as abstract basis functions associated with different graph structural features. In contrast to previous kernel methods, the prototypes are learned together with the GNN parameters in an end-to-end manner.

---

Under review as a conference paper at ICLR 2021
Toxic?
(a)
(b)
(c)
Figure 1: Our OT-GNN prototype model computes graph embeddings from Wasserstein distances between (a) the set GNN node embeddings and (b) prototype embedding sets. These distances then used as the molecular representation (c) for supervised tasks, e.g. property prediction. We assume that a few prototypes, e.g. some functional groups, highlight key facets or structural features of graphs relevant a particular downstream task at hand. We express graphs by relating them these abstract prototypes represented as free point cloud parameters.
Our notion of prototypes inspired from the vast prior work on prototype learning which we highlight in appendix F In our case, prototypes are not required to be the mean of a cluster of data, but instead they are entities living in the data embedding space that capture helpful information for the task under consideration. The closest analogy are the centers of radial basis function networks
(Chen et al. 1991 Poggio & Girosi 1990), but we also inspire from learning vector quantization approaches Kohonen 1995) and prototypical networks (Snell al., 2017).
Our model improves upon traditional aggregation by explicitly tapping into the full set node embeddings without collapsing them first to single vector. We theoretically prove that, unlike standard GNN aggregation, our model defines a class of set functions that is a universal approximator.
Introducing prototype points clouds as free parameters trained using combinatorial optimal transport solvers creates a challenging optimization problem. Indeed, as the models are trained end-to-end, the prototypes
the primary signal is initially available only in aggregate form. If trained often collapse to single points, reducing the Wasserstein distance between point clouds to Euclidean comparisons their means. To counter this effect, we introduce contrastive regularizer which effectively prevents the model from collapsing (Section 3.2). We demonstrate its merits empirically.
Our contributions. First, we introduce an efficiently trainable class of graph neural networks enhanced with optimal transport (OT) primitives for computing graph representations based on relations with abstract prototypes. Second, we are the first to successfully train parametric graph models together with combinatorial OT distances, despite optimization difficulties. A key element is our noise contrastive regularizer that prevents the model from collapsing back standard summation, thus fully exploiting the OT geometry. Third, we provide theoretical justification of the increased representational power compared to the standard GNN aggregation method. Finally, our model shows consistent empirical improvements previous state-of-the- -ht molecular datasets, yielding also smoother graph embedding spaces.
PRELIMINARIES
2.1 DIRECTED MESSAGE PASSING NEURAL NETWORKS (D-MPNN)
We briefly remind here of the simplified D-MPNN Dai et al. 2016) architecture which was adapted for state-of-tt the-art molecular property prediction by Yang et al. (2019). This model takes as input directed graph G (V, E), with node and edge features denoted by x, and eaw respectively, for w in the vertex set and v > w in the edge set E. The parameters D-MPNN are the matrices [W, W m. W.). It keeps track of messages m and hidden states h for each step t, defined as follows. An initial hidden state is set to h = ReLU (LLLÜWWWcoo Wxcat(x,,, Cow)) where "cat" denotes concatenation. Then, the updates are:
H+1 P2u
Σ
h
h I+1 Re + W m(1)
(+1)
(1)
(EN(
2

---

Under review as a conference paper at ICLR 2021
where (T)) = {k E V (k,, v) E} denotes e's incoming neighbors. After T steps of message passing, node embeddings are obtained by summing edge embeddings:
m Σ
h, = (W C.eLU(W.cat(xxs ocat(xx,
(2)
wEN(2)
final graph embedding then obtained h EvEy hy, which usually fed multilayer perceptron (MLP) for classification or regression.
2.2 OPTIMAL TRANSPORT GEOMETRY
mathematical framework that defines distances or similari- ties between objects such as probability distributions, either discrete or continuous, as the cost of an optimal transport plan from one the other.
Optimal Transport (Peyré al. 2019)
W(X. Y) mingelgy Enj T,,o(x. Vy)
Ti W
X2
X3
X4
X5
33
Figure 2: We illustrate, for a given 2D point cloud, the optimal transport plan obtained from minimizing the Wasserstein costs; e( denotes the Euclidean distance. A higher dotted-line thickness illustrates greater mass transport.
Wasserstein distance for point clouds. Let a point cloud X = {x4]==1 of size n be a set of n points x; E Rdd. Given point clouds X, Y of respective sizes n, m, a transport plan (or coupling) is matrix T of size n X with entries in [0, 1J. satisfying the two following marginal construints:
T1m " 1, and T 1,, = Im Intuitively, the marginal constraints mean that T preserves mass from X Y. We denote the set of such couplings CXY.
Given a cost function on Rd, its associated Wasserstein discrepancy is defined as
W(X,)
min TECXY
Te,((x,,
(3)
f) We further describe the shape of optimal transports for point clouds same sizes in Appendix B.3
MODEL PRACTICE
3.1 ARCHITECTURE ENHANCEMENT
Reformulating standard architectures. As mentioned at the end of Section 2.1 the final graph embedding h Evev h, obtained by aggregating node embeddings is usually fed performing matrix-multipliattion whose i-th component is (Rh); (r:, h). where r; is the i-th row matrix R. Replacing (sr) by a distance/kernel k((:.) allows the processing of more general graph representations than just vectors in RRd, such point clouds adjacency tensors.
MLP
From a single point to a point cloud. We propose to replace the aggregated graph embedding h Evev h, by the point cloud (of unaggregated node embeddings) H [ev and the inner-products (h, rx) by the below written Wasserstein discrepancy:
W(H, Q = min TECHQ,
Σ q(),
vi
where Q; { bi.......... N), Vi represented as a set N embeddings as free trainable parameters, and the cost is chosen as c -|.-. 1 r) Note that both options yield identical optimal transport plans.
(1, (1...., M) represent M prototype point clouds each being
3

---

Under review as a conference paper at ICLR 2021
Greater representational power. We formulate mathematically in Section that this kernel has strictly greater representational power than the kernel corresponding to standard inner-produce on top of sum aggregation, to distinguish between different point clouds.
Final architecture. Finally, the vector of all Wasserstein distances in eq. (4) becomes the input to final MLP with a single scalar as output. This can then be used as the prediction for various downstream tasks. This model is depicted in fig. and motivated theoretically in section 4.1
3.2 CONTRASTIVE REGULARIZATION
What would happen to W(H, Q.) if all points qf belonging to point cloud Q2 would collapse to the same point q.? All transport plans would yield the same cost, giving for c c=-c-=--- (..)):
W(H, Q.)
Σ Tho -h, q./lV)).
(5)
In this scenario, our proposition would simply over-parametrize the standard Euclidean model.
(a) No contrastive regularization
(b) Using contrastive regularization (weight 0.1)
Figure 3: 2D embeddings of prototypes and of real molecule samples with (right) and without (left) contrastive regularization for runs using the same random seed, Points in the prototypes tend to cluster and collapse more when regularization is used, suggesting that the optimal transport plan no longer remains uniquely discriminative. Prototype 1 (red) is enlarged for clarity: without regularization is clumped together (left), but with regularization it is distributed across the space (right).
first obstacle and cause. Our first empirical trials with OT-enhanced GNNs showed that model trained with only the Wasserstein component would sometimes perform similarly to the Euclidean baseline both train and validation settings, in spite of its greater representational power.
Further investigation revealed that the Wasserstein model would naturally displace the points in each of its prototype point clouds in such way that the optimal transport plan obtained by maximizing En Tey (ba, q was not discriminative, i.e. many other transports would yield similar Wasserstein cost. Indeed, as shown in Eq. (5). if each point cloud collapses to its mean, then the Wasserstein geometry collaspses to Euclidean geometry. In this scenario, any transport plan yields the same Wasserstein cost. Further explanations are provided in Appendix A. and Figure
Contrastive regularization. To address this difficulty, we consider adding regularizer which encourages the model to displace its prototype point clouds such that the optimal transport plans would be discriminative against chosen contrastive transport plans. Namely, consider point cloud Y of node embeddings and let T be optimal transport plan obtained in the computation of W(Y, Q.). For each T we then build a set (T)( C CyQ. of noisy/cottrastive transports. If we denote by (r Y) == Ex Tx yr) the Wasserstein cost obtained for the particular transport T, then our contrastive regularization consists maximizing the term:
Σ log (- ETENeg(Ty + ( -
e WV.(...Q.)
-
(6)
4

---

Under review as a conference paper at ICLR 2021
which can be interpreted the log-likelihood that the correct transport T, be (as it should) better minimizer of Wr(Y, Y. Q₂) than its negative samples. This can be considered as an approximation of log( (Tr(TT Y. Q.)), where the partition function is approximated by selection of negative examples, as done e.g. by Nickel & Kiela (017). Its effect is shown in Figure 3
Remarks. The selection of negative examples should reflect the trade-off: (i) not be too large, for computational efficiency while (ïi) containing sufficiently meaningful and challenging contrastive samples. Details about choice of contrastive samples are given in Appendix A.2 Note that replacing the set Neg (Tg with singleton {T} for contrastive random variable T lets us rewrite Eg. 6 a
Σ. log (WWr WT:), reminiscent of noise contrastive estimation (Gutmann & Hyvärinen 2010
3.3 OPTIMIZATION & COMPLEXITY
Backpropagating gradients through optimal transport (OT) has been the subject of recent research investigations: Genevay et al. (2017) explain how to unroll and differentiate through the Sinkhorn procedure solving OT, which was extended by [Schmitz al. (2018) to Wasserstein barycenters. However, more recently, Xu (2019) proposed to simply invoke the envelop theorem (Afriat) 1971) to support the idea of keeping the optimal transport plan fixed during the back pack-propagttion of gradients through Wasserstein distances. For the sake of simplicity and training stability, we resort to the latter procedure: keeping T fixed during back- proppppopagttion. We discuss complexity in appendix
4 THEORETICAL ANALYSIS
In this section we show that the standard architecture lacks fundamental property of universal approximation of functions defined on point clouds, and that our proposed architecture recovers this property. We will denote by X the set of point clouds X {x.) of size in R&.
4.1 UNIVERSALITY
As seen in Section 3.1 we have replaced the sum aggregation - followed by the Euclidean inner- product by Wasserstein discrepancies. How does this affect the function class and representations?
common framework used to analyze the geometry inherited from similarities and discrepancies is that of kernel theory. A kernel & on a set X is symmetric function k X x X + R, which can either measure similarities or discrepancies. An important property of a given kernel on a space X' is whether simple functions defined top of this kernel can approximate any continuous function on the same space. This is called universality: a crucial property to regress unknown target functions.
Universal kernels. A kernel & defined on XT said be universal if the following holds: for any compact subset X C XY. the set of functions in the form Σÿ= a 0,) + B,) is dense in the set C(X') of continuous functions from X' IR w.r.t the sup norm || |loo. o denoting the sigmoid. Although the notion of universality does not indicate how easy it is in practice to learn the correct function, at least guarantees the absence of fundamental bottleneck of the model using this kernel.
In the following we compare the aggregating kernel agg(X, Y) = ( x,, E,
}3) (used by popular
GNN models) with the Wasserstein kernel, where
= min Σ |x, - 1,,1 Waot (X, Y) = max ΣTT((x, TECXY 8j
(7)
TECxY
Theorem 1. We have that:
1. The aggregation kernel agg is not universal.
2. The Wasserstein kernel WL, is universal.
Proof: See appendix B.I Universality of the Wia kernel comes from the fact that its square-root defines a metric, and from the axiom of separation of distances: if d(x,, y) = 0 then r = y.
where a(-)) is the sigmoid function. For E N 0,8, R and 8, Xz.
5

---

Under review as a conference paper at ICLR 2021
Implications. Theorem states that proposed OT-GNN model is strictly more powerful than the state of the art GNN models that use summatton or averaging of node embeddings. Nevertheless, this implics Weisfeiler-Lehman distinguishable graphs in the case of graph convolutional networks.
only distinguish graphs that have distinct multi-sets of node embeddings, e.g. all
In practice, the shape of the aforementioned functions that posses universal approximation capabilities gives an indication of how one should leverage the vector of Wasserstein distances to prototypes to perform graph classification - e.g. using a multilayer perceptron (MLP) on top.
4.2 DEFINITENESS
For the sake of simplified mathematical analysis, similarity kernlss often required be positive definite (p.d.,, which corresponds to discrepancy kernels being conditionally negative definite (c.n.d.). Although such a property has the benefit of yielding the mathematical framework of Reproducing Kernel Hilbert Spaces, essentially implies linearity, i.e. the possibility to embed the geometry defined by that kernel a linear vector space.
We now discuss that, interestingly. the Wasserstein kernel we used does not satisfy this property, and hence constitutes an interesting instance of a universal, non p.d. kernel. Let us remind these notions.
Tn Kernel definiteness. A kernel k is positive definite (p.d.) on X' if for E N", II,,, E X and C..... Cn E E. we have Eis c,cyk(2,, xy) 2 O. It conditionally negative definite (c.n.d.) on X for n E N", I1, ..... Fn E X' and C1, Cp. E such that E, C: we have Eij
S
These two notions relate to each other via the below result Boughorbel al. (0055
Proposition 1. Let k be symmetrie kernel X, let zo X and define the kernel:
k(II v) == -1 - k(æ,, xo) ) k(y, mo) + k(x0, ro)].
(8)
Then is p.d. if and only if k is c.n.d. Example: k == - 11 and xo yield le = (:5
One can easily show that agg also defines a p.d. kernel, and that agg(;;)) 2ennnWW (..)) However, the Wasserstein kernel is not p.d., as stated in different variants before (e.g. Vert 2008 and reminded by the below theorem. We here give novel proof in Appendix [..2
Theorem 2. have that:
1. The (similarity) Wasserstein kernel Waot not positive definite;
2. The (discrepancy) Wasserstein kernel W1,2 is not conditionally negative definite.
EXPERIMENTS
5.1 EXPERIMENTAL SETUP
We experiment on 4 benchmark molecular property prediction datasets (Yang al. 2019) including both regression (ESOL, Lipophilicity) and classification (BACE, BBBP) tasks. These datasets cover different complex chemical properties (e.g. ESOL water solubility, LIPO octanol/water distribution coefficient, BACE inhibition of human B-secretase, BBBP blood-brain barrier penetration).
We provide results for our implementation of 4 different widely used graph-based models.
Fingerprint + MLP applies a MLP over the input features which are hashed graph structures (called a molecular fingerprint). GIN the Graph Isomorphism Network from (Xu et al., (019). which is variant of GNN. The original GIN does not account for edge features, so we adapt their algorithm to our setting. Next, GAT the Graph Attention Network from (Veličković et al.) 2017). which uses multi-head attention layers propagate information. The original GAT model does not account for edge features, so we adapt their algorithm to our setting. More details about our implementation of the GIN and GAT models can be found in the appendix D.2 Finally. Chemprop D-MPNN (Yang et al., 2019) is a graph network that exhibits tfate-of-the-art performance for molecular representation learning across multiple classification and regression datasets. Empirically

---

Under review as a conference paper at ICLR 2021
we find that this baseline is indeed the best performing, and so is used to obtain node embeddings in all our prototype models. Its architecture is described section 2.1
Different variants of OT-GNN prototype model are described below:
Proto ProtoW-L2// LL2/Dot Dot is the model that treats point clouds as point sets, and computes the Wasserstein distances each point cloud (using either L.2 distance (minus) dot product cost functions) the molecular embedding. ProtoS-L2 is a special case of ProtoW-L2 in which the point clouds have single point and instead of using Wasserstein distances, we just compute simple Euclidean distances between the aggregated graph embedding and point clouds. Here, we omit using dot product distances, that model is mathematically equivalent the GNN model.
We use the the POT library (Flamary & Courty) 2017) to compute Wasserstein distances using the network simplex algorithm (Earth Movers distance), which we find empirically be faster than the Sinkhorn algorithm for our datasets. We define the cost matrix by taking the pairwise L2 or negative dot product distances. As mentioned Section 3.3 we fix the transport plan, and only backpropagate through the cost matrix for computational efficiency. Additionally, to account for the variable size of each input graph, we multiply the OT distance between two point clouds by their respective sizes. To avoid the problem of point clouds collapsing, we employ the contrastive regularizer defined in More details about experimental setup are presented in Appendix D.T
Section
3.2
ESOL (RMSE)
Lipo (RMSE)
BACE (AUC)
BBBP (AUC) # ggrhhs= 2039
Models
# ggrphs = 1128
# #grphs= 4199
#grphs= 1512
Baselines
Fingerprint + MLP
.92 .017
.855 .017
.80 + .007
9911 + .005
GIN
.665 = .026
.658 + .019
.861 = .013
.900 = .014
GAT
.654 .028
.888 + .047
.800 .011
.888 = .015
D-MPNN
.635 = .027
.646 + .041
.865 .013
.955 + .010
ProtoS-L2
.611 .034
5880
865 = .010
.918 = .009
ProtoW-Dot (no reg.)
.608 = .029
.637 + .018
867 = .014
.999 = .009
Ours
ProtoW-Dot
.594 + .031
.629 + .015
.81 + .014
.19 + .009
ProtoW-L2 (no reg.)
.616 = .028
.615 + .025
.870 + .012
.920 - .010
ProtoW-L2
.605 = .029
.604 + .014
.873 - .015
.920 + .010
Table 1: Results the property prediction datasets. Best model is in bold, second best is underlined.
Lower RMSE is better, while higher AUC is better. Wasserstein models are by default trained with contrastive regularization as described in section 3.2 and outperform those without.
5.2
EXPERIMENTAL RESULTS
5.2. REGRESSION AND CLASSIFICATION
Results are shown in Table T Our prototype models outperform state-of-the art GNN/D-MPNN baselines on all 4 property prediction tasks. Moreover, the prototype models using Wasserstein distance (Proto W L2 /Dot) achieve better performance on 3 out of 4 of the datasets compared to the prototype model using only Euclidean distances (ProtoS-L2). This indicates that Wasserstein distance confers greater discriminative power compared traditional aggregation methods.
Table 2: The Spearman and Pearson correlation coefficients the ESOL dataset for the GNN and ProtoW-T.2 model w.r.t. the pairwise dif- ference in embedding vectors and labels.
Spearman
Pearson r
D-MPNN
.424 + .029
.393 .049
(a) D-MPNN model
(b) ProtoW-I.2 model
ProtoS-L2
.561 + .087
.414 + 114
Figure 4: Comparison the correlation between graph embedding distances (X axis) and label distances (Y axis) on the ESOL dataset.
ProtoW-D
.592 + .150
.559 + .216
ProtoW-13
.815 .026
.828 .020
7

---

Under review as a conference paper at ICLR 2021
(a) ESOL D-MPNN
(b) ESOL P-o
(c) LIPO D-MPNN
(d) LIPO ProtoW-L.2
projections of molecular embeddings (before the last linear layer) w.r.t. their associated predicted labels. Heat colors are interpolations based only on the test molecules from each dataset. Comparing (a) vs (b) and (c) vs (d), we can observe a smoother space our model compared to the D-MPNN baseline (sce also main text).
Figure 5: 2D heatmaps of T-SNE Maaten & Hinton (0088)
5.2.2 NOISE CONTRASTIVE REGULARIZER
Without any constraints, the Wasserstein prototype model will often collapse the set of points in a point cloud into a single point. As mentioned in Section 3.2 we use a contrastive regularizer to force the model to meaningfully distribute point clouds in the embedding space. We show 2D embeddings in Fig. illustrating that without contrastive regularization, prototype point clouds are often displaced close to their mean, while regularization forces them nicely scatter. Quantitative results in Table also highlight the benefit of this regularization.
5.2.3 LEARNED EMBEDDING SPACE: QUALITATIVE AND QUANTITATIVE RESULTS
We further examine the learned embedding space of the best bascline (i.e. D-MPNN) and our best Wasserstein model. We claim that our models learn smoother latent representations. We compute the pairwise difference in embedding vectors and the labels for each test data point on the ESOL dataset. Then, we compute two measures of rank correlation, Spearman correlation coefficient (p) and Pearson correlation coefficient (r). This is reminiscent of evaluation tasks for the correlation of word embedding similarity with human labels (Luong et al. 2013).
Our ProtoW-1.2 achieves better p and r' scores compared to the D-MPNN model (Table 2p. that indicating Wasserstein model constructs more meaningful embeddings with respect the label distribution. Indeed, Figure 4 plots the pairwise scores for the D-MPNN model (left) and the ProtoW-L2 model (right). Our ProtoW-L2 model, trained to optimize distances in the embedding space, produces more meaningful representations with respect to the label of interest.
Moreover, qualitatively shown in Figure 5 our model provides more robust molecular embeddings compared to the bascline, in the following sense: we observe that a small perturbation a molecular embedding corresponds small change in predicted property value desirable phenomenon that holds rarely for the baseline D-MPNN model. Our Proto-W-LL model yields smoother heatmaps.
Average Train Time per Epoch (sec)
Average Total
Train Time (min)
ProtoW
ProtoW +
D-MPNN
ProtoS
ProtoW
D-MPNN
ProtoS
ProtoW
NC reg
NC reg
ESOL
5.5
5,4
13.5
31.7
4.3
5.3
15.4
30.6
BACE
32.0
31.4
25.9
39.7
12.6
12.4
25.1
49.4
BBBP
41.8
42,4
51.2
96.4
12.6
13.2
31.5
62.2
LIPO
33.3
35.5
27
69.4
28.6
31.8
74.7
146.9
Table 3: Training times for each model and dataset.
Model
Fingerprint MLP
GIN
GAT
D-MPNN
ProtoS
ProtoW
# Parameters
401k
626k
671k
100k
65k
66k
Table Number parameters per model. Corresponding hyperparameters are appendix D.I
8

---

Under review as a conference paper at ICLR 2021
5.2.4 ADDITIONAL EXPERIMENTAL DETAILS
Model Sizes Using MPNN hidden dimension as 200, and the final output MLP hidden dimension as 100, the number of parameters for the models are given by table 4 The fingerprint used dimension was 2048, explaining why the MLP has a large number of parameters. The D-MPNN model is much smaller than GIN and GAT models because it shares parameters between layers, unlike the others. Our prototype models are even smaller than the D-MPNN model because we do not require the large MLP at the end, instead parameters used for these point clouds). The dimensions of the prototype embeddings are also smaller compared the node embedding dimensions of the D-MPNN and other baselines. We did not see significant improvements in quality by increasing any of the hyperparameter values.
compute distances
few small prototypes (small number of overall
Remarkably, our model outperforms all the baselines using between 10 to 1.5 times less parameters.
Runtimes. We report the average total training time (number of epochs might vary depending on the early stopping criteria), as well as average training epoch time for the D-MPNN and our prototype models in table 3 We note that our method is between to 7.1 times slower than the D-MPNN baseline which mostly happens due to the frequent calls to the Earth Mover Distance OT solver.
6 RELATED WORK
Graph Neural Networks were introduced by Gori et al. (2005) and Scarselli et al. (2008) as form of recurrent neural networks. Graph convolutional networks (GCN) appeared later on various forms. Duvenaud et al. (201)); Atwood & Towsley (0166) proposed propagation rule inspired from convolution and diffusion, but these methods do not scale graphs with either large degree distribution or node cardinality. Niepert et al. (2016) defined GCN as ID convolution on a chosen node ordering. Kearnes et al. (2016) also used graph convolutions to generate high quality molecular fingerprints. Efficient spectral methods were proposed by Bruna et al. (2013); Defferrard et al. (2016). Kipf& & Welling (2017) simplified their propagation rule, motivated from spectral graph theory (Hammond et al. 22011). Different such architectures were later unified into the message passing neural networks framework by Gilmer et al. (2017). and applied to molecular property prediction. A directed variant of message passing from Dai et al. (2016) was later used improve state-of of the art in molecular property prediction on a wide variety of datasets by (Yang et al. (019). Other applications include recommender systems (Ying et al. (0088).. Inspired by DeepSets Zaheer et al. (2017), Xu et al.) (2119) suggest a simplified architecture called GIN, which theoretically can discriminate between any different local neighborhoods. Other recent approaches modify the uum-gggregation of node embeddings in the GCN architecture with the aim to preserve more information Kondor et al. (2018); this category there also the recently growing class of hierarchical graph pooling methods which typically either use deterministic and non-differentiable node clustering (Defferrard al. 2016,1 Jin et al.) 2018), or differentiable pooling (Ying al.) 2018b Noutahi et al.) 2019 Gao & Ji 2019). However, these methods are still strugling with small labelled graphs such as molecules where global and local node interconnections cannot be easily cast as a hierarchical interaction.
Pei et al. (2020).
Other recent geometry-inspired GNNs include adaptations to non-Euclidean spaces (L.iu et al. 2019
Chami et al. 2019 Bachmann et al. 2019).
We additionally discuss related work on prototype learning in appendix F
CONCLUSION
We propose OT-GNN: one of the first parametric graph models that leverages optimal transport to learn graph representations. It learns abstract prototypes as free parametric point clouds that highlight different aspects of the graph. Empirically, property prediction tasks, while the learned representations also exhibit stronger correlation with the target labels. Finally, universal approximation theoretical results enhance the merits of our model.
outperform popular baselines in different molecular

---

Under review as a conference paper at ICLR 2021
REFERENCES
SN Afriat. Theory of maxima and the method of lagrange. SIAM Journal on Applied Mathematics, 2033):43-357, 1971.
James Atwood and Don Towsley. Diffusion-convolutional neural networks. In Advances in neural information processing systems, pp. 1993-2001, 2016.
Gregor Bachmann, Gary Bécigneul, and Octavian- Eugen Ganea. Constant curvature graph convolu- tional networks. arXiv preprint arXiv:1911.0500 2019.
Sabri Boughorbel, J-P Tarel, and Nozha Boujemaa. Conditionally positive definite kernels for svm based image recognition. In 2005 IEEE International Conference on Multimedia and Expo, pp. 113-116. IEEE, 2005.
Joan Bruna, Wojeiech Zaremba, Arthur Szlam, and Yann LeCun. Spectral networks and locally connected networks graphs. arXiv preprint arXiv:1312.6203. 2013.
Kerstin Bunte, Petra Schneider, Barbara Hammer,
Schleif, Thomas Villmann, and Michael Bichl. Limited rank matrix learning, discriminative dimension reduction and visualization. Neural Networks, 26: 159-173, 2012.
Ines Chami, Zhitao Ying, Christopher Ré, and Jure Leskovec. Hyperbolic graph convolutional neural networks. In Advances in Neural Information Processing Systems, pp. 4869-4880, 2019.
Sheng Chen, Colin FN Cowan, and Peter M Grant, Orthogonal least squares learning algorithm for radial basis function networks. IEEE Transactions on neural networks, 2(2):302-309, 1991.
Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. In Advances in neural information processing systems, pp. 2292-2300, 2013.
George Cybenko. Approximation by superpositions of a sigmoidal function. Mathematics of control, signals and systems, 2(4):303-314, 1989.
Hanjun Dai, Bo Dai, and Le Song. Discriminative embeddings of latent variable models for structured data. In International conference machine learning, pp. 2702-2711, 2016.
Michaël Defferrard, Xavier Bresson, and Pierre Vandergheynst. Convolutional neural networks on graphs with fast localized spectral filtering. In D. D. Lee, M. Sugiyama, U. V. Luxburg, Guyon, and R. Garnett (eds.), Advances in Neural Information Processing Systems 29, pp. 3844-3852 Curran Associates, Inc., 2016.
Robert PW Duin and Elžbieta Pękalska. The dissimilarity space: Bridging structural and statistical pattern recognition. Pattern Recognition Letters, 33(7):826 832, 2012.
David K Duvenaud, Dougal Maclaurin, Jorge Iparraguirre, Rafael Bombarell, Timothy Hirzel, Alán Aspuru-Guzik, and Ryan P Adams. Convolutional networks on graphs for learning molecular fingerprints. In Advances in neural information processing systems, pp. 2224-2232. 2015.
Pavel Dvurechensky, Alexander Gasnikov, and Alexey Kroshnin. Computational optimal transport: Complexity by accelerated gradient descent is better than by sinkhorn's 's algorithm. arXiv preprint 1802.04367.. 2018.
Rémi Flamary and Nicolas Courty. Pot python optimal transport library. GitHub: http:://github com//flamary/POr. 2017.
Hongyang Gao and Shuiwang Ji. Graph u-nets. In International Conference Machine Learning, pp. 2083-2092 2019.
Aude Genevay, Gabriel Peyré, and Marco Cuturi. Learning generative models with sinkhorn diver- gences. arXiv preprint arXiv:1706.0023 00292, 2017.
Justin Gilmer, Samuel S Schoenholz, Patrick F Riley, Oriol Vinyals, and George E Dahl. Neural message passing for quantum chemistry. In Proceedings of the 34th International Conference on Machine Learning- Volume 70, pp. 1263-1272. JMLR. org, 2017.
10

---

Under review as a conference paper at ICLR 2021
Marco Gori, Gabriele Monfardini, and Franco Scarselli. A new model for learning in graph domains. In Proceedings. 2005 IEEE International Joint Conference on Neural Networks, 2005., volume 2, pp. 729-734, IEEE, 2005.
Michael Gutmann and Aapo Hyvärinen. contascon estimation: new estimation principle for unnormalized statistical models. In Proceedings of the Thirteenth International Conference on Artificial Intelligence and Statistics, pp. 297-304, 2010.
Barbara Hammer and Thomas Villmann. Generalized relevance learning vector quantization. Neural Networks, 15(8-9): 10991059-1000 1068, 2002.
David Hammond, Pierre Vandergheynst, and Rémi Gribonval. Wavelets on graphs via spectral graph theory. Applied and Computational Harmonic Analysis, 30(2):129-150, 2011.
Jørgen Hoffmaann 319-326, 1976.
Measures which agree on balls. Mathematica Scandinavica, 37(2):
Wengong Jin, Regina Barzilay, and Tommi Jaakkola. Junction tree variational autoencoder for molecular graph generation. arXiv preprint wrXiv:1802.04364 2018.
Steven Kearnes, Kevin McCloskey, Mare Berndl, Vijay Pande, and Patrick Riley. Molecular graph convolutions: moving beyond fingerprints. Journal of computer-aided molecular design, 30(8): 595-608, 2016.
Thomas N Kipf and Max Welling. Semi- Semi-supereised classification with graph convolutional networks. International Conference on Learning Representations (ICLR), 2017.
Teuvo Kohonen. Learning vector quantization. In Self-organizing maps, pp. 175-189. Springer, 1995.
Risi Kondor, Hy Truong Son, Horace Pan, Brandon Anderson, and Shubhendu Trivedi. Covariant compositional networks for learning graphs. In International conference on machine learning, 2018.
Qi Liu, Maximilian Nickel, and Douwe Kiela. Hyperbolic graph neural networks. In Advances in Neural Information Processing Systems, pp. 8228-8239, 2019.
Minh-Thann Thang Luong, Richard Socher, and Christopher D) Manning. Better word representations with recursive neural networks for morphology. In Proceedings of the Seventeenth Conference on Computational Natural Language Learning, 104-113, 2013.
Laurens van der Maaten and Geoffrey Hinton. Visualizing data using t-sne. Journal of machine learning research, 9(Nov):2579-2660 2008.
Maximillian Nickel and Douwe Kiela. Poincaré embeddings for learning hierarchical representations. 2017.
In Advances neural information processing systems, pp.
Mathias Niepert, Mohamed Ahmed, and Konstantin Kutzkov. Learning convolutional neural networks for graphs. In International conference machine learning, pp. 2014-2023, 2016.
Emmanuel Noutahi, Dominique Beaini, Julien Horwood, Sébastien Giguère, and Prudencio Tossou. Towards interpretable sparse graph representation learning with laplacian pooling, arXiv preprint arXiv:1905.1 115.7 2019.
Hongbin Pei, Bingzhe Wei, Chen-Chuar Chang Kevin, Yu Lei, and Bo Yang. Geom-gen: Geometric graph convolutional networks. In International conference machine learning, 2020.
Ofir Pele and Michael Werman. Fast and robust earth mover's distances. In 2009 IEEE 12th International Conference on Computer Vision, pp. 460-467. IEEE, 2009.
Gabriel Peyré, Marco Cuturi, et al. Computational optimal transport. Foundations and Trends in Machine Learning, 1 11(-6))355-607, 2019.
Tomaso Poggio and Federico Girosi. Networks for approximation and learning. Proceedings of the IEFF, 78(9): 1491- 1990.
11

---

Under review as a conference paper at ICLR 2021
Atsushi Sato and Keiji Yamada. Generalized learning vector quantization. Advances in neural information processing systems, 8:423 4299,4429110 1995.
Franco Scarselli, Marco Gori, Ah Chung Tsoi, Markus Hagenbuchner, and Gabriele Monfardini. The graph neural network model. IEEE Transactions on Neural Networks, 20(1):61-80, 2008.
Morgan Schmitz, Matthieu Heitz, Nicolas Bonneel, Fred Ngole, David Cocurjolly, Marco Cuturi, Gabriel Peyré, and Lean-Luc Starck. Wasserstein dictionary learning: Optimal transport- based unsupervised nonlinear dictionary learning. SIAM Journal on Imaging Sciences, 11()):4436678 2018.
Petra Schneider, Michael Biehl, and Barbara Hammer. Adaptive relevance matrices in learning vector quantization. Neural computation, 21(12)3532-3561, 2009.
Nino Shervashidze, Pascal Schweitzer, Erik Jan van Leeuwen, Kurt Mehlhorn, and Karsten M Borgwardt. Weisfeiler-lehman graph kernels. Journal of Machine Learning Research, 12(Sep): 2539-2561, 2011.
Jake Snell. Kevin Swersky, and Richard Zemel. Prototypical networks for few-shot learning. In Advances in neural information processing systems, pp. 4077-4087, 2017.
Matteo Togninalli, Elisabetta Ghisu, Felipe Llinares-López, Bastian Rieck, and Karsten Borgwardt. Wasserstein weisfeiler-lehman graph kernels. In Advances in Neural Information Processing Systems, pp. 6436-6446, 2019.
Jessica Vamathevan, Dominic Clark, Paul Czodrowski, Ian Dunham, Edgardo Ferran, George Lee, Bin Li, Anant Madabhushi, Parantu Shah, Michaela Spitzer, et al. Applications of machine learning in drug discovery and development. Nature Reviews Drug Discovery, 18(6):463-777, 2019.
Petar Veličković, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pietro Lio, and Yoshua Bengio. Graph attention networks, arXiv preprint wrXiv1170.11000 2017.
Jean-Philippe Vert. The optimal assignment kernel is not positive definite. arXiv preprint arXiv:0801.4061, 2008
S Vichy N Vishwanathan, Nicol N Schraudolph, Risi Kondor, and Karsten M Borgwardt. Graph kernels. Journal of Machine Learning Research, 1 121212) 2010.
Andrew Gordon Wilson Zhiting Hu, Ruslan Salakhutdinov, and Eric P Xing. Deep kernel learning. In Artificial intelligence and statistics, pp. 370-378, 2016.
Hongteng Xu. Grosov-wasserstein factorization models for graph clustering. arXiv preprint 1.Xiv:9111.08531 2019.
Keyulu Xu, Weihua Hu, Jure Leskovec, and Stefanie Jegelka. How powerful are graph neural networks? International Conference on Learning Representations (ICLR), 2019.
Kevin Yang, Kyle Swanson, Wengong Jin, Connor Coley, Philipp Eiden, Hua Gao, Angel Guzman- Perez, Timothy Hopper, Brian Kelley, Miriam Mathea, et al. Analyzing learned molecular representations for property prediction. Journal of chemical information and modeling. 59(8): 3370-3388, 2019.
Rex Ying, Ruining He, Kaifeng Chen, Pong Eksombatchai, William L. Hamilton, and Jure Leskovec. Graph convolutional neural networks for web-scale recommender systems. In Proceedings of the 24th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, pp. 974-983, 2018a.
Zhitao Ying, Jiaxuan You, Christopher Morris, Xiang Ren, Will Hamilton, and Jure Leskovec. Hierarchical graph representation learning with differentiable pooling. In Advances neural information processing systems, pp. 4800-4810, 2018b.
Manzil Zaheer, Satwik Kottur, Siamak Ravanbakhsh, Baraabas Poczos, Russ R Salakhutdinov, and Alexander Smola. Deep sets. In Advances in neural information processing systems, pp- 3391-3401, 2017.
12

---

Under review as a conference paper at ICLR 2021
A FURTHER DETAILS ON CONTRASTIVE REGULARIZATION
A.1 MOTIVATION
One may speculate that it was locally easier for the model to extract valuable information if it would behave like the Euclidean component, preventing it from exploring other roads of the optimization landscape. To better understand this situation, consider the scenario which a subset of points in prototype point cloud "collapses", i.e. become close to each other (see Figure [3). thus sharing similar distances to all the node embeddings of real input graphs. The submatrix of the optimal transport matrix corresponding to these collapsed points can be equally replaced by any other submatrix with the same marginals (i.e. same two vectors obtained by summing rows or columns), meaning that the columns
optimal transport matrix is not discriminative. In general, we want avoid any two in the Wasserstein cost matrix being proportional. An additional problem of point collapsing is that it a non-escaping situation when using gradient-based learning methods. The reason is that gradients of these collapsed points would become and remain identical, thus nothing will encourage them to "separate" in the future.
A.2 ON THE CHOICE OF CONTRASTIVE SAMPLES
Our experiments were conducted with ten negative samples for each correct transport plan. Five of them were obtained by initializing a matrix with uniform i.i.d entries from (0, 10) and performing around five Sinkhorn iterations (Cuturi 2013) in order make the matrix satisfy the marginal constraints. The other five were obtained by randomly permuting the columns of the correct transport plan. The latter choice has the desirable effect of penalizing the points of prototype point cloud Q; to collapse onto the same point. Indeed, the rows of T E CHQ. index points in H, while its columns index points in Q,
B THEORETICAL RESULTS
B.1 PROOF OF THEOREM I
1. Let first justify why agg is not universal. Consider function f E C(X) such that there exists
X, Y X' satisfying both f(X) * f(Y) and Σx Eiy. Clearly, any function of the form
Σ. (agg(( + 0.) would take equal values on X and Y and hence would not approximate arbitrarily well.
2. To justify that W is universal, we take inspiration from the proof of universality of neural networks Cybenko 1989).
Notation. Denote by (()) the space of finite, signed regular Borel measures on X.
Definition. We say that a is discriminatory w.r.t a kernel l if for a measure μ E M
J a (k + (X) = 0
for all Y E X7 and E R implies that y =
We start by reminding lemma coming from the original paper on the universality of neural networks by Cybenko Cybenko (1989).
Lemma. If a is discriminatory w.r.t. then is universal.
Proof: Let 5 be the subset of functions of the form Dis (go Q:) + 0.) for any 0, E R. Qi E XT and m E N and denote by S the closure of S in C(2)) Assume by contradiction that
Wtrtt the topology defined by the sup norm ||f supxex [f(x).
13

---

Under review as a conference paper at ICLR 2021
+ C(X). By the that for all S. L(h) = 0 and such that there exists E (()) s.t. L(h') * By the Riesz representation theorem, this bounded linear functional is of the form:
theorem, there exists a bounded linear functional L on C(Y) such
L(h) = - XxE h(X)du(X),
for all h E C()), for some ja E M(x) Since or (Q. 0) S, we have
J (M( 0)dp =
for all Q X3 and 0 R. Since is discriminatory w.r.t. this implies that p and hence I 0, which is contradiction with L(M) + O. Hence $ = C(X), i.e. S is dense in C(2) and is universal.
Now let us look at the part of the proof that is new.
Lemma. σ is discriminatory w.r.t. WL2-
Proof: Note that for any X, Y, A so when À - +oc we have that o (A(WL2 (X, Y) + + 9) goes
to if (X2(X, Y) 0, to 0 if WL X, Y) + 0 < 0 and to a (5p) if WLa X, Y)+0= +
Denote by Ty (X X ( Y) 0 0) and BY, 0 for f 2 0 and for < 0. By the Lebesgue Bounded Convergence Theorem we have:
(X X v M Y) <
0 f lim aAMMLa(X, Y) -) ()d
xEX XEX A++00 (()))((Iv..)) u(x' Bx,v))
Since this is true for any io, it implies that u(lly 0) = u(X'' Bx, 0. From x(X)= 0 (because By,vō 0 for 0 0), also have ya( B (x, = Hence l
all balls defined by the
metric VW2--
From the Hahn decomposition theorem, there exist disjoint Borel sets P, N such that X' PU N and == - μ where (A) == p(An (A) == μ(AA n N) for any Borel set A with yt, being positive measures. Since μ² and μ coincide on all balls on finite dimensional metric space, they coincide everywhere -orrgennn-- 1976) and hence μ = 0.
Combining the previous lemmas with W1.2 concludes the proof.
B.2 PROOF THEOREM 2
1. We build counter example. We consider 4 point clouds of size n 2 and dimension d 2. First, define u; (Li/2), (2 :%2) for i E {0, (0.....). 3). hand, if W(X,, X,)
Then take X {uo, un X2 = {uo, } 0, then all vectors in the
X3 {un, us} } and X, = {u, ). On the two point clouds are orthogonal. which can only happen for {i.j} = {1,2). On the other hand, W(X,X,) = then either = or = = 4. This yields the following Gram matrix
1
0
1
1
0
1
1
2
1
2
whose determinant is -1/16, which implies that this matrix has negative eigenvalue.
14

---

Under review as a conference paper at ICLR 2021
2. This comes from proposition 1 Choosing WL2 and 2o to be the trivial point cloud made of times the zero vector yields k = Wdot- Since k is not positive definite from the previous point of the theorem, k is not conditionally negative definite from proposition
B.3 SHAPE OF THE OPTIMAL TRANSPORT PLAN FOR POINT CLOUDS OF SAME SIZE
The below result describes the shape of optimal transport plans for point clouds of same size. For the sake of curiosity, we also illustrate in Figure 2 the optimal transport for point clouds of different sizes. We note that non-square transports seem to remain relatively sparse as well. This is in line with our empirical observations.
Proposition 2. For X, Y E Xn,d there exists a rescaled permutation matrix ,, (oio()) Jgassn which an optimal transport plan, iee.
1
(X2(X, Y)
[x01)) yy
Waon (X, Y) (xo)):
(10)
Proof. It is well known from Birkhoff's theorem that every squared boubly-stochastic matrix is a convex combination of permutation matrices. Since the Wasserstein cost for a given transport T is linear function, it is also a convex/concave function, and hence it is maximized/minimized over the convex compact set of couplings at one of its extremal points, namely one of the rescaled permutations, yielding the desired result.
C COMPLEXITY
C.I WASSERSTEIN
Computing the Wasserstein optimal transport plan between two point clouds consists in the mini- mization of a linear function under linear constraints. can either be performed exactly by using network simplex methods or interior point methods done by (Pele & Werman 2009) in time Ö(n²), or approximately up via the Sinkhorn algorithm (Cuturi 2013 in time (n²/ /e ( More recently, (Dvurechensky et al., 2018) proposed an algorithm solving OT to e with time complexity ( (n,"=="n""/)") via a primal-dual method inspired from accelerated gradient descent.
In our experiments, we used the Python Optimal Transport (POT) library (Flamary & Courty 2017). We noticed empirically that the Earth Mover Distance (EMD) solver yielded faster and more accurate solutions than Sinkhorn for datasets, because the graphs and point clouds were small enough (< 30 elements). However, Sinkhorn may take the lead for larger graphs.
Significant speed up could potentially be obtained by rewritting the POT library for it to solve OT in batches over GPUs. In our experiments, we ran all jobs on CPUs.
15

---

Under review as a conference paper at ICLR 2021
D FURTHER EXPERIMENTAL DETAILS
D.1 SETUP OF EXPERIMENTS
Each dataset is split randomly 5 times into 80%:10%::00% train, validation and test sets. For each of the 5 splits, we run each model 5 times to reduce the variance in particular data splits (resulting in each model being run 25 times). We search hyperparameters for each split of the data, and then take the average performance over all the splits. The hyperparameters are separately searched for each data split, so that the model performance based on completely unseen set, and that there is no data leakage across data splits. The models are trained for 150 epochs with early stopping if the validation error has not improved in 50 epochs and a batch size of 16. We train the models using the Adam optimizer with a learning rate of 5e-4. For the prototype models, we use different learning rates for the GNN and the point clouds (5e-4 and 5e-3 respectively), because empirically we find that the gradients are much smaller for the point clouds. The molecular datasets used for experiments here are small size (varying from 1 1-4k data points), so this is a fair method of comparison, and is indeed what is done in other works molecular property prediction Yang et al. (2019).
Parameter Name n_epochs batch size (5c-4)
Search Values
Description
{150}
Number of epochs trained
{16}
Size of each batch
Ir
Overall learning rate for model
___pe
{5e-3}
Learning rate for the prototype embeddings Number of layers the GNN Size of hidden dimension in GNN Size of the output feed forward layer Dropout probability for GNN
layers
{5}
n_ hidden
(50, 200}
n_
(1e2, 1e3, le4]
dropout_gnn
(0.)
dropout fnn
0.1
Dropout probability for feed forward layer
n_pe (M)
(10
Number of the prototypes (point clouds)
pe_size (N)
{10}
Number of points in each prototype (point cloud)
pe_ hidden (d)
{5, 10}
Size of hidden dimension of each point in each prototype
nc_coef
{o., 0.010,...111
Coefficient for noise contrastive regularization
Table 5: The parameters for our models (the prototype models all use the same GNN base model), and the values that we used for hyperparameter search. When there only single value in the search list, it means we did not search over this value, and used the specified value for all models.
D.2 BASELINE MODELS
Both the GIN Xu et 2019) and GAT (Veličković al. 2017) models were originally used for graphs without edge features. Therefore, we adapt both these algorithms to our use-case, in which edge features are a critical aspect of the prediction task. Here, we expand on the exact architectures that we use for both models.
First we introduce common notation that we will use for both models. Each example is defined by a set of vertices and edges (V, F). Let Uj E V denote the ith node in the graph, and let eij E E denote (k) the edge between nodes (i. j). Let fas be the feature representation of node ii at layer k. Let her, be the input features for the edge between nodes (i. j). and is static because we do updates only on nodes. Let N (v:) denote the set of neighbors for node i, not including node i; let / (v:) denote the set of neighbors for node i as well as the node itself.
GIN
The update rule for GIN is then defined as:
h = MML(()
(()))
+
Σ
[k(-1) + w(()) her!))
(1I)
mEN((m)
16

---

Under review as a conference paper at ICLR 2021
As with the original model, the final embedding hig is defined as the concatenation of the summed node embeddings for each layer.
h CONCAT
Σ l 0, 0,
(12)
GAT
For our implementation of the GAT model, compute the attention scores for each pairwise node
ij as follows.
(k)
heml))) +er E Eo,ed(o.) exp (LeakyReLU(a(K)) (w()) Aa (x-1) (| (uns ( 1) + with hessssss))
exp (LeakyReLU(al)) [w/Kk) (W) (k= (k-1)) ( (n!s-1)
(R) = 9
(13)
{W(*), W(A), w/ final attention score for each pair of nodes. Note that here we do attention across all of a node's neighbors well as the node itself.
are layer specific feature transforms, while (*) is vector that computes the
The updated node embeddings are as follows:
h
Σ
an ho
(14)
(E.)(
The final graph embedding is just simple aggregation of the node representations the last layer (hg = E,, N As with (Veličković et al. 2017). we also extend this formulation to utilize multiple attention heads.
E WHAT TYPES OF MOLECULES DO PROTOTYPES CAPTURE ?
To better understand if the learned prototypes offer interpretability, we examined the ProtoW-Dot model trained with NC regularization (weight 0.1). For each of the 10 learned prototypes, we computed the set of molecules in the test set that are closer in terms of the corresponding Wasserstein distance to this prototype than to any other prototype. While we noticed that one prototype is closest to the majority of molecules, there are other prototypes that are more interpretable as shown in fig,
F RELATED WORK ON PROTOTYPE LEARNING
Learning prototypes to solve machine learning tasks has been extensively studied, Generalized learning vector quantization (GLVQ) methods (Kohonen 1995 Sato & Yamada, 1995) perform classification by assigning to each data point the class of the closest neighbor prototype measured using some distance function, typically Euclidean. Each class has set of prototypes that are optimized together such that the closest wrong prototype is moved away, while the correct prototype brought closer. Several extensions of GLVQ (Hammer & Villmann, 2002, Schneider al. 2009) Bunte et al., 2012) introduce feature weights and parameterized input transtormations to leverage more flexible and adaptive metric spaces. Nevertheless, such models are limited to classification tasks and might suffer from extreme gradient sparsity.
On the other hand, more similar to our work are the radial basis function (RBF) networks (Chen et al. 1991) that perform classification regression based RBF kernel similarities to all prototypes. One such similarity vector is computed for each data point and used together with a shared linear output layer to obtain the final predictions. Prototypes typically set an unsupervised fashion, e.g. via k-means clustering, or using the Orthogonal Least Square Learning algorithm, unlike being learned using backpropagation as in our case.
Combining the non non-noonnann parametric mathematical power of kernel methods with the flexibility of deep learning models have resulted in even more expressive and scalable similarity functions that have
17

---

Under review as a conference paper at ICLR 2021
Prototype
Prototype
Prototype
Figure 6: The closest molecules to some particular prototypes terms of the corresponding Wasser- stein distance. One observe that some prototypes closer insoluble molecules containing rings (Prototype 2), while others prefer more soluble molecules (Prototype 1).
been conveniently trained with backpropagation by maximizing the likelihood of a Gaussian process (Wilson et al. 2016). Recently, learning parametric data embeddings and prototypes was also investigated for few-shot and zero-shot classification scenarios (Snell al. (2017).
Finally, using distances to prototypes as opposed to p.d. kernels, while not as common, was analyzed in the past (Duin & Pekalska 2012: Snell et al. 2017).
In contrast with the above line of work, our research focuses on learning parametric prototypes for graphs trained jointly with graph embedding functions for both graph classification and regression problems. Prototypes are modeled as sets (point clouds) of embeddings, while graphs are represented by sets of unaggregated node embeddings obtained using graph neural network models. Disimilarities between prototypes and graph embeddings are then quantified via set distances computed using optimal transport. Additional challenges arise due to the combinatorial nature of the Wasserstein distances between sets, hence our discussion on introducing the noise contrastive regularizer.
18


Reviews:
[{"review_id": "bEVChHmSYj", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "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.\n\nPros:\n1. The empirical study is thorough and the model has good empirical performance. \n2. 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. \n\nAreas to improve:\n1. 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?\n2. 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)?\n3. 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? \n4. 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?", "extra_scores": {}, "extra_text": {}}, {"review_id": "WJjMZ-dDXr", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "**Summary**\nThe 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.\n\n**Pros**\n1. 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.\n2. 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.\n3. The computational complexity is analyzed for Wasserstein optimal transport in the OT-GNN. \n\n**Cons**\n1. 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?\n2. 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?\n3. 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.\n4. The authors also need to try more public graph datasets, such as Open Graph Benchmark, https://ogb.stanford.edu/.\n5. Line 4 below (6): the last sentence seems not finished. ", "extra_scores": {}, "extra_text": {}}, {"review_id": "DFrMf-okCg", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "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.\n\n\nBelow I would like to list my main concerns or confusions regarding the paper.\n\nThe 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.\nThe 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. \nThe 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?\nWhat is the partition function you mentioned in the 3rd line after Eq.6? Do you use different terminology from in Nickel & Kiela (2017)? \nThe 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. \nFollowing the last point, I would expect more recent research to be included in the Related Work section.\nSome figures could be redesigned for better presentation. For example, it’s hard to read the axis information in Figure 4.\n\nMinor typos:\n\n(line 2, Section 2.1) “…property prediction by Yang et al. Yang et al. (2019).”\n(line 1, Section 2.2) “Optimal Transport (OT) Peyré et al. (2019) is a mathematical framework…”\n\nGiven the problems stated as above, I would suggest the paper to be carefully polished before it’s ready to publish.", "extra_scores": {}, "extra_text": {}}, {"review_id": "6AQr6RMJncj", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "## Summary\n\nThe 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.\n\n## Strengths\n\nThe paper has several strengths worth highlighting:\n\n* 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.\n* 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.\n* 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)\n* 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.\n\n## Weaknesses\n\nThat being said, I also see some issues which could be addressed to make the paper even stronger.\n\n* Most importantly, from my perspective, the relation to existing work warrants much more discussion in two directions:\n    * 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.\n    * 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]\n* I see several issues with how prototypes are integrated into the model.\n    * 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.\n    * 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.\n    * 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.\n    * 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?\n* 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).\n* 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.\n* 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.\n\nThere are also a few smaller issues which could be clarified:\n\n* 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?\n* Page 3: The vectors r_i are not introduced. Are these supposed to be the weight vectors of neurons in the MLP?\n* 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.\n\n## Judgment\n\nOverall, 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.\n\n## References\n\n* Sato, A., & Yamada, K. (1995). Generalized Learning Vector Quantization. Proceedings of NIPS 1995. [Link][GLVQ]\n* de Vries, H., Memisevic, R., & Courville, A. (2016). Deep Learning Vector Quantization. Proceedings of ESANN 2016. [Link][DLVQ]\n* 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]\n* 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]\n* Wilson, A., Hu, Z., Salakhutdinov, R., and Xing, E. (2016). Deep Kernel Learning. Proceedings of AISTATS 2016. [Link][DeepKernels]\n\n[GLVQ]:https://papers.nips.cc/paper/1113-generalized-learning-vector-quantization \"Sato, A., & Yamada, K. (1995). Generalized Learning Vector Quantization. Proceedings of NIPS 1995.\"\n[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.\"\n[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).\"\n[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).\"\n[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\"\n\n## Clarifications\n\nWhile 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:\n\n* 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?\n* 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?\n* The optimal transport plan T is treated as constant for backprop, but is recomputed after each gradient step, correct?\n* Are the limitations to Theorem 1 described above correct or did I misunderstand the Theorem?\n\n## Typos\n\nBeyond 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.\n\n* Abstract: The abbreviation 'OT' is not introduced\n* 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\n* page 2: 'that is universal approximator.' -> that is a universal approximator\n* page 2: The abbreviation 'D-MPNN' is not introduced\n* 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?\n* As far as I know, ICLR permits an appendix. As such, it would be good to move the appendix into the main paper.", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

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