O
PTIMAL T RANSPORT G RAPH N EURAL N ETWORKS
Anonymous authors
Paper under double-blind review
A BSTRACT
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 ﬁrst to successfully combine optimal transport (OT) 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 satisﬁes 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.
I NTRODUCTION
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 et al., 2019; Vamathevan et al., 2019).
Historically, graphs were decomposed into features such as molecular ﬁngerprints, or turned into non-parametric graph kernels (Vishwanathan 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 et al., 2019).
Despite these successes, graph neural networks are often underutilized in whole graph prediction tasks such as molecule property prediction.
Speciﬁcally, while GNNs produce node embeddings for each atom in the molecule, these are typically aggregated via simple operations such as a sum or average, turning the molecule into a single vector prior to classiﬁcation 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 a kernel function over graphs by directly comparing non-parametric node embeddings as point clouds through optimal transport (Wasserstein distance).
Their non-parametric model yields better empirical performance over popular graph kernels, but this idea hasn’t been extended to the more challenging parametric case where optimization difﬁculties have to be reconciled with the combinatorial aspects of optimal transport solvers.
Motivated by these observations and drawing inspiration from prototypical networks (Snell et al., 2017), we introduce a new class of GNNs where the key representational step consists of comparing each input graph to a set of abstract prototypes (ﬁg. 1).
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 ﬁrst encoded into a set of 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.
The prototypes as abstract basis functions can be understood as keys that highlight property values 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.
Our model improves upon traditional aggregation by explicitly tapping into the full set of node embeddings without collapsing them ﬁrst to a single vector.
We theoretically prove that, unlike standard GNN aggregation, our model deﬁnes 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 primary signal is initially available only in aggregate form.
If trained as is, the prototypes often collapse to single points, reducing the Wasserstein distance between point clouds to Euclidean comparisons of their means.
To counter this effect, we introduce a contrastive regularizer which effectively prevents the model from collapsing (Section 3.2).
We demonstrate its merits empirically.
Our contributions.
First, we introduce an efﬁciently 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 ﬁrst to successfully train parametric graph models together with combinatorial OT distances, despite optimization difﬁculties.
A key element is our noise contrastive regularizer that prevents the model from collapsing back to standard summation, thus fully exploiting the OT geometry.
Third, we provide a theoretical justiﬁcation of the increased representational power compared to the standard GNN aggregation method.
Finally, our model shows consistent empirical improvements over previous state-of-the-art on molecular datasets, yielding also smoother graph embedding spaces.
P RELIMINARIES 2.
[Equation]
We brieﬂy remind here of the simpliﬁed D-MPNN (Dai et al., 2016) architecture which was successfully adapted for state-of-the-art molecular property prediction by Yang et al.
This model takes as input a directed graph G =
( V, E ) , with node and edge features denoted by x v and e vw respectively, for v , w in the vertex set V and v → w in the edge set E .
The parameters of D-MPNN are the weight matrices { W i , W m , W o } .
It keeps track of messages m tvw and hidden states h tvw for each step t , deﬁned as follows.
An initial hidden state is set to h 0 vw :=
ReLU ( W i cat( x v , e vw )) where “ cat ” denotes concatenation.
Then, the message passing operates as
[Equation]
where N ( v ) =
{ k ∈ V | ( k, v ) ∈ E } denotes v ’s incoming neighbors.
After T steps of message passing, node embeddings are obtained by summing edge embeddings:
[Equation]
A ﬁnal graph embedding is then obtained as h = (cid:80) v ∈ V h v , which is usually fed to a multilayer perceptron (MLP) for classiﬁcation or regression.
2. O PTIMAL T RANSPORT G EOMETRY
Optimal Transport (Peyré et al., 2019) is a mathematical framework that deﬁnes distances or similarities between objects such as probability distributions, either discrete or continuous, as the cost of an optimal transport plan from one to the other.
Wasserstein distance for point clouds.
Let a point cloud X
= { x i } ni
=1 of size
n be a set of n points x i ∈ R d .
Given point clouds X , Y of respective sizes n, m , a transport plan (or coupling ) is a matrix T of size n × m with entries in [0 , 1] , satisfying the two following marginal constraints :
T1 m
= 1 n 1 n and T T 1 n = 1 m 1 m .
Intuitively, the marginal constraints mean that T preserves the mass from X to Y .
We denote the set of such couplings as C XY .
Given a cost function c on R d , its associated Wasserstein discrepancy is deﬁned as
[Equation]
We further describe the shape of optimal transports for point clouds of same sizes in Appendix B.3.
M ODEL & P RACTICE 3. A RCHITECTURE E NHANCEMENT
Reformulating standard architectures.
As mentioned at the end of Section 2.1, the ﬁnal graph embedding h = (cid:80) v ∈ V h v obtained by aggregating node embeddings is usually fed to a MLP performing a matrix-multiplication whose i-th component is ( Rh ) i
=
(cid:104) r
i , h (cid:105) , where r i is the i-th row of matrix R .
Replacing (cid:104)· , ·(cid:105) by a distance/kernel k ( · , · ) allows the processing of more general graph representations than just vectors in R d , such as point clouds or adjacency tensors.
From a single point to a point cloud.
We propose to replace the aggregated graph embedding h = (cid:80) v ∈ V h v by the point cloud (of unaggregated node embeddings)
H
=
{ h v } v ∈ V , and the inner-products (cid:104) h , r i (cid:105) by the below written Wasserstein discrepancy :
[Equation]
where Q i = { q ji } j ∈{ 1 ,...,N } , ∀ i ∈ { 1 , . . , M } represent M prototype point clouds each being represented as a set of N embeddings as free trainable parameters, and the cost is chosen as c = (cid:107) · − · (cid:107) 22 or c = −(cid:104)· , ·(cid:105) .
Note that both options yield identical optimal transport plans.
Greater representational power.
We formulate mathematically in Section 4 that this kernel has a strictly greater representational power than the kernel corresponding to standard inner-product on top of a sum aggregation, to distinguish between different point clouds.
However, there are additional practical concerns that will be discussed in the next subsection. Final architecture.
Finally, the vector of all Wasserstein distances in eq.
(4) becomes the input to a ﬁnal MLP with a single scalar as output.
This can then be used as the prediction for various downstream tasks.
This model is depicted in ﬁg. 1 and motivated theoretically in section 4.1.
C ONTRASTIVE R EGULARIZATION
What would happen to W ( H , Q i ) if all points q ji belonging to point cloud Q i would collapse to the same point q i ?
All transport plans would yield the same cost, giving for c = −(cid:104)· , ·(cid:105) :
[Equation]
In this scenario, our proposition would simply over-parametrize the standard Euclidean model.
[Figure]
A ﬁrst obstacle and its cause.
Our ﬁrst empirical trials with OT-enhanced GNNs showed that a model trained with only the Wasserstein component would sometimes perform similarly to the
Euclidean baseline in 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 a way that the optimal transport plan T obtained by maximizing (cid:80) vj T vj (cid:104) h v , q ji (cid:105) was not discriminative , i.e. many other transports would yield a 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.1 and Figure 3..
Contrastive regularization.
To address this difﬁculty, we consider adding a regularizer which encourages the model to displace its prototype point clouds such that the optimal transport plans it computes would be discriminative against chosen contrastive transport plans .
Namely, consider a point cloud Y of node embeddings and let T i be an optimal transport plan obtained in the computation of W ( Y , Q i ) .
For each T i , we then build a set Neg ( T i ) ⊂
C YQ i of noisy/contrastive transports.
If we denote by W T ( X , Y ) := (cid:80)
kl T kl c ( x k , y l ) the Wasserstein cost obtained for the particular transport T , then our contrastive regularization consists in maximizing the term:
[Equation]
which can be interpreted as the log-likelihood that the correct transport T i be (as it should) a better minimizer of W T ( Y , Q i ) than its negative samples.
This can be considered as an approximation of log(Pr( T i | Y , Q i )) , where the partition function is approximated by our selection of negative examples, as done e.g. by Nickel & Kiela (2017).
Its effect is shown in Figure 3.
Remarks.
The selection of negative examples should reﬂect the trade-off: (i) to not be too large, for computational efﬁciency while (ii) containing sufﬁciently meaningful and challenging contrastive samples.
Details about choice of contrastive samples are given in Appendix A.2.
Note that replacing the set Neg ( T i ) with a singleton { T } for a contrastive random variable T lets us rewrite Eq.
(6) as 1 (cid:80) i log σ ( W T − W T i ) , reminiscent of noise contrastive estimation (Gutmann & Hyvärinen, 2010).
3. O PTIMIZATION & C OMPLEXITY
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 et 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 ﬁxed during the back-propagation of gradients through Wasserstein distances.
For the sake of simplicity and training stability, we resort to the latter procedure: keeping T ﬁxed during back-propagation.
We discuss complexity in appendix C.
T HEORETICAL A NALYSIS
In this section we show that the standard architecture lacks a fundamental property of universal approximation of functions deﬁned on point clouds, and that our proposed architecture recovers this property.
We will denote by X nd the set of point clouds X
= { x i } ni
=1 of size n in R d .
4. U NIVERSALITY
As seen in Section 3.1, we have replaced the sum aggregation − followed by the Euclidean innerproduct − by Wasserstein discrepancies.
How does this affect the function class and representations?
A common framework used to analyze the geometry inherited from similarities and discrepancies is that of kernel theory.
A kernel k on a set X is a symmetric function k :
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 deﬁned on 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 k deﬁned on X nd is said to be universal if the following holds: for any compact subset X ⊂ X nd , the set of functions in the form 2 (cid:80) mj
=
α j σ ( k ( · , θ j ) + β j ) is dense in the set C ( X ) of continuous functions from X to R , w.r.t the sup norm (cid:107) · (cid:107) ∞ , X , σ denoting the sigmoid.
Although the notion of universality does not indicate how easy it is in practice to learn the correct function, it at least guarantees the absence of a fundamental bottleneck of the model using this kernel. In the following we compare the aggregating kernel agg ( X , Y ) : = (cid:104) (cid:80) i x i , (cid:80) j y j (cid:105) (used by popular GNN models) with the Wasserstein kernel, where
[Equation] Theorem 1.
We have that:
The aggregation kernel agg is not universal .
The Wasserstein kernel W L2 is universal .
Proof: See Appendix B.1.
Universality of the Wasserstein kernel W L2 essentially comes from the fact that its square-root deﬁnes a metric, and in particular from the axiom of separation of distances:
if d ( x, y ) = 
 then x = y .
Implications.
Theorem 1 states that our proposed OT-GNN model is strictly more powerful than the state of the art GNN models that use summation or averaging of node embeddings. Nevertheless, this implies we can only distinguish graphs that have distinct multi-sets of node embeddings, e.g. all Weisfeiler-Lehman distinguishable graphs in the case of graph convolutional networks. 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 classiﬁcation – e.g. using a multilayer perceptron (MLP) on top.
4. D EFINITENESS
For the sake of simpliﬁed mathematical analysis, similarity kernels are often required to be positive deﬁnite (p.d.), which corresponds to discrepancy kernels being conditionally negative deﬁnite (c.n.d.).
Although such a property has the beneﬁt of yielding the mathematical framework of Reproducing
Kernel Hilbert Spaces, it essentially implies linearity , i.e. the possibility to embed the geometry deﬁned by that kernel in 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.
Kernel deﬁniteness.
A kernel k is positive deﬁnite (p.d.) on X if for n ∈ N ∗ , x 1 , ..., x n ∈ X
and c 1 , ..., c n ∈ R , we have (cid:80) ij c
i c j k ( x i , x j ) ≥ 0 .
It is conditionally negative deﬁnite (c.n.d.)
on X if for n ∈ N ∗ , x 1 , ..., x n ∈ X and c 1 , ..., c n ∈ R such that (cid:80) i c i = 0 , we have (cid:80) ij c
i c j k ( x i , x j ) ≤ 0 .
These two notions relate to each other via the below result Boughorbel et al.
Proposition 1.
Let k be a symmetric kernel on X , let x 0 ∈ X and deﬁne the kernel:
˜ [Equation]
Then ˜ k is p.d. if and only if k is c.n.d.
Example: k = (cid:107) · − · (cid:107) 22 and x 0 = 0 yield ˜ k
= (cid:104)· , ·(cid:105) .
One can easily show that agg also deﬁnes a p.d. kernel, and that agg ( · , · ) ≤ n 2 W ( · , · ) .
However, the
Wasserstein kernel is not p.d., as stated in different variants before (e.g. Vert (2008)) and as reminded by the below theorem.
We here give a novel proof in Appendix B.2.
Theorem 2.
We have that:
The (similarity) Wasserstein kernel W dot is not positive deﬁnite ; 2.
The (discrepancy) Wasserstein kernel W L2 is not conditionally negative deﬁnite .
E XPERIMENTS 5. E XPERIMENTAL S ETUP
We experiment on 4 benchmark molecular property prediction datasets (Yang et al., 2019) including both regression (ESOL, Lipophilicity) and classiﬁcation (BACE, BBBP) tasks.
These datasets cover different complex chemical properties (e.g. ESOL - water solubility, LIPO - octanol/water distribution coefﬁcient, BACE - inhibition of human β -secretase, BBBP - blood-brain barrier penetration).
We provide results for our implementation of 4 different widely used graph-based models.
Fingerprint + MLP is a simple graph baseline that applies a MLP over the input features which are hashed graph structures (called a molecular ﬁngerprint).
GIN is the Graph Isomorphism Network from (Xu et al., 2019), which is a variant of a GNN.
The original GIN does not account for edge features, so we adapt their algorithm to our setting.
Next, GAT is the Graph Attention Network from (Veliˇckovi´c et al., 2017), which uses multi-head attention layers to 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 state-of-the-art performance for molecular representation learning across multiple classiﬁcation and regression datasets.
Empirically we ﬁnd that this baseline is indeed the best performing, and so is used as to obtain node embeddings in all our prototype models.
Its architecture is described in section 2.1.
Different variants of our OT-GNN prototype model are described below:
ProtoW - L2 /
Dot is the model that treats point clouds as point sets, and computes the Wasserstein distances to each point cloud (using either L2 distance or (minus) dot product cost functions) as the molecular embedding.
ProtoS - L2 is a special case of ProtoW - L2 , in which the point clouds have a 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, as that model is mathematically equivalent to 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 ﬁnd empirically to be faster than the
Sinkhorn algorithm for our datasets.
We deﬁne the cost matrix by taking the pairwise L2 or negative dot product distances.
As mentioned in Section 3.3, we ﬁx the transport plan, and only backpropagate through the cost matrix for computational efﬁciency.
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 deﬁned in
Section 3.2.
More details about experimental setup are presented in Appendix D.1.
[Table]
u r s
[Table]
Lower RMSE is better, while higher contrastive as in
5. E XPERIMENTAL R ESULTS 5.2. R EGRESSION AND C LASSIFICATION
Results are shown in Table 1.
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 ( ProtoW - 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 to traditional aggregation methods. [Figure] [Table] on the the GNN and
- L2 model the in
Spearman
[Table] Pearson r D-MPNN .424 ± . .393 ± .
ProtoS - L .561 ± .
.414 ± .
ProtoW - Dot .592 ± .
.559 ± .
ProtoW - L .815 ± .
.828 ± .
[Figure]
t. their based only on the test molecules from each dataset.
Comparing (a) vs (b) and (c) vs (d), we can observe a smoother space of our model compared to the D-MPNN baseline (see also main text)
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.
3, illustrating that without contrastive regularization, prototype point clouds are often displaced close to their mean, while regularization forces them to nicely scatter.
Quantitative results in Table 1 also highlight the beneﬁt of this regularization.
5.2. L EARNED E MBEDDING S PACE : Q UALITATIVE AND Q UANTITATIVE R ESULTS
We further examine the learned embedding space of the best baseline (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 coefﬁcient ( ρ ) and Pearson correlation coefﬁcient ( r ).
This is reminiscent of evaluation tasks for the correlation of word embedding similarity with human labels (Luong et al., 2013).
Our ProtoW - L2 achieves better ρ and r scores compared to the D-MPNN model (Table 2), that indicating our Wasserstein model constructs more meaningful embeddings with respect to 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, as qualitatively shown in Figure 5, our model provides more robust molecular embeddings compared to the baseline, in the following sense: we observe that a small perturbation of a molecular embedding corresponds to a small change in predicted property value – a desirable phenomenon that holds rarely for the baseline D-MPNN model.
Our Proto-W-L2 model yields smoother heatmaps. [Table] [Table]
5.2. A DDITIONAL E XPERIMENTAL D ETAILS
Model Sizes Using size of MPNN hidden dimension as 200, and the ﬁnal output MLP hidden dimension as 100, the number of parameters for the models are given by table 4. The ﬁngerprint 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 we compute distances to a few small prototypes (small number of overall parameters used for these point clouds).
The dimensions of the prototype embeddings are also smaller compared to the node embedding dimensions of the D-MPNN and other baselines.
We did not see signiﬁcant improvements in quality by increasing any of the hyperparameter values.
Remarkably, our model outperforms all the baselines using between 10 to 1.5 times less parameters.
Runtimes.
The runtimes for the D-MPNN and our prototype models are given in table 3.
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.
One can notice that our method is between 1 to 7.1 times slower than the D-MPNN baseline which mostly happens due to the frequent calls to the Earth Mover
Distance (EMD) OT solver.  R ELATED W ORK Graph Neural Networks were introduced by Gori et al.
and Scarselli et al.
(2008) as a form of recurrent neural networks.
Graph convolutional networks (GCN) appeared later on in various forms.
Duvenaud et al.
; Atwood & Towsley (2016) proposed a propagation rule inspired from convolution and diffusion, but these methods do not scale to graphs with either large degree distribution or node cardinality.
Niepert et al.
(2016) deﬁned a GCN as a 1D convolution on a chosen node ordering.
Kearnes et al.
(2016) also used graph convolutions to generate high quality molecular ﬁngerprints.
Efﬁcient spectral methods were proposed by Bruna et al.
Defferrard et al.
Kipf & Welling (2017) simpliﬁed their propagation rule, motivated from spectral graph theory (Hammond et al., 2011).
Different such architectures were later uniﬁed 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 to improve state-of-the-art in molecular property prediction on a wide variety of datasets by (Yang et al., 2019).
Other applications include recommender systems (Ying et al., 2018a).
Inspired by DeepSets Zaheer et al.
(2017), Xu et al.
(2019) suggest a simpliﬁed architecture called GIN, which theoretically can discriminate between any different local neighborhoods.
Other recent approaches modify the sum-aggregation of node embeddings in the GCN architecture with the aim to preserve more information Kondor et al.
Pei et al.
In this category there is also the recently growing class of hierarchical graph pooling methods which typically either use deterministic and non-differentiable node clustering (Defferrard et al., 2016; Jin et al., 2018), or differentiable pooling (Ying et 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.
Other recent geometry-inspired GNNs include adaptations to non-Euclidean spaces (Liu et al., 2019; Chami et al., 2019; Bachmann et al., 2019).
C ONCLUSION
We propose OT - GNN : one of the ﬁrst 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, we outperform popular baselines in different molecular 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. A F URTHER D ETAILS ON C ONTRASTIVE R EGULARIZATION A. M OTIVATION
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 in which a subset of points in a 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 optimal transport matrix is not discriminative.
In general, we want to avoid any two rows or columns in the Wasserstein cost matrix being proportional.
An additional problem of point collapsing is that it is 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. O N THE C HOICE OF C ONTRASTIVE S AMPLES
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 ﬁve Sinkhorn iterations (Cuturi, 2013) in order to make the matrix satisfy the marginal constraints.
The other ﬁve were obtained by randomly permuting the columns of the correct transport plan.
The latter choice has the desirable effect of penalizing the points of a prototype point cloud Q i to collapse onto the same point.
Indeed, the rows of T i ∈ C HQ i index points in H , while its columns index points in Q i .
B T HEORETICAL R ESULTS B. P ROOF OF T HEOREM
Let us ﬁrst justify why agg is not universal.
Consider a function f ∈ C ( X ) such that there exists
X , Y ∈ X satisfying both f ( X ) (cid:54) =
f ( Y )
and (cid:80) k x k = (cid:80)
l y l .
Clearly, any function of the form (cid:80)
i
α i σ ( agg ( W i , · ) + θ i ) would take equal values on X and Y and hence would not approximate f arbitrarily well.
To justify that W is universal, we take inspiration from the proof of universality of neural networks Cybenko (1989).
Notation.
Denote by M ( X ) the space of ﬁnite, signed regular Borel measures on X .
Deﬁnition.
We say that σ is discriminatory w.r.t a kernel k if for a measure µ ∈ M ( X ) ,
[Equation]
for all Y ∈ X nd and θ ∈ R implies that µ ≡ 0 .
We start by reminding a lemma coming from the original paper on the universality of neural networks by Cybenko Cybenko (1989).
Lemma.
If σ is discriminatory w.r.t.
k then k is universal.
Proof: Let S be the subset of functions of the form (cid:80) mi =1 α i σ ( k ( · , Q i ) + θ i ) for any θ i ∈ R ,
Q i ∈
X nd and m ∈ N ∗ and denote by ¯ S the closure 3 of S in C ( X ) .
Assume by contradiction that
¯ S (cid:54) =
C ( X ) .
By the Hahn-Banach theorem, there exists a bounded linear functional L on C ( X ) such that for all h ∈ ¯ S , L ( h )
= 0 and such that there exists h (cid:48)
∈ C ( X ) s.t.
L ( h (cid:48) )
(cid:54) = 0 .
By the Riesz representation theorem, this bounded linear functional is of the form:
[Equation]
for all h ∈ C ( X ) , for some µ ∈ M ( X ) .
Since σ ( k ( Q , · ) + θ ) is in ¯ S , we have
[Equation]
for all Q ∈ X nd and θ ∈ R .
Since σ is discriminatory w.r.t.
k , this implies that µ = 0 and hence
L ≡ 0 , which is a contradiction with L ( h (cid:48) )
(cid:54) = 0 .
Hence ¯ S = C
( X ) , i.e. S is dense in C ( X ) and k is universal.
[Equation]
Now let us look at the part of the proof that is new.
Lemma.
σ is discriminatory w.r.t.
W L2 .
Proof: Note that for any X , Y , θ, ϕ , when λ → + ∞ we have that σ ( λ ( W L2 ( X , Y )
+ θ ) + ϕ ) goes to 1 if W L2 ( X , Y )
+ θ > 0 , to 0 if W L2 ( X , Y )
+ θ < 0 and to σ ( ϕ ) if W L2 ( X , Y )
+ θ = 0 .
Denote by Π Y ,θ := { X ∈ X | W L2 ( X , Y ) − θ = 0 } and B Y ,θ :=
{ X ∈ X | (cid:112) W L2 ( X , Y ) <
θ } for θ ≥ 0 and ∅ for θ < 0 .
By the Lebesgue Bounded Convergence Theorem we have:
[Equation]
Since this is true for any ϕ , it implies that µ (Π Y ,θ ) =
µ ( X \ B Y , √ θ ) = 0 .
From µ ( X ) = 0 (because
B Y , √ θ = ∅ for θ < 0 ), we also have µ ( B Y , √ θ ) = 0 .
Hence µ is zero on all balls deﬁned by the metric √W L2 .
From the Hahn decomposition theorem, there exist disjoint Borel sets P, N such that X
= P ∪ N and µ = µ + − µ − where µ + ( A ) :
= µ ( A ∩ P ) , µ − ( A ) := µ ( A ∩ N ) for any Borel set A with µ + , µ − being positive measures.
Since µ + and µ − coincide on all balls on a ﬁnite dimensional metric space, they coincide everywhere Hoffmann-Jørgensen (1976) and hence µ ≡ 0 .
[Equation]
Combining the previous lemmas with k = W L2 concludes the proof.
(cid:3)
B. P ROOF OF T HEOREM
We build a counter example.
We consider 4 point clouds of size n = 2 and dimension d = 2 .
First, deﬁne u i =
( (cid:98) i/ 2 (cid:99) , i %2) for i ∈ { 0 , ..., 3 } .
Then take X 1 = { u 0 , u 1 } , X 2 = { u 0 , u 2 } ,
{+X 
+}
= { u 0 , u 3 } and X 4 = { u 1 , u 2 } .
On the one hand, if W ( X i , X j ) = 0 , then all vectors in the two point clouds are orthogonal, which can only happen for { i, j } = { 1 , 2 } .
On the other hand, if
W ( X i , X j ) = 1 , then either i =
j = 3 or i =
j = 4 .
This yields the following Gram matrix
[Equation]
whose determinant is − 1 / 16 , which implies that this matrix has a negative eigenvalue.
This comes from proposition 1.
Choosing k = W L2 and x 0 = 0 to be the trivial point cloud made of n times the zero vector yields ˜ k
= W dot .
Since ˜ k is not positive deﬁnite from the previous point of the theorem, k is not conditionally negative deﬁnite from proposition 1.
(cid:3)
B. S HAPE 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 ∈ X n,d there exists a rescaled permutation matrix 1 n ( δ iσ ( j ) ) 1 ≤ i,j ≤ n which is an optimal transport plan, i.e.
[Equation]
Proof.
It is well known from Birkhoff’s theorem that every squared doubly-stochastic matrix is a convex combination of permutation matrices.
Since the Wasserstein cost for a given transport
T is a 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 C OMPLEXITY C. W ASSERSTEIN
Computing the Wasserstein optimal transport plan between two point clouds consists in the minimization of a linear function under linear constraints.
It can either be performed exactly by using network simplex methods or interior point methods as done by (Pele & Werman, 2009) in time ˜
O ( n 3 ) , or approximately up to ε via the Sinkhorn algorithm (Cuturi, 2013) in time ˜ O ( n 2 /ε 3 ) .
More recently, (Dvurechensky et al., 2018) proposed an algorithm solving OT up to ε with time complexity ˜
O (min { n 9 / 4 /ε, n 2 /ε 2 } ) 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 our datasets, because the graphs and point clouds were small enough ( < 30 elements).
However, Sinkhorn may take the lead for larger graphs.
Signiﬁcant 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.
D F URTHER E XPERIMENTAL D ETAILS D. S ETUP OF E XPERIMENTS
Each dataset is split randomly 5 times into 80%:10%:10% 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 is based on a completely unseen test 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 ﬁnd that the gradients are much smaller for the point clouds.
The molecular datasets used for experiments here are small in size (varying from 1-4k data points), so this is a fair method of comparison, and is indeed what is done in other works on molecular property prediction Yang et al.
[Table]
D. B ASELINE MODELS
Both the GIN (Xu et al., 2019) and GAT (Veliˇckovi´c et 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 deﬁned by a set of vertices and edges ( V, E ) .
Let v i ∈
V denote the i th node in the graph, and let e ij ∈ E denote the edge between nodes ( i, j ) .
Let h ( k ) v
i be the feature representation of node v i at layer k .
Let h e ij be the input features for the edge between nodes ( i, j ) , and is static because we do updates only on nodes.
Let N ( v i ) denote the set of neighbors for node i , not including node i ; let ˆ N ( v i ) denote the set of neighbors for node
i as well as the node
GIN
The update rule for GIN is then deﬁned as:
[Equation]
[Equation]
As with the original model, the ﬁnal embedding h G is deﬁned as the concatenation of the summed node embeddings for each layer.
[Equation]
GAT
For our implementation of the GAT model, we compute the attention scores for each pairwise node α ( k ) ij as follows.
[Equation]
{ W ( k ) 1 , W ( k ) 2 , W ( k ) b } are layer speciﬁc feature transforms, while a ( k ) is a vector that computes the ﬁnal attention score for each pair of nodes.
Note that here we do attention across all of a node’s neighbors as well as the node itself.
The updated node embeddings are as follows:
[Equation]
The ﬁnal graph embedding is just a simple sum aggregation of the node representations on the last layer ( h G = (cid:80) v i h
Kv i ).
As with (Veliˇckovi´c et al., 2017), we also extend this formulation to utilize multiple attention heads.
E W HAT 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 ﬁg. 6.