Optimal Transport Graph Neural Networks

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 ( 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 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 . 1 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 et al. , 2019 ; Vamathevan et al. , 2019 ) . Historically , graphs were decomposed into features such as molecular fingerprints , 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 . Specifically , 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 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 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 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 a set of abstract prototypes ( fig . 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 first 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 notion of prototypes is 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 et al. , 2017 ) . Our model improves upon traditional aggregation by explicitly tapping into the full set of node embeddings without collapsing them first to a 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 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 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 to standard summation , thus fully exploiting the OT geometry . Third , we provide a theoretical justification 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 . 2 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-the-art molecular property prediction by Yang et al . ( 2019 ) . This model takes as input a directed graph G = ( V , E ) , with node and edge features denoted by xv and evw respectively , for v , w in the vertex set V and v → w in the edge set E. The parameters of D-MPNN are the matrices { Wi , Wm , Wo } . It keeps track of messages mtvw and hidden states htvw for each step t , defined as follows . An initial hidden state is set to h0vw : = ReLU ( Wicat ( xv , evw ) ) where “ cat ” denotes concatenation . Then , the updates are : mt+1vw = ∑ k∈N ( v ) \ { w } htkv , h t+1 vw = ReLU ( h 0 vw + Wmm t+1 vw ) , ( 1 ) 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 : mv = ∑ w∈N ( v ) hTvw , hv = ReLU ( Wocat ( xv , mv ) ) . ( 2 ) A final graph embedding is then obtained as h = ∑ v∈V hv , which is usually fed to a multilayer perceptron ( MLP ) for classification or regression . 2.2 OPTIMAL TRANSPORT GEOMETRY . Optimal Transport ( Peyré et al. , 2019 ) is a mathematical framework that defines 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 = { xi } ni=1 of size n be a set of n points xi ∈ Rd . 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 : T1m = 1 n1n and T T1n = 1 m1m . Intuitively , the marginal constraints mean that T preserves the mass from X to Y . We denote the set of such couplings as CXY . Given a cost function c on Rd , its associated Wasserstein discrepancy is defined as W ( X , Y ) = min T∈CXY ∑ ij Tijc ( xi , yj ) . ( 3 ) We further describe the shape of optimal transports for point clouds of same sizes in Appendix B.3 . 3 MODEL & PRACTICE . 3.1 ARCHITECTURE ENHANCEMENT . Reformulating standard architectures . As mentioned at the end of Section 2.1 , the final graph embedding h = ∑ v∈V hv obtained by aggregating node embeddings is usually fed to a MLP performing a matrix-multiplication whose i-th component is ( Rh ) i = 〈ri , h〉 , where ri is the i-th row of matrix R. Replacing 〈· , ·〉 by a distance/kernel k ( · , · ) allows the processing of more general graph representations than just vectors in Rd , such as point clouds or adjacency tensors . From a single point to a point cloud . We propose to replace the aggregated graph embedding h = ∑ v∈V hv by the point cloud ( of unaggregated node embeddings ) H = { hv } v∈V , and the inner-products 〈h , ri〉 by the below written Wasserstein discrepancy : W ( H , Qi ) : = min T∈CHQi ∑ vj Tvjc ( hv , q j i ) , ( 4 ) where Qi = { qji } 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 = ‖ · − · ‖22 or c = −〈· , ·〉 . 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 . Final architecture . Finally , the vector of all Wasserstein distances in eq . ( 4 ) becomes the input to a 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 . 1 and motivated theoretically in section 4.1 . 3.2 CONTRASTIVE REGULARIZATION . What would happen toW ( H , Qi ) if all points qji belonging to point cloud Qi would collapse to the same point qi ? All transport plans would yield the same cost , giving for c = −〈· , ·〉 : W ( H , Qi ) = − ∑ vj Tvj〈hv , qji 〉 = −〈h , qi/|V |〉 . ( 5 ) In this scenario , our proposition would simply over-parametrize the standard Euclidean model . A first obstacle and its cause . Our first 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∑ vj Tvj〈hv , q j i 〉 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 difficulty , we consider adding a 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 a point cloud Y of node embeddings and let Ti be an optimal transport plan obtained in the computation of W ( Y , Qi ) . For each Ti , we then build a set Neg ( Ti ) ⊂ CYQi of noisy/contrastive transports . If we denote byWT ( X , Y ) : = ∑ kl Tklc ( xk , yl ) the Wasserstein cost obtained for the particular transport T , then our contrastive regularization consists in maximizing the term : ∑ i log ( e−WTi ( Y , Qi ) e−WTi ( Y , Qi ) + ∑ T∈Neg ( Ti ) e −WT ( Y , Qi ) ) , ( 6 ) which can be interpreted as the log-likelihood that the correct transport Ti be ( as it should ) a better minimizer ofWT ( Y , Qi ) than its negative samples . This can be considered as an approximation of log ( Pr ( Ti | Y , Qi ) ) , 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 reflect the trade-off : ( i ) to not be too large , for computational efficiency while ( ii ) 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 ( Ti ) with a singleton { T } for a contrastive random variable T lets us rewrite Eq . ( 6 ) as1∑ i log σ ( WT −WTi ) , reminiscent of noise contrastive estimation ( Gutmann & Hyvärinen , 2010 ) .

Summary: 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.