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