Template based Graph Neural Network with Optimal Transport Distances
Abstract
Current Graph Neural Networks (GNN) architectures generally rely on two important components: node features embedding through message passing, and aggregation with a specialized form of pooling.
The structural (or topological) information is implicitly taken into account in these two steps.
We propose in this work a novel point of view, which places distances to some learnable graph templates at the core of the graph representation.
This distance embedding is constructed thanks to an optimal transport distance: the Fused Gromov-Wasserstein (FGW) distance, which encodes simultaneously feature and structure dissimilarities by solving a soft graph-matching problem.
We postulate that the vector of FGW distances to a set of template graphs has a strong discriminative power, which is then fed to a non-linear classiﬁer for ﬁnal predictions.
Distance embedding can be seen as a new layer, and can leverage on existing message passing techniques to promote sensible feature representations.
Interestingly enough, in our work the optimal set of template graphs is also learnt in an end-to-end fashion by differentiating through this layer.
After describing the corresponding learning procedure, we empirically validate our claim on several synthetic and real life graph classiﬁcation datasets, where our method is competitive or surpasses kernel and GNN state-of-the-art approaches.
We complete our experiments by an ablation study and a sensitivity analysis to parameters.
Introduction
Attributed graphs are characterized by i) the relationships between the nodes of the graph (structural or topological information) and ii) some speciﬁc features or attributes endowing the nodes themselves.
Learning from those data is ubiquitous in many research areas [3], e.g. image analysis [21, 8], brain connectivity [30], biological compounds [23] or social networks [69], to name a few.
Various methodologies approach the inherent complexity of those data, such as signal processing [54],
[Table]
BBB@ 0 .
CCCA <
[Table]
C i
[Table]
h
i
[Table] [Table]
[Table] [Table]
Learnable
[Table]
/987cznYWhr6K8PVd0uDtVf2Pi3iCZ7iuUxgjTPZwIgnnuIjPuGz88X55vx0fv2lLnSs5jFmlvP7DyRa3aA=</latexit>
Bayesian and kernel methods on graphs [41, 29] or more recently Graph Neural Networks (GNN) [64] in the framework of the geometric deep learning [8, 7].
We are interested in this work in the classiﬁcation of attributed graphs at the instance level .
One existing approach consists in designing kernels that leverage topological properties of the observed graphs [5, 17, 19, 53].
For instance, the popular Weisfeiler-Lehman (WL) kernel [52] iteratively aggregates for each node the features of its k -hop neighborhood.
Alternative approaches aim at learning vectorial representations of the graphs that can encode the graph structure ( i.e. graph representation learning [10]).
In this domain, GNN lead to state-of-the-art performances with endto-end learnable embeddings [64].
At a given layer, these architectures typically learn the node embeddings via local permutation-invariant transformations aggregating its neighbour features [37, 25, 20, 67].
In order to obtain a representation of the whole graph suitable for classiﬁcation, GNNs ﬁnally operate a pooling [26, 40] of the node embeddings, either global ( e.g summation over nodes [67]), or hierarchical ( e.g by iteratively clustering nodes [71, 70, 32]).
Another line of works targets the construction of meaningful distances that integrate simultaneously the structural and feature information, and that are based on optimal transport (OT) [61, 46].
Originally designed to compare probability distributions based on a geometric notion of optimality, it allows deﬁning very general loss functions between various objects, modeled as probability distributions.
In a nutshell, it proceeds by constructing a coupling between the distributions that minimizes a speciﬁc cost .
Some approaches dealing with graphs rely on non-parametric models that ﬁrst embed the graphs into a vectorial space and then match them via OT [43, 56, 27, 35].
Recently [11] proposed the
OT-GNN model, that embeds a graph as a vector of the Wasserstein distances between the nodes’ embeddings (after GNN pre-processing) and learnt point clouds, acting as templates.
Building further from OT variants, the Gromov-Wasserstein (GW) distance [39] directly handles graphs through the symmetric matrix C that encodes the distance/similarity between every pairs of nodes ( e.g. adjacency, shortest path) and a weight vector h on the nodes encoding the nodes’ relative importance.
GW has proven to be useful for tasks such as graph matching and partitioning [66, 15] or unsupervised graph dictionary learning [65, 63, 62].
GW has been also extended to directed graphs [14] and to attributed graphs via the Fused Gromov-Wasserstein (FGW) distance [59, 58], that realizes a trade-off between an OT distance with a cost on node features and the GW distance between the similarity matrices.
Despite its recent successes on complex unsupervised tasks such as graph clustering [65, 63, 62], FGW has never been explored as part of an end-to-end model for graph classiﬁcation.
In this work, we ﬁll this gap by introducing a novel “layer” that embeds an attributed graph into a vector, whose coordinates are FGW distances to few (learned) graph templates.
While FGW can be performed directly on raw data ( i.e. the input structured graph without any pre-processing), we also consider the case where features representations are learnt from a GNN, similarly to OT-GNN [11], and thus also realizing a particular type of aggregation.
Contributions.
We introduce a new GNN layer, named TFGW for Template-based FGW and illustrated in the center of Figure 1.
From an input graph, it computes a vector of FGW distances to learnable graph templates.
This layer that can be seen as an alternative to global pooling layers and can be integrated into any neural network architecture.
We discuss its properties and the associated invariances.
We detail the optimization strategy that enables learning simultaneously GNN preprocessing layers and graph templates relevant for a downstream task in an end-to-end fashion.
We empirically demonstrate the relevance of our model in terms of performances compared to several state-of-the-art architectures.
Remarkably, we show that a simple GNN model leveraging on our new layer can surpass state-of-the-art performances by a relatively large margin.
Finally, we also provide some illustrative interpretations of our method and a sensitivity analysis of our model parameters.
Fused Gromov-Wasserstein template based layer
In order to describe our novel template-based GNN layer
we ﬁrst introduce more formally the
FGW distance and its properties.
In the following we denote by Σ n := { h ∈ R n + | (cid:80) i h i
= 1 } the probability simplex with n -bins, and by S n ( A ) the set of symmetric matrices of size
n taking values in A ⊂ R .
2.1 Fused Gromov-Wasserstein distance
An undirected attributed graph G with n nodes can be modeled in the OT context as a tuple ( C , F , h ) , where C ∈ S n ( R ) is a matrix encoding relationships between nodes, F = ( f 1 , ..., f n ) (cid:62) ∈ R n × d is a node feature matrix and h ∈ Σ n is a vector of weights modeling the relative importance of the nodes within the graph (Figure 1, left).
Without any prior knowledge, uniform weights can be chosen ( h = 1 n /n ).
The matrix C can be the graph adjacency matrix, the shortest-path matrix or any other description of the node relationships (i.e. the topology) of the graph [47, 58, 15].
Let us now consider two such graphs ( C , F , h ) and ( C , F , h ) , of respective sizes n and n (with possibly n (cid:54) = n ).
The
Fused Gromov-Wasserstein ( FGW ) distance is deﬁned for α ∈ [0 , 1] as [58, 59]:
[Equation]
where U ( h , h ) :
= { T ∈ R n × n + | T 1 n = h , T
[Equation]
1 n = h } is the set of admissible coupling between h and h .
FGW aims at ﬁnding an optimal coupling T (cid:63) by minimizing a trade-off cost, via α , between a
Wasserstein (W) cost on the features and a Gromov-Wasserstein (GW) cost on the similarity matrices, both sharing the same coupling.
The optimal coupling T (cid:63) acts as a soft matching of the nodes, which tends to associate pairs of nodes that have similar pairwise relations in C and C (GW cost), and similar features in F and F (W cost).
Interestingly, FGW deﬁnes a metric on the space of attributed graphs.
In particular, if C and C are shortest-path matrices, the FGW distance vanishes if and only if the two attributed graphs are the same up to a permutation [59, Theorem 3.2].
Such invariance involves that two graphs strongly isomorphic according to Weisfeiler-Lehman base tests [33, 56] will have a zero FGW distance for any α and, more importantly,
FGW α = 0 implies that the graphs are strongly isomorphic 1 .
When
C , C are any symmetric matrices, we can mention that GW ( α = 1 ) also deﬁnes a pseudo-distance [55, Theorem 5.8] with respect to the notion of weak isomorphism [55, 15].
Solving for FGW.
The optimization problem 1 is a non-convex quadratic program [59, equation 6], whose non-convexity comes from the GW cost.
A possible optimization procedure to solve this problem is a Conditional Gradient (CG) algorithm, which is known to converge to a local optimum [31].
The computational complexity of each iteration is O ( n 2 n + n 2 n ) [47].
Thus, if two graphs of considerably different sizes are considered, the complexity is quadratic with respect to the largest size.
Existing attempts to reduce this computational cost either exploit entropic regularization of OT [47, 50] or graph partitioning [66, 13].
2.2 Template-based (T)FGW Graph Neural Networks
Building upon the FGW distance and its properties, we propose a simple layer for a GNN that takes a graph ( C , F , h ) as input and computes its FGW distances to a list of K template graphs
G :=
{ ( C k , F k , h k ) } k ∈
[[ K ]] as follows :
[Equation]
We postulate that this graph representation can be discriminant between the observed graphs due to
FGW .
This claim relies on the theory of [2] allowing one to learn provably strongly discriminant classiﬁers based on the distances from the observed graphs and templates that are sampled from the dataset (see e.g. [48] adopting the Wasserstein distance).
However such an approach often requires a large amount of templates which might be prohibitive if the distance is costly to compute.
Instead, we propose to learn the graph templates G in a supervised manner.
In the same way, we also learn the trade-off parameter α on the data.
As such, the TFGW layer can automatically adapt to the data whose discriminating information can be discovered either in the features or in the structure of the graphs, or in a combination of the two.
Moreover, the template structures can leverage on any type of input representation C i since they are learnt directly from the data.
Indeed, in the numerical experiments we implemented the model using either adjacency matrices (ADJ) that provide more interpretable templates (component C i,j ∈ [0 , 1] can be seen as a probability of link between nodes) or shortest path matrices (SP) that are more complex to interpret but encode global relations between the nodes.
The TFGW layer can be used directly as a ﬁrst layer to build a graph representation feeding a fully connected network (MLP) for e.g. graphs classication.
In order to enhance the discriminating power of the model, we propose to put a GNN (denoted by φ u and parametrized by u ) on top of the
TFGW layer.
We assume in the remainder that this GNN model φ u is injective in order to preserve isomorphism relations between graphs (see [67] for more details).
With a slight abuse of notation, we write φ u ( F ) to denote the feature matrix of an observed graph after being processed by the GNN.
Learning with TFGW-GNN.
We focus on a classiﬁcation task where we observe a dataset D of
I graphs {G i =
( C i , F i , h i ) } i ∈
[[ I ]] with variable number of nodes { n i } i ∈
[[ I ]] and where each graph is assigned to a label y i ∈ Y , with Y a ﬁnite set.
The full model is illustrated in Figure 1.
We ﬁrst process the features of the nodes of the input graphs via the GNN φ u , then use the TFGW layer to represent the graphs as vectors in R K .
Finally we use the ﬁnal MLP model ψ v : R K → Y parameterized by v , to predict the label for any input graph.
The whole model is learned in a end-to-end fashion by minimizing the cross-entropy loss on the whole dataset leading to the following optimization problem :
[Equation]
Notable parameters of (3) are the template graphs in the embeddings { ( C k , F k , h k ) } and more precisely their pairwise node relationship C k , node features F k and the distribution on the nodes h k on the simplex.
The last parameter reweighs individual nodes in each template and performs nodes selection when some weights are exactly 0 [63, 62].
Finally, the global parameter α is also learnt from the whole dataset.
Although it is possible to learn a different α per template, we observed that this extra level of ﬂexibility is prone to overﬁtting, and we will not consider it in the experimental section.
Optimization and differentiation of TFGW.
We propose to solve the optimization problem in (3) using stochastic gradient descent.
The FGW distances are computed by adapting the conditional gradient solver implemented in the POT toolbox [18].
The solver was designed to allow backward propagation of the gradients w.r.t.
all the parameters of the distance and was adapted to also compute the gradient w.r.t.
the parameter α .
The gradients are obtained using the Envelop Theorem [1] allowing to keep T (cid:63) constant.
We used Pytorch [45] to implement the model.
The template structure
C k , node weights h k and α are updated with a projected gradient respectively on the set of symmetric matrices S n k
( R + ) ( S n k ([0 , 1]) when C i are adjacency matrices), the simplex Σ n k and [0 , 1] .
The projection onto the probability simplex of the node weights leads to sparse solutions [16], therefore the size of each ( C k , F k , h k ) can decrease along iterations hence reducing the effective number of their parameters to optimize.
This way the numerical solver can leverage on the fact that many computations are unnecessary as soon as the weights are set to zero.
Note that the FGW solver from
POT uses an OT solver implemented in C++ on CPU which means that it comes with some overhead (memory transfer between GPU and CPU) when training the model on GPU.
Still the multiple FGW distances computation has been implemented in parallel on CPU with a computational time that remains reasonable in practice (see experimental section 3.3).
While a GPU solver can be found when using entropy regularized FGW, it introduces a new parameter related to the regularization strength which is more cumbersome to set, and that we did not consider it in the experiments.
Properties of the TFGW layer.
We now discuss a property of the proposed layer resulting from the properties of FGW (see Section 2.1).
We have the following result:
Lemma
The TFGW embeddings are invariant to strong isomorphism.
This lemma directly stems from the fact that FGW is invariant to strong isomorphism of one of its inputs.
This proposition implies that two graphs with any aforementioned representation which only differ by a permutation of the nodes will share the same TFGW embedding.
Moreover such a property holds for any mapping φ u which is injective, such as a Multi-Layer Perceptron (MLP) [22] or any GNN with a sum aggregation scheme as described in [67].
Moreover, the optimal coupling T (cid:63) resulting from TFGW between ( C i , φ u ( F i ) , h i ) and the template ( C k , F k , h k ) , will encode correspondances between the nodes of the graph and the nodes of the template that will be propagated during the backward operation.
The size of the inputs, the size of the templates and their respective weight h k will play a crucial role regarding this operation.
Also note that, since the templates are estimated here to optimize a supervised task, they will promote discriminant distance embedding instead of graph reconstruction quality as proposed in other FGW unsupervised learning methods [63, 62].
Numerical experiments
This section aims at illustrating the performances of our approach for graph classiﬁcation in synthetic and real-world datasets.
First, we showcase the relevance of our TFGW layer on existing synthetic datasets known to require expressiveness beyond the WL-test (Section 3.1).
Then we benchmark our model with state-of-the-art approaches on well-known real-world datasets (Section 3.2).
We ﬁnally discuss our results through a sensitivity analysis of our models (Section 3.3).
3.1 Synthetic datasets beyond WL test
[Table]
Identiﬁcation of graphs beyond the WL test is one important challenge faced by the GNN community.
In order to test the ability of TFGW to handle such fundamentally difﬁcult problems we consider two synthetic datasets: 4-CYCLES [34, 44] contains graphs with (possibly) disconnected cycles where the label y
i is the presence of a cycle of length 4; SKIP-CIRCLES [12] contains circular graphs with skip links and the labels (10 classes) are the lengths of the skip links among { 2 , 3 , 4 , 5 , 6 , 9 , 11 , 12 , 13 , 16 } .
We compare the performances of the TFGW layer for embedding such graphs with GIN [67] designed to be at least as expressive as the WL test, and DropGIN [44] which proposed a successful dropout technique to overcome some drawbacks of GIN.
We replicate the benchmark of [44] by considering for both GIN and DropGIN, 4 GIN layers for 4-CYCLES, and 9 GIN layers for SKIPCIRCLES as the skip links can form cycles of up to 17 hops.
Since the graphs do not have features we use directly the TFGW on the raw graph representation with α = 1 hence computing only the GW distance.
The GNN methods above artiﬁcially adds a feature equal to 1 on all nodes as they have the same degree.
For these experiments we use adjacency matrices for C i and we
[Table] [Table] [Table] [Table]
[Table] [Table] [Table] [Table]
[Table] ) WWL 86. [Table] [Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table] [Table] [Table] [Table]
[Table]
) Gain with
[Table] [Table]
investigate two ﬂavours of TFGW : 1) in TFGW -ﬁx we ﬁx the templates by sampling one template per class from the training dataset (this can be seen as a simpler FGW feature extraction); 2)
[Figure]
for TFGW we learn the templates from the training data (as many as the number of classes) as proposed in the previous sections.
Results are averaged over 10 runs and reported in Table 1.
TFGW based methods perform very well on both datasets with impressive results on SKIPCIRCLE when GNN have limited performances.
This is due to the fact that different samples from one class of SKIP-CIRCLE are generated by permuting nodes of the same graph and FGW distances are invariant to these permutations.
4-CYCLES has a more complex structure with intra-class heterogeneity and requires more than two templates to perform as good as DropGIN.
To illustrate this we have computed the accuracy on this dataset as a function of the number of templates K in Figure 2.
We can see that a perfect classiﬁcation is reached up to K = 4 for
TFGW , while TFGW -ﬁx still struggles to generalize at K = 20 .
This illustrates that learning the templates is essential to keep K (and numerical complexity) small while ensuring good performances.
3.2 Graph classiﬁcation benchmmark
We now evaluate and compare the performances of our TFGW GNN with a number of state-of-the-art graph classiﬁers, from kernel methods to GNN.
The numerical experiments are conducted on real life graph datasets to provide a fair benchmark of all methods on several heterogeneous graph structures.
Datasets.
We use 8 well-known graph classiﬁcation datasets [24]: 5 bioinformatics datasets among which 3 have discrete node features (MUTAG, PTC, NCI1 [28, 52]) and 2 have continuous node features (ENZYMES, PROTEINS[6]) and 3 social network datasets (COLLAB, IMDB-B, IDBM-M [69]).
In order to analyse them with all methods, we augment unattributed graphs from social networks with node degree features.
Detailed description and statistics on these datasets are reported in the supplementary material.
Baselines.
We benchmark our approaches to the following state-of-the-art baselines for graphs classiﬁcation, split into 3 categories: i) kernel based approaches , including FGW [59] operating on adjacency and shortest-path matrices, the WL subtree kernel [52, WL] and the Wasserstein WL kernel [56, WWL].
For these methods that do not require a stopping criterion dependent on a validation set, we report results using for parameter validation a 10-fold nested cross-validation [59, 29] repeatedtimes.
ii)
OT based representation learning models, including WEGL [27] and OT-GNN [11].
iii)
GNN models , with global or more sophisticated pooling operations, including PATCHY-SAN [42],
DIFFPOOL [70], PPGN [36], GIN [67] and its augmented version through structure perturbations
DropGIN [44].
For all these methods, we adopt the hyper-parameters suggested in the respective papers, but with a slightly different model selection scheme, as detailed in the next paragraph.
Benchmark settings.
Recent GNN literature [37, 67, 36, 44] successfully addressed many limitations in terms of model expressiveness compared to the WL tests.
Within that scope, they suggested to benchmark their models using a 10-fold cross-validation (CV) where the best average accuracy on the validation folds was reported.
We suggest here to quantify the generalization capacities of
GNN based models by performing a 10-fold cross validation with a holdout test set never seen during training.
For each split, we track the accuracy on the validation fold every 5 epochs, then the model whose parameters maximize that accuracy is retained.
Finally, the model used to predict on the holdout test set is the one with maximal validation accuracy averaged across all folds.
This setting is more realistic than a simple 10-fold CV and allows a better understanding of the generalization performances [4].
This point explains why some existing approaches have here different performances than those reported in their original paper.
For all the TFGW based approaches we empirically study the impact of the input structure representation by considering adjacency (ADJ) and shortest-path (SP) matrices C i .
For all template based models, we set the size of the templates to the median size of the observed graphs.
We validate the number of templates K in { β |Y|} β , with β ∈ { 2 , 4 , 6 , 8 } and |Y| the number of classes.
Only for ENZYMES with 6 classes of 100 graphs each, we validate β ∈ { 1 , 2 , 3 , 4 } .
All parameters of our TFGW layers highlighted in red in Figure 1 are learned while φ u is a GIN architecture [67] composed of L = 2 layers aggregated using the Jumping Knowledge scheme [68] known to prevent overﬁtting in global pooling frameworks.
For OT-GNN we validate the number of GIN layers in L ∈ { 2 , 4 } .
Finally for fairness, we validate the number of hidden units within the
GNN layers and the application of dropout on the ﬁnal MLP for predictions, similarly to GIN and
DropGIN.
Results analysis.
The results of the comparisons in terms of accuracy are reported in Table 2.
Our
TFGW approach consistently outperforms with signiﬁcant margins the state-of-the-art approaches from all categories.
Even if most of the benchmarked models can perfectly ﬁt the train sets by learning implicitly the graphs structure [67, 44, 36], enforcing such knowledge explicitly as our TFGW layer does (through FGW distances) leads to considerably stronger generalization performances.
On 7 out of 8 datasets, TFGW leads to better performances while operating on adjacency matrices (TFGW
ADJ) than on shortest-path ones (TFGW SP).
Interestingly, this ranking with respect to those input representations does not necessarily match the one of the FGW kernel which extracts knowledge from the graph structures C i through FGW , as our TFGW layer.
These different dependencies to the provided inputs may be due to the GNN pre-processing of node features which suggests the study of its ablation.
[Table]
Finally to complete this analysis, we report in Tablethe number of parameters of best selected models across various methods, for the dataset PTC.
We refer the reader interested in an extended benchmark to the supplementary material.
Our TFGW leads to better classiﬁcation performances while having comparable number of parameters than these competitors.
We also reported for these models, their averaged prediction time per graph.
These measures were taken on CPUs (Intel Core i9-9900K CPU, 3.60 GHz) in order to fairly compare the numerical complexity of these methods, as OT solver used in TFGW and OT-GNN are currently limited to these devices (see the detailed discussion in Section 2.1).
Although the theoretical complexity of our approach is at most cubic in the number of nodes, we still get in practice a fairly good speed for classifying graphs, in comparison to the other competitive methods.
3.3 Ablation study, sensitivity analysis and discussions
In this section we inspect the role of some of the model parameters ( α in FGW , weights estimation in the templates, depth of the GNN φ u ) in terms of the classiﬁcation performance.
To this end, we ﬁrst conduct on all datasets an ablation study on the graph template weights h k and the number of
GIN layers in φ u .
Then, we take a closer look at the estimated trade-off parameters α and provide a sensitivity analysis w.r.t.
the number of templates and the number of GIN layers.
Ablation Study.
Following the same procedure as in Section 3.2, we benchmark the following settings for our TFGW models: for adjacency (ADJ) and shortest path (SP) representations C i , we[-learn the distance layers either directly, on the raw data ( i.e. L = 0 , φ u = id ), or after embedding the data with ( L = 1) or ( L = 2) GIN layers. For L = 
 we either ﬁx the graph template weights h k-]
[Table]
81.1(2.5) TFGW (L=2)
ADJ learnt 96.4(3.3) 72.4(5.7) 73.8(4.6) 82.9(2.7) 88.1(2.5) 78.3(3.7) 56.8(3.1) 84.3(2.6)
[Figure]
[Figure] [Figure] [Figure] [Figure]
[Figure]
[Figure] [Figure]
learn the distance layers either directly, on the raw data ( i
[Figure]
{+e. L = 0 , φ u = id ), or after embedding the data with ( L = 1) or ( L = 2) GIN layers. For L = 
 we either ﬁx the graph template weights h k+} uniformly or learn them.
The results (test accuracy) are reported in Table 4.
Learning the weights systematically improves the generalization capabilities of our models of at least 1% or 2% for both
ADJ and SP graph representations.
For a given number of graph templates, the weights learning allows to better ﬁt the speciﬁcities of the classes (e.g. varying proportion of nodes in different parts of the graphs).
Moreover, as weights can become sparse in the simplex during training they also allow the model to have templates whose number of nodes adapts to the classiﬁcation objective, while bringing computational beneﬁts as discussed in Section 2.2.
Those observations explain why we learnt them by default in the benchmark of the previous subsection.
Next, we see in Table 4 that using GNN layers as a pre-processing for our TFGW layer enhances generalization powers of our models, whose best performances are obtained for L = 2 .
Interestingly, for L = 0 , TFGW with SP matrices outperforms TFGW with ADJ matrices, meaning that the shortest path distance brings more discriminant information on raw data.
But when L ≥ 1 ( i.e. when a GNN pre-processes the node features), TFGW with ADJ matrices improves the accuracy.
An explanation could be that the GNN φ u can somehow replicate (and outperform) a SP metric between nodes.
This emphasizes that the strength of our approach clearly exhibited in Table 2 lies in the inductive bias of our FGW distance embedding.
This last point is further reinforced by additional experiments reported in the supplementary material, where GIN layers (used by default for our TFGW model) are replaced by GAT layers [60].
[Figure]
Importance of the structure/feature aspect of
FGW.
To the best of our knowledge we are the ﬁrst to actually learn the trade-off parameter α of FGW in a supervised way.
For this matter, we verify that our models did not converge to degenerated solutions where either the structure ( α = 0 ) or the features ( α = 1 ) are omitted.
To this end we report in Figure 4 the distributions of the estimated α for some models learnt on datasets PTC and IMDB-B, where features are respectively existing in the dataset or created using node degrees.
We can see that for both kinds of input graph representations, α parameters are strictly between 0 and 1 .
One can notice the variances of those distributions illustrating the non-uniqueness of this trade-off parameter coming from the non-convexity of our optimization problem (a given value of α can potentially be compensated by the scaling of the GNN output).
Unfortunately, the analysis of the real relative importance between structures and features can not be achieved only by looking at those values as the node embeddings and templates are different across models and data splits.
Visualizing the TFGW embedding.
In order to interpret the TFGW embedding, we illustrate in
Figure 3 the PCA projection of our distance embeddings learned on PTC with L = 0 and L =and the number of templates K varying in { 4 , 8 } .
For this experiment, we have chosen the PCA because it allows to have a more interpretable low dimensional projection that preserves the geometry compared to local neighbourhood based embeddings such as TSNE [57] or UMAP [38].
As depicted in the ﬁgure, the learned templates are extreme points in the embedding space of the PCA.
This result is particularly interesting because existing unsupervised FGW representation learning methods tend toward estimating templates that belong to the data manifold, or to form a “convex enveloppe” of the data to ensure good reconstruction [63, 65].
On the contrary, the templates learned through our approach seem to be located on a plane in the PCA space while the samples evolve orthogonally to this plane (when the FGW distance increases).
In a classiﬁcation context, this means that the learned templates will not actually represent realistic graphs from the data but might encode “exaggerated” or “extreme” features in order to maximize the margin between classes in the embedding.
To reinforce this intuition, we added plots of the estimated templates in the supplementary.
Finally, in the ﬁgure, the samples are coloured w.r.t.
their class and the templates are coloured by their predicted class.
Interestingly, the classes are already well separated with 4 templates but the separation is clearly non-linear whereas using GNN pre-processing and a larger number of templates leads to a linear separation of the two classes.
[Figure]
Sensitivity to the number of templates and GNN layers.
To illustrate the sensitivity of our TFGW layer to the number of templates K and the number of GNN layers L in φ u , we learned our models on the PTC dataset with L varying in { 0 , 1 , 2 , 3 , 4 } .
We follow the same procedure than in the benchmark of
Section 3.2 regarding the validation of K and the learning process, while ﬁxing the number of hidden units in the GNN layers to 16 .
The test accuracy distributions for all settings are reported in Figure 5.
Two phases are clearly distinguishable.
The ﬁrst one for L ≤ 2 , where for each L we see that the performance across the number of templates steadily increases, and the second for L > 2 where this performance progressively decreases as a function of L .
Moreover, in the ﬁrst phase performances are considerably dependent on K to compensate for a simple node representation, while this dependency is mitigated in the second which exhibits a slight overﬁtting.
Note that these deeper models still lead to competitive results in comparison with benchmarked approaches in Table 2, with best averaged accuracies of 70 .
( L = 3) and 67 . 4 ( L = 4) .
On one hand, these observations led us to set our number of layers to
L = 2 for all benchmarked datasets which lead to strong generalization power.
On the other hand, deeper models might be a way to beneﬁt from our FGW embeddings with very few templates which can be interesting from a computational perspective on larger graph datasets.
Conclusion
We have introduced a new GNN layer whose goal is to represent a graph by its distances to template graphs, according to the optimal transport metric FGW .
The proposed layer can be used directly on raw graph data as the ﬁrst layer of a GNN or can also beneﬁt from more involved node embedding using classical GNN layers.
In a graph classiﬁcation context, we combined this TFGW layer with a simple MLP model.
We demonstrated on several benchmark datasets that this approach compared favorably with state-of-the-art GNN and kernel based classiﬁers.
A sensitivity analysis and an ablation study were presented to justify the choice of several parameters explaining the good generalization performances.
We believe that the new way to represent complex structured data provided by TFGW will open the door to novel and hopefully more interpretable GNN architectures.
From a practical perspective, future works will be dedicated to combine TFGW with fast GPU solvers for network ﬂow [51].
This would greatly accelerate our approach and more generally OT based deep learning methods.
We also believe that the FGW distance and its existing extensions can be used with other learning strategies including semi-relaxed FGW [62] for sub-graph detection.
Acknowledgements
This work is partially funded through the projects OTTOPIA ANR20-CHIA-0030 and 3IA Côte d’Azur Investments ANR-19-P3IA-0002 of the French National Research Agency (ANR).
This research was supported by 3rd Programme d’Investissements d’Avenir ANR-18-EUR-0006-02.
This action beneﬁted from the support of the Chair "Challenging Technology for Responsible Energy" led by l’X – Ecole polytechnique and the Fondation de l’Ecole polytechnique.
This work is supported by the ACADEMICS grant of the IDEXLYON, project of the Université de Lyon, PIA operated by ANR-16-IDEX-0005.
This project was supported in part by the AllegroAssai ANR project
ANR-19-CHIA-0009.
The authors are grateful to the OPAL infrastructure from Université Côte d’Azur for providing resources and support.
Checklist
For all authors... (a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’s contributions and scope?
[Yes] Every key points are encapuslated in the contributions part at the end of the introduction.
(b) Did you describe the limitations of your work?
[Yes] Yes the limitations of our work mostly lie in its computational complexity and the lack of a current implementation designed for GPU parallelization of the FGW distance solver in section 2.1 and 3.2.
(c) Did you discuss any potential negative societal impacts of your work?
[N/A] (d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
[Yes] 2.
If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?
[-[Equation] Lemma 
 The TFGW embeddings are invariant to strong isomorphism. Proof-][Yes] Indeed the full set of assumptions of all theoretical results are in the paper.
Few technical deﬁnitions which are common in the OT litterature have been reported to the supplementary material.
(b) Did you include complete proofs of [-graphs G 
 = ( C 1 , F 1 , h 1 ) and G 2 = ( C 2 , F 2 , h 2 ) deﬁned as in the subsection 6.1. Assume that G 1 and G 2 are strongly isomorphic .-]all theoretical results?
[-and there exists a permutation matrix P ∈ { 0 , 1 } n × n such that C 
 = P C 1 P (cid:62) , F 2 = P F 1 and h 2 = P h-][Yes] A simple intuition for Lemma 1 is given in the paper and more details are provided in the supplementary.
[-where we used C 
 = ( P (cid:62) C 2 P ) 2 = P (cid:62) C 22 P and-]If you ran experiments... (a) Did you include the code, data, and [-φ ( F 2 ) 
 = ( P φ ( F 1 )) 2 =-]instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
[Yes] The code is provided in the supplementary material.
A read me has been made to reproduce these experiments more easily.
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
[Yes] For each kind of experiments most settings are provided in the core of the paper (see section 3).
Complementary details on the experiments are reported in the supplementary material (c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[Yes] Yes, each tables and ﬁgures contain such information on the distributions when relevant.
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?
[No] 4.
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?
[Yes] Our code mostly depends on the POT and Pytorch librairies which are quoted in section 2.
(b) Did you mention the license of the assets?
[Yes] Those assets have been following the format prescribed in their respective websites/papers.
(c) Did you include any new assets either in the supplemental material or as a URL?
[Yes] Some implementations of the methods reproduced in the benchmark section 3.2 depend on assets which were not used initially for our own method.
Therefore we quoted them in the supplementary (d) Did you discuss whether and how consent was obtained from people whose data you’re using/curating?
[Yes] Used datasets are public ones, we quoted their designers as sign of consent.
(e) Did you discuss whether the data you are using/curating contains personally identiﬁable information or offensive content?
[No] Used datasets are well-known public ones that do not contain such content.
If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable?
[N/A] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?
[N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
[N/A]