Wasserstein Embedding for Graph Learning
Authors: Soheil Kolouri
Navid Naderializadeh
Gustavo K. Rohde
Heiko Hoffmann
Venue: ICLR 2021 Poster (2021)
Keywords: Wasserstein
graph embedding
graph-level prediction

Abstract: We present Wasserstein Embedding for Graph Learning (WEGL), a novel and fast framework for embedding entire graphs in a vector space, in which various machine learning models are applicable for graph-level prediction tasks. We leverage new insights on defining similarity between graphs as a function of the similarity between their node embedding distributions. Specifically, we use the Wasserstein distance to measure the dissimilarity between node embeddings of different graphs. Unlike prior work, we avoid pairwise calculation of distances between graphs and reduce the computational complexity from quadratic to linear in the number of graphs. WEGL calculates Monge maps from a reference distribution to each node embedding and, based on these maps, creates a fixed-sized vector representation of the graph. We evaluate our new graph embedding approach on various benchmark graph-property prediction tasks, showing state-of-the-art classification performance while having superior computational efficiency. The code is available at https://github.com/navid-naderi/WEGL.

Published as a conference paper at ICLR 2021
WASSERSTEIN EMBEDDING FOR GRAPH LEARNING
Soheil Kolouri' Navid Naderializadeh", Gustavo K. Rohde', & Heiko Hoffmann" HHRL Laboratories, LLC., +University of Virginia
(skclcuri, nnaderia zadeh, fmann [hrl. con, gustavofvirginia.ce
ABSTRACT
We present Wasserstein Embedding for Graph Learning (WEGL), a novel and fast framework for embedding entire graphs in vector space, in which various machine learning models are applicable for graph-level prediction tasks. We leverage new insights defining similarity between graphs as function of the similarity between their node embedding distributions. Specifically, we use the Wasserstein distance measure the dissimilarity between node embeddings of different graphs. Unlike prior work, we avoid pairwise calculation of distances between graphs and reduce the computational complexity from quadratic linear the number of graphs. WEGL calculates Monge maps from reference distribution to each node embedding and, based on these maps, creates fixed-sized vector representation of the graph. We evaluate on various benchmark graph- graph-property prediction tasks, showing state-of-the-ar classification performance while having superior computational efficiency. The code is available at htttps://gitii
graph embedding approach
com
INTRODUCTION
Many exciting and practical machine learning applications involve learning from graph-structured data. While images, videos, and temporal signals (c.g., audio or biometrics) are instances of data that are supported on grid-like structures, data in social networks, cyber-physical systems, communication networks, chemistry, and bioinformatics often live on irregular structures (Backstrom Leskovee 2011; Sadreazami et al.] 2017] Jim et al.) 2017, Agrawal et al., 2018 Naderializadeh et al. (0200). One can represent such data as (attributed) graphs, which are universal data structures. Efficient and generalizable learning from graph-strun strct-structured data opens the door to vast number of applications, which were beyond the reach of classic machine learning (ML) and, more specifically, deep learning (DL) algorithms.
Analyzing graph-structured data has received significant attention from the ML, network science, and signal processing communities over the past few years. On the one hand, there has been rush toward extending the success deep neural networks to graph-structured data, which has led to variety of graph neural network (GNN) architectures. On the other hand, the research kernel approaches (Gärtnnr et al. 2003), perhaps most notably the random walk kernel (Kashima et al. 2003) and the Weisfeiler-Lehmal (WL) kernel (Shervashidze et al.. 2011 Rieck 2019: 2020), remains an active field of study and the methods developed therein provide competitive performance in various graph representation tasks (see the recent survey by Kriege et al. (2020)).
2019 Morris et al.
To learn graph representations, GNN-based frameworks make use of three generic modules, which provide i) feature aggregation, ii) graph pooling (i.e., readout), and iii) classification (Hu* al. 2020). The feature aggregator provides a vector representation for each node of the graph, referred to as node embedding. The graph pooling module creates a representation for the graph from its node embeddings, whose dimensionality is fixed regardless of the underlying graph size, and which can then be analyzed using a downstream classifier of choice. On the graph kernel side, one leverages a kernel to measure the similarities between pairs of graphs, and uses conventional kernel methods to perform learning on a set of graphs (Hofmann et al.) 2008). A recent example of such methods is the framework provided by (Tognnaall et al. (2019), in which the authors propose novel node embedding inspired by the WL kernel, and combine the resulting node embeddings with the
" Denotes equal contribution.

---

Published as a conference paper at ICLR 2021
Kolouri et al.) 2017) to measure the dissimilarity between two graphs. Afterwards, they leverage conventional kernel methods based on the pairwise-mest measured dissimilarities to perform learning on graphs.
Wasserstein distance (Villani 2008
Considering the ever- vver-increasing scale of graph datasets, which may contain tens of thousands of graphs millions to billions of nodes per graph, the issue of scalability and algorithmic efficiency becomes of vital importance for graph learning methods (Hernandez & Brown, 2020; Hu et al.) 2020). However, both of the aforementioned paradigms of GNNs and kernel methods suffer in this sense. On the GNN side, acceleration of the training procedure is challenging and scales poorly as the graph size grows (Bojchevski et al., 2019). On the graph kernel side, the need for calculating the matrix of all pairwise similarities can be burden in datasets with a large number of graphs, especially if calculating the similarity between each pair of graphs is computationally expensive. For instance, in the method proposed in (Togninalli et al. 2019), the computational complexity of each calculation of the Wasserstein distance cubic in the number of nodes (or linearithmic for the entropy- regulay-regularized distance).
To overcome these issues, inspired by the linear optimal transport framework of (Wang et al. 2013), we propose a linear Wasserstein Embedding for Graph Learning, which we refer to as WEGL. Our proposed approach embeds a graph into a Hilbert space, where the lg distance between two embedded graphs provides a true metric between the graphs that approximates their 2- 2-Wasserstein distance. For a set of M graphs, the proposed method provides:
1. Reduced computational complexity of estimating the graph Wasserstein distance Togninalli jet al. (2119) for a dataset of M graphs from a quadratic complexity in the number of graphs, i.e.,
M(M calculations, to linear complexity, i.e. M calculations of the Wasserstein
distance; and
2. An explicit Hilbertian embedding for graphs, which is not restricted kernel methods, and therefore can be used in conjunction with any downstream classification framework.
We show that compared to multiple GNN and graph kernel baselines, WEGL achieves either state- of-the-art or competitive results on benchmark graph-level classification tasks, including classical graph classification datasets (Kersting et al., 2020) and the recent molecular property-predice benchmarks (Hu et al., 2020). We also compare the algorithmic efficiency of WEGL with two baseline GNN and graph kernel methods and demonstrate that it is much more computationally efficient relative those algorithms.
2
BACKGROUND AND RELATED WORK
In this section, we provide a brief background on different methods for deriving representations for graphs and an overview Wasserstein distances by reviewing the related work in the literature.
2.1 GRAPH REPRESENTATION METHODS
Let G = (V, E) denote a graph, comprising a set of nodes and set of edges & C V2, where two nodes u, E V are connected to each other if and only if (u, v) E E For each node E V. we define its set of neighbors as N, & {u V (u, v) E E}. The nodes of the graph G may have categorical labels and/or continuous attribute vectors. We use a unified notation of x, E FR to denote the label and/or attribute vector of node Moreover, we use Weo E IR to denote the edge feature vector for any edge (u, v) E E, where E' denotes the edge feature dimensionality. Node and edge features may be present depending on the graph dataset under consideration.
V. where F denotes the node feature dimensionality.
To learn graph properties from the graph structure and its node/edge features, one can use function ab : G + H to map any graph G in the space of all possible graphs G to an embedding (6G)) in Hilbert space H. Kernel methods have been among the most popular ways of creating such graph embeddings. A graph kernel is defined as function k G2 + R, where for two graphs G and G'. k(G, G*) represents the inner product of the embeddings (G) and (G") over the Hilbert space H. The mapping v could be explicit, as in graph convolutional neural networks, or implicit as in the case
Note that this definition includes both directed and undirected graphs, where in the latter case, for each edge (u, v) E E, the reverse edge (v, u) is also included in E.
2

---

Published as a conference paper at ICLR 2021
of the kernel similarity function k((;)) .) (i.e., the kernel trick). Kriege et al. (0144) provide thorough discussion on explicit and implicit embeddings for learning from graphs.
Kashima et al. (2003) introduced graph kernels based on random walks labeled graphs. Subse- quently, shortest- -pth kernels were introduced in (Borgwardt & Kriegel, 2005). These works have
been followed by graphlet and Weisfeiler-Lehmar subtree kernel methods (Shervashidze et al.] (009
2011; Morris et al. 2017). More recently, kernel methods using spectral approaches (Kondor &
Pan 2016), assignment-based approaches (Kriege et al. 2016 Nikolentzos al. (017)), and graph decomposition algorithms (Nikolentzos et al. 2018) have also been proposed in the literature.
Despite being successful for many years, kernel methods often fail leverage the explicit continuous features that are provided for the graph nodes and/or edges, making them less adaptable to the underlying data distribution. To alleviate these issues, and thanks in part to the prominent success of deep learning in many domains, including computer vision and natural language processing, techniques based on graph neural networks (GNNs) have emerged as an alternative paradigm for learning representations from graph-based data. In general, a GNN comprises multiple hidden layers, where at each layer, each node combines the features of its neighboring nodes the graph to derive a new feature vector. At the GNN output, the feature vectors of all nodes are aggregated using readout function (such as global average pooling), resulting in the final graph embedding (G) More details on the combining and readout mechanisms of GNNs are provided in Appendix A.4
Kipf and Welling (Kipl & Welling 2016) proposed a GNN architecture based on a graph convo- lutional network (GCN) framework. This work, alongside other notable works geometric deep learning (Defferrard et al., 2016), initiated a surge of interest in GNN architectures, which has led several architectures, including the Graph Attention network (GAT) (Veličković et al. 2017). Graph SAmple and aggreGatE (GraphSAGE) (Hamilton et al.) 2017). and the Graph Isomorphism Network (GIN) (Xu et al., 2019). Each of these architectures modifies the GNN combining and readout functions and demonstrates state-of-the-an performance in variety of graph representation learning tasks.
2.2
WASSERSTEIN DISTANCES
Let μi denote a Borel probability measure with finite ndh moment defined on Z C Rd, with corre- sponding probability density function pt, i.e., (=) = (z)dd... The 2- 2-Waseerstein distance between H4; and Hy defined on 2, Z' C Rd is the solution to the optimal mass transportation problem with &
transport cost (Villani 2008):
( inf (ET(()))) 2x
/ ||= 11=-22(p(7(e,2)))
W2(ps, (bj
(1)
where ['44 (4y) is the set of all transportation plans 7 E [[ (4y such that ((A x Z') = (A) and <Z x B)= = (B) for any Borel subsets A S Z and B C Z'. Due to Brenier's theorem (Brenier 1991), for absolutely continuous probability measures 14 and Hyy (with respect to the Lebesgue measure), the 2- W-sssresneein distance can be equivalently obtained from
inf (EM / || [1=== ([=((())(du((e))
Wa(ui,)))
(2)
where M P( (2))= = ( :2 Z'ISmm = {5) and f represents the pushforward of measure
Ms. characterized as
f dps(22)
Jf dz((z) for any Borel subset B C Z'.
The mapping / is referred to as a transport map (Kolouri et al. 2017). ), and the optimal transport map called the Monge map. For absolutely continuous measures, the differential form of the above equation takes the following form, de (f))) pz(2), which is referred to as the Jacobian cquation. For discrete probability measures, when the transport plan 7 is a deterministic optimal coupling, such transport plan is referred to Monge coupling (Villani 2008).
Recently, Wasserstein distances have been used for representation learning on graphs and images (Togninalli et al. 2019 Zhang et al., 2020 Bécigneul et al. 2020). In particular, (Toginaall et al.) 2019) proposed a Wasserstein kernel for graphs that involves pairwise calculation of the Wasserstein distance between graph representations. Pairwise calculation of the Wasserstein distance, however,
3

---

Published as a conference paper at ICLR 2021
could be expensive, especially for large graph datasets. In what follows, apply the linear optimal transportation framework (Wang et al. 2013) to define a Hilbertian embedding, in which the lz distance provides true metric between the probability measures that approximates W2. We show that a dataset containing M graphs, this framework reduces the computational complexity from calculating [MM-1]
linear programs to M.
2
3 LINEAR WASSERSTEM EMBEDDING
Wang et al. (0013) and the follow-up works (Seguy & [aturi][2015) Kolouri et al. 2016 Courty al. 2018) describe frameworks for isometric Hilbertian embedding of probability measures such that the Euclidean distance between the embedded images approximates W2. We leverage the prior work and introduce the concept of linear Wasserstein embedding for learning graph embed- dings.
Molnaael
PH(2)
Py.
p(pp) =,0
3.1 THEORETICAL FOUNDATION
[o(p]) ol (0(m) Welpopu)
We adhere to the definition of the linear Wasserstein em- bedding for continuous measures. However, all deriva- tions hold for discrete measures as well. More precisely, let No be reference probability measure defined on 2 S RR with a positive probability density function po, s.t. dx = Po (z)ddz and Po(z) > O for Vz E Z. Let fi denote the Monge map that pushes μi into μu, i.e.,
[464] - (teg) la =s Wolmus)
Graphical representation of the linear Wasserstein embedding frame- work, where the probability distributions are mapped to the tangent space with respect to a fixed reference distribution. The figure is adapted from Kolouri et al. (2017).
Figure I:
ft argmin JEMP(HON) / f(=))l2dpu()). (4)
Define x(4)) A (fs id) /pn. where id(=) = 2, E 2 is the identity function. In cartography, such mapping is known as the equidistant azimuthal projection, while in differential geometry, is called the logarithmic map. The mapping o(-) has the following characteristics (partially illustrated in Figure D):
1. <(-) provides an isometric embedding for probability measures, i.e., using the Jacobian equation p; (e) (fo where *64) + id.
vPa
2. ((Ho) = 0, i.c., the reference is mapped to zero.
3. [(((4) - 6(H0) 2l2 = [6(()) (2 po). i.e., the mapping preserves distances to μo
4. [((())) - (Hy)]l2 se (2 My), i.e., the lz distance between 6(14) and @(2)), while being true metric between μı and py, is approximation of W2 My).
Embedding probability measures ( via o(-) requires calculating M Monge maps. The fourth characteristic above states that Φ(-) provides linear embedding for the probability measures. There- fore, we call the linear Wasserstein embedding. The mapping ((4) could be thought as the Reproducing Kernel Hilbert Space (RKHS) embedding of the measure, μ4; (Muandet et al. (017). In practice, for discrete distributions, the Monge coupling is used, which could approximated from the Kantorovich plan (i.e., the transport plan) via the so-called barycenteric projection (Wang al. 2013). We here acknowledge the concurrent work by Mialon et al. (202T), where the authors use a similar idea to the linear Wasserstein embedding as pooling operator for learning from sets of features. The authors demonstrate the relationship between their proposed optimal transport-basee pooling operation and the widespread attention pooling methods in the literature. A detailed description of
the capabilities the linear Wasserstein embedding framework is included in Appendix A. We
next provide the numerical details of the barycenteric projection (Ambrosio et al. 2008 Wang et 2013).
4

---

Published as a conference paper at ICLR 2021
3.2 NUMERICAL DETAILS
E IR be an arry Let us define Po to be a reference distribution, with Zo = [= [=f...... , = E Nxd where = LM ΣM N; and 2) E R ~ Po, Vj {I..... N}. The optimal transport plan between p; and po, denoted by Ty E IR RNN,, is the solution to the following linear program,
Consider a set of probability distributions (p: 1=1- and let Z4 = [
containing N; i..dd samples from distribution p,, i.e., z E R pi, Vk E fI,...,,W., No).
N
T = argmin EEIT
Tyn 1= - =//1²,
(5)
where IT, - {or RNNN, No Tik
Tyk 1, Ak (1, N.....,,,,,, N.).j (I....,
The Monge map then approximated from the optimal transport plan by barycentric projection via
F = N 2,) E E Nxd
(6)
Note that the transport plan n, could split the mass in z E Zo and distribute it on zis. The barycentric projection calculates the center of mass of the transportation locations for 2] to ensure that no mass splitting is happening (see Figure 5). and hence it approximates Monge coupling. Finally, the embedding can be calculated by (Zz) (F: - Zo)/ / Rxd With slight abuse of notation, we use ((p)) and ((Z,) interchangeably throughout the paper. Due to the barycenteric projection, here, <6-) is only pseudo-invertible
4 WEGL: A LINEAR WASSERSTEIN EMBEDDING FOR GRAPHS
The application of the optimal transport problem to graphs is multifaceted. For instance, some works focus on solving the "structured" optimal transport concerning an optimal probability flow, where the transport cost comes from distances
often unchanging underlying graph (Léonard et al. 2016
Essid & Solomon 2018) Titouan et al.) 2019). Here, we are interested in applying optimal transport
to measure the dissimilarity between two graphs (Maretic al., 2019 Togninalli al. 2019 Dong
& Sawin 2020). Our work significantly differs from (Maretic et al. 2019 Dong & Sawin 2020).
which measure the dissimilarity between non-attributed graphs based on distributions defined by their
Laplacian spectra and is closer to (Togninalli et al. 2019).
Our proposed graph embedding framework, termed Wasserstein Embedding for Graph Learning (WEGL). combines node embedding methods for graphs with the linear Wasserstein embedding explained in Section 3 More precisely, let {G, (V,, E,) M denote a set of M individual graphs, each with a set of possible node features heeee and set of possible edge features {War E,,))E,
Let h(.) be an arbitrary node embedding process, where h(G.) = Z, = [=i...., 21:1 E R| Having the node embeddings {Z we can then calculate reference node embedding Zo (see Section 4.2 for details), which leads to the linear Wasserstein embedding ((Z,) with respect to Zo. as described in Section 3 Therefore, the entire embedding for each graph G., E {1...., MM, is obtained by composing (.)) and h(-), i.e., (G.))) = (()().)) Figure 2 visualizes this process.
4.1 NODE EMBEDDING
There are many choices for node embedding methods (Chami et al. 2020p. These methods in
general could be parametric or non-parametric, e-g. as in propagation//diffsion bassd embeddings. Parametric embeddings often implemented via GNN encoder. The encoder capture different graph properties depending on the type of supervision (e.g., supervised or unsupervised), Self- supervised embedding methods have also been recently shown to be promising (Hu* al. 2020).
In this paper, for our node embedding process h(-). we follow a similar non non-parametric propagation//iiftsion based encoder in (Togninalli et al. 2019).. One of the appealing advantages of this framework is its simplicity, as there are no trainable parameters involved. In short, given a graph G (V,E) with node features {I., hev and scalar edge features (wov) (u,))E,, the
5

---

Published as a conference paper at ICLR 2021
Monge Coupling
Transport
k(G
Displacements
( ((((G)
G
Monge Coupling
Transport Displacements
h(G))
Node embedding, k())
Linear Wasserstein Embedding, (()
Figure 2: proposed graph embedding framework, WEGL, combines node embedding methods with linear Wasserstein embedding framework described in Section Given a graph G; = (V; &x). we first embed the graph nodes into a d-dimensional Hilbert space and obtain an array of node embeddings, denoted by h(..) = Z, E Kilxd We then calculate linear Wasserstein embedding of with respect a reference Zo. Le., (Z.), to derive the final graph embedding.
following instantiation of equation 12 to define the combining function as 11-1), A f1..... L}. VOEV,
Σ
Wuv
/deg(u)deg(1)
WEN,U(M)
where for any node v E its degree deg(v) is defined its number of neighbors in G augmented with seff-connections, i.e., deg(v) _ [W Note that the normalization of the messages between graph nodes by the (square root of) the two end-point degrees in equation 7 have also been used in other architectures, including GCN (Kipf Welling 2016). For the cases where the edge weights are not available, including self-connection weights {wee Juev, set them one. In Appendix A.5 show how we use an extension of equation to treat graphs with multiple edge features/labels. Finally, (1)
we let 2, g (ax ) represent the resultant embedding for each node V, where g(-) is local pooling process on single node (not a global pooling), e.g., concatenation or averaging.
4.2 CALCULATION OF THE REFERENCE DISTRIBUTION
To calculate the reference distribution,
the k-means clustering algorithm on U 7, with calculate the Wasserstein barycenter CCuturi & Doucet 2014) of the node embeddings or simply use N samples from normal distribution. While approximation of the 2-Wasserstein distance in the tangent space depends on the reference distribution choice, surprisingly, we see stable performance of WEGL. for different references. In Appendix A.2 we compare the performance of WEGL. with respect to different references.
N
M D= N; centroids. Alternatively,
M
EXPERIMENTAL EVALUATION
In this section, we discuss the evaluation results of our proposed algorithm on multiple benchmark graph classification datasets. We use the PyTorch Geometric framework (Fey & Lenssen 2019) for implementing WEGL. In all experiments, we scikit-learn for the implementation of our downstream classifiers on the embedded graphs (Buitinck et al. (013).
5.1 MOLECULAR PROPERTY PREDICTION ON THE OPEN GRAPH BENCHMARK
We first evaluate our algorithm on the molecular property prediction task on the ogbg-me m dataset. This dataset is part of the Open Graph Benchmark (Hu et al. 2020), which involves nove-level, link-level, and graph-level learning and prediction tasks on multiple datasets spanning diverse problem domains. The ogbg-molhiv dataset, in particular, is molecular tree-like dataset, consisting 41, 127 graphs, with an average number of 25.5 nodes and 27.5 edges per graph. Each graph molecule, with nodes representing atoms and edges representing bonds between them, and

---

Published as a conference paper at ICLR 2021
Method
Validation ROC-AUC
Test COOC-AUC (%)
GCN
Kipf Welling
2016
83.8 = 0.9
76.0 + 1.2
GIN Virtual Node
2019
84.8 0.7
77.1
GNN
DCeperGCN
2020
+: 0.6
78.6
HIMP Fey at al. 2020
78.8 0.8
GCN GraphNorm (Cai et
2020
79.0 + 1.1
78.8 +: 1.0
WBGL Random Forest
79.2 +: 2.2
75.5 1.5
Ours
Virtual
Random Forest
1.3
76.5
WEGL + Virtual Node + AutoML.
81.6 +: 0.6
79.1 0.3
GAP Virtual Node Random Forest
74.9 + 2.4
72.1 + 1.7
Table Graph classification results on the ogbg-molie dataset. The results for GCN and GIN are reported from (Hu et al. 2020). The best validation and test results are shown in bold.
it includes both node and edge attributes, characterizing the atom and bond features. The goal is to predict a binary label indicating whether or not molecule inhibits HIV replication.
To train and evaluate our proposed method, we use the scaffold split provided by the dataset, and report the mean and standard deviation of the results across 10 different random seeds. We perform grid search over set of hyperparameters and report the configuration that leads to the best validation performance. We also report virtual node variant of the graphs in our evaluations, where each graph is augmented with an additional node that is connected to all the original nodes in the graph. This node serves as a shortcut for message passing among the graph nodes, bringing any pair of nodes within at most two hops of each other. The complete implementation details can be found in Appendix A.5)
Table shows the evaluation results of WEGL the ogbg-molhiv dataset in terms of the ROC-AUC (i.e., Receiver Operating Characteristic Area Under the Curve), alongside multiple GNN- based baseline algorithms. Specifically,
show the results using two classifiers: A random forest classifier (Breiman, 200T) and an automated machine learning (AutoMI.) classifier using the Auto- Sklearn 2.0 library (Feurer al. 2020). As the table demonstrates, while WEGL. embeddings combined with random forest achieve a decent performance level, using AutoML further enhances the performance and achieves state-of-the-an test results on this dataset, showing the high expressive power of WEGL in scage graph datasets, without the need for end-to-end training.
Moreover, as an ablation study, we report the evaluation results on the ogbg-mo mooog-moll 1 hii dataset using a random forest classifier, where the Wasserstein embedding module after the node embedding process replaced with global average pooling (GAP) among the output node embeddings of each graph to derive the graph-level embedding. As the table demonstrates, there is a significant performance drop when using GAP graph embedding, which indicates the benefit of our proposed graph embedding method as opposed to average readout.
5.2 TUD BENCHMARK DATASETS
We also consider a set of social network, bioinformatics and molecule graph datasets (Kersting et al.. 2020). The social network datasets (IMDB BINARY, MMDDB-MULTE COLLAB, REDDIT BINARY, and REDDEEDDI IT ---LMMLT 5K) lack both node and edge features. Therefore, in these datasets we use one-hot representation of the node degrees as their initial feature vectors, as also used in prior work, and REDDIT-MULTI,
e.g., (Xu et al. (019). To handle the large scale of the and datasets, we clip the node degrees 500.
Moreover, for the molecule (PTC-MR) and bioinformatics (ENZ YMES and PROTEINS) datasets, we use the readily-provided node labels in (Kersting et al. 2020) as the initial node feature vectors. Besides, for PTC-MR which has edge labels, as explained in Appendix A.5] we use an extension of equation 7 to use the one-hot encoded edge features in the diffusion process. To evaluate the performance of WEGL, we follow the methodology used in (Yanardag & Vishwanathan 20155 Niepert et al. 2016 Xu et al. 2019), where for each dataset, we perform 10-fold cross- coass--vlidation with random splitting the entire dataset, conducting grid search over the desired set of hyperparameters as mentioned in Appendix [A.5] and we then report the mean and standard deviation of the validation accuracies achieved during cross-validation. For this experiment,
ensemble classifiers,
7

---

Published as a conference paper at ICLR 2021
Method
IMDD 3
IMDB M.
COLLAD
BC B
BB MEK
2IC NR
BNZYMD3
PROTEINE
DGCNN
al.
2018a
6.02.13.00
3..61:.4
71221.9
87.8.2.5
49.21112
58.6
38.915.77
72.913.5
GraphSAGE
Hamilion et 2017
68.8.14.5
4..6...55
7.99.1.7
843119
50.0.113
63.9.7.7
75.9.13.2
GNN
GIN Xuctal al. 2019
75.1.5..1
2.3.2.88
80.2119
924125
57.5115
64.6.17.0
76.2.12.8
GNTK Du al. 2019)
76.9433.6
52.8:466
83.6
67,9:6.9
75.6:4422
CapsUNN Kinyi & Chen
2019
73.14:4.8
50.3=2.6
79.6
52.915
54.745.7
76.3+:3.6
GraphNoom (Cai 2020
80.2+110
93542.
64.9
77.414499
DGK
Yanardap & Vidiwanttan
2015
67.0410.
44.6=0..5
73.1
7800104
41.3102
57.311.11
27.110.88
71.71:0.5
WT.
Shervschibee.tt al
2B11
73.8:1399
0..8=0.5
74.8+0.2
68.2402
51.24033
$7.042.0
53.24:111
72.9.0.66
RetGK
Zhane al.
2018b
71,0::0.6
46.7=0.6
73.6:0.3
90.402
54.240.3
67.91144
59.1+
75.2+00.
AWE
Burnaey
74.545.8
51.5=3.6
73.9119
87.9225
50.5419
35.8415.99
WWL
Togninalla et
2019
74.4:10.8
66.3:1:122
59.11:0.8
74.3::0.6
WEGL
SVM-RBF
7..4
51,7=311
78.6
92.1
86.1423
63.425.3
57.34:4.2
76. 76.044:44
Ours
WEGL Random Forest
7..4::.0
52.0=4.1
7..81:..5
920108
55.1125
67.511.77
60.5.1.5.9
7:,5
WEGL
75.245.0
52.3-2.9
2.06::.0
55.4166
66.2.16.9
60.0:16..
76.313.9
Table 2: Graph classification accuracy (%) of our method and comparison with the state-of-the GNNs and graph kernels (GKs) various TUD graph classification tasks. The results for DGCNN are reported from (Errica ct al. 2020). The top-three performers on each dataset are shown in bold.
namely Random Forest and Gradient Boosted Decision Tree (GBDT), together with kemel-SVM with an RBF kernel (SVM-RBF). Given that the Euclidean distance in the embedding space approximates the Wasseesttei distance, the SVM-RBF classification results are comparable with those reported by Togninalli et al. (019).
Table shows the classification accuracies achieved by WEGL on the aforementioned datasets as compared with several GNN and graph kernel baselines, whose results are extracted from the corresponding original papers. As the table demonstrates, our proposed algorithm achieves either state-of- the-art competitive results across all the datasets, and in particular, it is among the top-three performers in all them. This shows the effectiveness of the proposed linear Wasserstein embedding for learning graph-level properties across different domains.
5.3 COMPUTATION TIME
As mentioned before, one of the most important advantages of WEGL. compared to other graph rep- resentation learning methods its algorithmic efficiency. To evaluate that, we compare the wall-clock training and inference times of WEGL with those of GIN and the Wasserstein Weisfeiler-Lehman (WWL) graph kernel on five different TUD datasets (IMDB B, MUTAG, PTC MR, PROTEINS. and NCI1 For WEGL and WWL, -ngrlyy-rguularized version). We carry out our experiments for WEGL and WWL on 2.3 GHz Intel x E5-2670 v3 CPU, while we 16 GB NVIDIA @ Tesla P100 GPU for GIN. As reference, we implement GIN on the aforementioned CPU hardware as well.
the exact linear programming solver (as opposed to the
VigureBshoww how the training and inference run-times of WEGL,, WWL and GIN scale with the number of graphs in the dataset, the average number of nodes per graph, and the average number of edges per graph. As the figure illustrates, while having similar or even better performance, training WEGL is several orders of magnitude faster than WWL and GIN, especially for datasets with larger numbers of graphs. Note that training WWL becomes very inefficient as the number of graphs in the dataset increases due to the pairwise distance calculation between all the graphs.
During inference, WEGL is slightly slower than GIN, when implemented on a GPU. However, on most datasets, WEGL. is considerably faster in inference as compared to GIN implemented on CPU. Moreover, both algorithms are significantly faster than WWL. Using GPU-accelerated implementations of the diffuusin process in equation and and the entropy-regularized transport problem could potentially further enhance the computational efficiency WEGL during inference.
The complete classification results of WEGL on these datasets can he found in Appendix A.6
8

---

Published as a conference paper at ICLR 2021
PROTONS
Average training time (sec)
187
1800 2010
2018 J00 JS
Eeenge snbertepreee
(GFU)
2M (ON)
Average inference time per graph (90C)
TNE 2010 29/8
13
Figure 3: Average wall-clock time comparison of our proposed method, WEGI, with WWI TTgninalli al. 2019) and GIN (Xu et al. 2019). WEGL and WWL were implemented a 2.3 GHz Intel Xcon 3 E5-2670 CPU. Moreover, GIN was implemented twice, once using the aforementioned CPU and once using GB NVIDIA @ Tesla P100 GPU, with 5 hidden layers, 300 training epochs and batch size of 16 in both cases.
CONCLUSION
We considered the problem of graph property prediction and introduced the linear Wasserstein Embedding for Graph Learning, which we denoted as WEGL. Similar to (Togninalli et al. 2019). our approach also relies on measuring the Wasserstein distances between the node embeddings of graphs. Unlike (Togninali et al. 2019). however, we further embed the node embeddings of graphs into a Hilbert space, in which their Euclidean distance approximates their 2-Wasserstein distance. WEGL provides two significant benefits: i) it has linear complexity in the number of graphs (as opposed to the quadratic complexity of (Togninalli et al., 2019)). and ii) it enables the application of any ML algorithm of choice, such as random forest, gradient boosted decision tree, even AutoML. We demonstrated WEGL's superior performance and highly efficient training on wide range of benchmark datasets, including the ogbg mo lhiv dataset and the TUD graph classification tasks.
ACKNOWLEDGEMENT
We gratefully acknowledge funding by the United States Air Force under Contract No. FA8750-19 19- C-0098. Gustavo K. Rohde also acknowledges funding by NIH grant GM130825. Any opinions. findings, and conclusions or recommendations expressed in this material are those of the author(s) and not necessarily reflect the views of the United States Air Force and DARPA.
REFERENCES
Monica Agrawal, Marinka Zitnik, and Jure Leskovec. Large-scale analysis of disease pathways in the human interactome. In PSB, pp. 111-122. World Scientific, 2018.
Luigi Ambrosio, Nicola Gigli, and Giuseppe Savaré. Gradient flows: in metric spaces and in the space of probability measures. Springer Science & Business Media, 2008
Lars Backstrom and Jure Leskovec. Supervised random walks: predicting and recommending links
in social networks. In Proceedings of the Fourth ACM International Conference on Web Search and Data Mining, pp. 635-644, 2011.
Gary Bécigneul, Octavian- Eugen Ganea, Benson Chen, Regina Barzilay, and Tommi Jaakkola. Optimal transport graph neural networks. arXiv preprint arXiv:2006.04880 2020.
Aleksandar Bojchevski, Johannes Klicpera, Bryan Perozzi, Martin Blais, Amol Kapoor, Michal Lukasik, and Stephan Günnemann. Is pagerank all you need for scalable graph neural networks? In Proceedings of the 15th International Workshop on Mining and Learning with Graphs (MLG), 2019.

---

Published as a conference paper at ICLR 2021
Karsten M Borgwardt and Hans-Peter Kriegel. Shortest-path kernels graphs. In Fifth IEEE International Conference on Data Mining (ICDM'OS) IEEE, 2005.
Lco Breiman. Random forests. Machine learning, 45(1):5-32, 2001.
Yann Brenier. Polar factorization and monotone rearrangement of vector-valued functions. Commu- nications on pure and applied mathematics, 44(4):375-417, 1991.
Lars Buitinck, Gilles Louppe, Mathieu Blondel, Fabian Pedregosa, Andreas Mueller, Olivier Grisel, Vlad Niculae, Peter Prettenhofer, Alexandre Gramfort, Jaques Grobler, Robert Layton, Jake VanderPlas, Arnaud Joly, Brian Holt, and Gaël Varoquaux. API design for machine learning software: Experiences from the scikit-learn project. In ECML PKDD Workshop: Languages for Data Mining and Machine Learning, pp. 108-122, 2013.
Tianle Cai, Shengjie Luo, Keyulu Xu, Di He, Tie-yan Liu, and Liwei Wang. Graphnorm: A principled approach to accelerating graph neural network training. arXiv preprint arXiv:2009.03294 2020.
Ines Chami, Sami Abu-EI-Haija, Bryan Perozzi, Christopher Ré, and Kevin Murphy. Machine learning on graphs: A model and comprehensive taxonomy. arXiv preprint 20X5::2005.03675, 2020.
Nicolas Courty, Rémi Flamary, and Mélanic Ducoffe. Learning wasserstein embeddings. In Interna- fional Conference on Learning Representations, 2018.
Marco Cuturi and Arnaud Doucet. Fast computation of Wasserstein barycenters. In Proceedings of the 31 International Conference on Machine Learning, Proceedings of Machine Learning Research, pp. 685 693. PMLR, 2014. URL URLhttp:: httt:::/p/occeee/////// http:///aococcccnn.....rrrrrr/ep ht.m
Michaël Defferrard, Xavier Bresson, and Pierre Vandergheynst. Convolutional neural networks on graphs with fast localized spectral filtering. In Advances in Neural Information Processing Systems, pp. 3844-3852, 2016.
Yihe Dong and Will Sawin. COPT: Coordinated optimal transport on graphs. arXiv preprint arXiv:2003.00 038992 2020.
Simon S Du, Kangeheng Hou, Russ R Salakhutdinov, Barnabas Poezos, Ruosong Wang, and Keyulu Xu. Graph neural tangent kernel: Fusing graph neural networks with graph kernels. In Advances in Neural Information Processing Systems, pp. 5724-5734, 2019.
P Dvurechensky, A Gasnikov, and A Kroshnin. Computational optimal transport: Complexity by accelerated gradient descent is better than by sinkhorn's algorithm. In 35th International Conference on Machine Learning, ICML 2018, pp. 2196-2220, 2018
Federico Errica, Marco Podda, Davide Bacciu, and Alessio Micheli. A fair comparison of graph neural networks for graph classification. In International Conference on Leurning Representations, net forum? Hddrrnmmmid 6NFPB
2020. URL URL httpl/ttt//ppeeeea/aaaa
Montacer Essid and Justin Solomon. Quadratically regularized optimal transport on graphs. SIAM Journal on Scientific Computing, 4014):A1911-A1996 2018.
Matthias Feurer, Katharina Eggensperger, Stefan Falkner, Marius Lindauer, and Frank Hutter. Auto- skleam 2.0: The next generation. arXiv preprint arXiv:2007.04074 2020.
Matthias Fey and Jan E. Lenssen. Fast graph representation learning with PoTorch Geometric. In ICLR Workshop on Representation Learning on Graphs and Manifolds, 2019.
Matthias Fey, Jan-Gin Yuen, and Frank Weichert. Hierarchical inter-message passing for learning on molecular graphs. arXiv preprini arXiv:2006.1211 2020.
R'emi Flamary and Nicolas Courty. Pot python optimal transport library, 2017. URL https /onoth g git.....o/
Thomas Gärtner, Peter Flach, and Stefan Wrobel. On graph kernels: Hardness results and efficient alternatives. In Learning Theory and Kernel Machines, pp. 129-143. Springer, 2003.
10

---

Published as a conference paper at ICLR 2021
Will Hamilton, Zhitao Ying, and Jure Leskovec. Inductive representation learning on large graphs. In Advances in Neural Information Processing Systems. pp. 1024-1034 2017.
Danny Hernandez and Tom B Brown. Measuring the algorithmic efficiency of neural networks. arXiv preprint arXiv:2005. 00000000000035,, 2020.
Thomas Hofmann, Bernhard Schölkopf, and Alexander Smola. Kernel methods in machine learning. The Annals of Statistics, pp. 1171-1220 2008
Weihua Hu, Matthias Fey, Marinka Zitnik, Yuxiao Dong, Hongyu Ren, Bowen Liu, Michele Catasta, and Jure Leskovec. Open graph benchmark: Datasets for machine learning on graphs. arXiv preprint arXiv:2005.00687 2020.
Weihua Hu*, Bowen Liu*, Joseph Gomes, Marinka Zitnik. Percy Liang, Vijay Pande, and Jure Leskovec. Strategies for pre-training graph neural networks, In International Conference on Learn- ing Representations, 2020. htttss/t:///
/ttp:////peeeeiiieeeemmee / net. wow.ooiddd..llwiiii.ii..
Sergey Ivanov and Evgeny Burnaev. Anonymous walk embeddings. In International Conference on Machine Learning, pp. 2186 2195, 2018.
Wengong Jin, Connor Coley, Regina Barzilay, and Tommi Jaakkola. Predicting organic reaction outcomes with Weisfeiler-Lehman network. In Advances in Neural Information Processing Systems, pp. 2607-2616, 2017.
Hisashi Kashima, Koji Tsuda, and Akihiro Inokuchi. Marginalized kernels between labeled graphs. In Proceedings of the 20th International Conference on Machine Learning (ICML-03), pp. 321-328, 2003.
Kristian Kersting, Nils M. Kriege, Christopher Morris, Petra Mutzel, and Marion Neumann. Bench- mark data sets for graph kernels, 2020. URL http::/www: www.:rgraaphhleannnggg... httt::/www.rrappheeeerrre
io/
Thomas Kipf and Max Welling. Semi- Semi-supervised classification with graph convolutional networks. arXiv preprint a609v:6609.080 10000099977 2016.
Soheil Kolouri, Akif B Tosun, John A Ozolek, and Gustavo K Rohde. A continuous linear optimal transport approach for pattern analysis in image datasets. Pattern Recognition, 51:453-462, 2016.
Soheil Kolouri, Se Rim Park, Matthew Thorpe, Dejan Slepcev, and Gustavo K Rohde. Optimal mass transport: Signal processing and machine-learning applications. IEEE Signal Processing Magazine, 34(4):43-59, 2017.
Risi Kondor and Horace Pan. The multiscale Laplacian graph kernel. In Advances in Neural Information Processing Systems, pp. 2990-2998, 2016.
Nils Kriege, Marion Neumann, Kristian Kersting, and Petra Mutzel. Explicit versus implicit graph feature maps: A computational phase transition for walk kernels. In 2014 IEEE international conference on data mining, pp. 881-886. IEEE, 2014.
Nils M Kriege, Pierre-Louis Giscard, and Richard Wilson. On valid optimal assignment kernels and applications graph classification. In Advances in Neural Information Processing Systems, pp. 2016.
Nils M Kriege, Fredrik D Johansson, and Christopher Morris. A survey graph kemels. Applied Network Science, 5(1):1 2020.
Christian Léonard et al. Lazy random walks and optimal transport on graphs. The Annals of Probability, 44(3):1864-191 2016.
Guohao Li, Chenxin Xiong, Ali Thabet, and Bernard Ghanem. Deepergen: All you need to train deeper gens. arXiv preprint arXiv:2000607739,, 2020.
Hermina Petric Maretic, Mireille El Gheche, Giovanni Chierchia, and Pascal Frossard. GOT: An optimal transport framework for graph comparison. In Advances in Neural Information Processing Systems, pp. 13876-1 13887, 2019.
11

---

Published as a conference paper at ICLR 2021
Grégoire Mialon, Dexiong Chen, Alexandre d' Aspremont, and Julien Mairal. A trainable optimal transport embedding for feature aggregation and its relationship to attention. In International Conference Learning Representations, 2021.
iddz26vTvb844
Caroline Moosmüller and Alexander Cloninger. Linear optimal transport embedding: Provable fast wasserstein distance computation and classification for nonlinear problems. arXiv preprint arXiv:2008.08 1665, 2020.
Christopher Morris, Kristian Kersting, and Petra Mutzel. Glocalized Weisfeiler-Lehmall graph kernels: Global-local feature maps graphs. In 2017 IEEE International Conference on Data Mining (ICDM), pp- 327-336. IEEE, 2017.
Christopher Morris, Martin Ritzert, Matthias Fey, William L Hamilton, Jan Eric Lenssen, Gaurav Rattan, and Martin Grohe. Weisfeiler and leman go neural: Higher-order graph neural networks. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 33, pp.
2019.
Christopher Morris, Gaurav Rattan, and Petra Mutzel. Weisfeiler and leman go sparse: Towards scalable higher-order graph embeddings. Advances in Neural Information Processing Systems, 33, 2020.
K Muandet, Fukumizu, Sriperumbudur, and Schölkopf. Kernel mean embedding of distri- butions: A review and beyond. Foundations and Trends in Machine Learning, 10(1-2):1-144, 2017.
Navid Naderializadeh, Mark Eisen, and Alejandro Ribeiro. Wireless power control via counterfactual optimization of graph neural networks. arXiv preprint arXiv:2002.07631, 2020.
Mathias Niepert, Mohamed Ahmed, and Konstantin Kutzkov. Learning convolutional neural networks for graphs. In International Conference on Machine Learning, pp. 2014-2023, 2016.
Giannis Nikolentzos, Polykarpos Meladianos, and Michalis Vazirgiannis. Matching node embeddings for graph similarity. In Thirty-First AAAI Conference Artificial Intelligence, 2017.
Giannis Nikolentzos, Polykarpos Meladianos, Stratis Limnios, and Michalis Vazirgiannis. A degener- acy framework for graph similarity. In IJCAI, pp. 2595-2601, 2018
Bastian Rieck, Christian Bock, and Karsten Borgwardt. A persistent weisfeiler-lhhman procedure for graph classification. In International Conference on Machine Learning, pp. 9448-5458, 2019.
Hamidreza Sadreazami, Arash Mohammadi, Amir Asif, and Konstantinos N Plataniotis. Distributed- craph-based statistical approach for intrusion detection in cyber-physical systems. IEEE Transac- fions Signal and Information Processing Networks, 4(1):137-147,200
Vivien Seguy and Marco Cuturi. Principal geodesic analysis for probability measures under the optimal transport metric. In Advances in Neural Information Processing Systems, pp. 3312-3320, 2015.
Nino Shervashidze, SVN Vishwanathan, Tobias Petri, Kurt Mchlhorn, and Karsten Borgwardt. Efficient graphlet kernels for large graph comparison. In Artificial Intelligence and Statistics, pp. 498, 2009.
Nino Shervashidze, Pascal Schweitzer, Erik Jan Van Leeuwen, Kurt Mchlhorn, and Karsten M Borgwardt. Weisfeiler-Lchman graph kernels. Journal of Machine Learning Research, 12(77): 2539 2561, 2011.
Vayer Titouan, Nicolas Courty, Romain Tavenard, Chapel Laetitia, and Rémi Flamary. Opti- mal transport for structured data with application on graphs. In Proceedings of the 36th In- ternational Conference Machine Learning, volume 97 of Proceedings of Machine Learn- ing Research, pp. 6275-6284, Long Beach, California, USA, 09-15 Jun 2019. PMLR. URL
http:://proceddinne ngs. ml www.:mmwrr.../..mmmmmmmm
press wwwpss//w//ttttt///.
19 html]
12

---

Published as a conference paper at ICLR 2021
Matteo Togninalli, Elisabetta Ghisu, Felipe Llinares-Lppez, Bastian Rieck, and Karsten Borgwardt. Wasserstein Weisfeiler-Lehms graph kernels. In Advances in Neural Information Processing Systems, pp. 6436-6446, 2019.
Petar Veličković, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pietro Lio, and Yoshua Bengio. Graph attention networks. arXiv preprint 171.903700.00003, 2017.
Cédric Villani. Optimal transport: Old and new, volume 338. Springer Science & Business Media, 2008.
Wei Wang, Dejan Slepčev, Saurav Basu, John Ozolek, and Gustavo K Rohde. A linear optimal transportation framework for quantifying and visualizing variations in sets of images. International Journal of Computer Vision, 101()):254 269, 2013.
Zhang Xinyi and Lihui Chen. Capsule graph neural network. In International Conference on Learning Representations, 2019, URL/Ltttss//ppperevviie wytp://wwww.eee/orrrriiieggyyyaa 13BBRRcmmm
Keyulu Xu, Weihua Hu, Jure Leskovec, and Stefanie Jegelka. How powerful are graph neural
networks? In International Conference on Learning Representations, 2019. URL https
/oppnrevieww ntt/ feee//oommiiddrGGees A5Km
Pinar Yanardag and SVN Vishwanathan. Deep graph kernels. In Proceedings of the 21 ACM SIGKDD International Conference on Knowledge Discovery and Data Mining, pp. 1365-1374
2015.
Chi Zhang, Yujun Cai, Guosheng Lin, and Chunhua Shen. Deepemd: Few-shot image classifica- tion with differentiable earth mover's distance and structured classifiers. Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), June 2020.
Muhan Zhang, Zhicheng Cui, Marion Neumann, and Yixin Chen. An end-to-end deep learning architecture for graph classification. In Thirty-Second AAAI Conference on Artificial Intelligence, 2018a.
Zhen Zhang, Mianzi Wang, Yijian Xiang, Yan Huang, and Arye Nehorai. RetGK: Graph kernels based on return probabilities of random walks. In Advances in Neural Information Processing Systems, pp. 3964-3974, 2018b.
13

---

Published as a conference paper at ICLR 2021
A APPENDIX
Here we provide further details on the theoretical aspect of WEGL, our implementation details, and the sensitivity of the results to the choice of reference distribution. We also share implementation code on the ogbg-molhiv dataset to help with the review process.
A.1 DETAILED DISCUSSION ON LINEAR WASSERSTEIN EMBEDDING
The linear Wasserstein embedding used in WEGL is based on the Linear Optimal Transport (LOT) framework introduced in (Wang et al.) 2013). The main idea to compute the "projection" " of the manifold of probability measures to the tangent space at fixed reference measure. In particular, the tangent space at measure Ji is the set of vector fields VID {o 2 + Rd fz [=(=){=)uo(z)) < oc) such that the inner product the weighted lz:
(oi, Us)mo = mz(2) us())dpo(z).
(8)
We can then define v2(-) A f(-) id((), where /(-) is the optimal transport map from Mo to p4y. Note that 84 V Noo v0 0, and
|| l1:o = (wi, us) = J [15(z) -111 =1|(d)o(z)= W2(ps.Ho)
||0 Via-vollis voll2.
(9)
In the paper, use ((m,) es uty to turn the weighted-Ez into lz.
The discussion above assumes an absolutely continuous reference measure Po- A more interest- ing treatment of the problem is via the generalized geodesics defined in (Ambrosio et al. 2008), connecting μ4 and μú; and enabling us to use discrete reference measures. Following the notation in (Ambrosio et al. 2008), given the reference measure μlo let [[44, Mu) be the set of transport plans between 14 and μuu and let [(4 Hy. Au) be the set of all measures on the product space Z x x Z such that the marginals over μi, and μij are [[pj (o) and [(p,, po), respectively. Then the linearized optimal transport distance is defined as
inf Lo (14s (bj) TET(Mi-H)PP)) ExEx2
| 2' 12-2211/d7(22 221,d)((,,2',""
(10)
N.
N,
In discrete setting, where μ4
Σ
Sen Mj
Σ
82h and No 1/2 Σ 02g we have
No
N,
m=1
1=1
N: N
N
L (Pi; Mb))
min TET(Augg)o))
Tnml
En 2 112
(11)
N,N,N
n
See Figure 4 for depiction of Equation equation II meaning. Finally, the idea of barycenteric projection used to approximate Monge couplings and provide fixed-size representation is shown in F..
Yuml
23
[bjle
99
[p)]
Mi
μi
Pi
Hj
Reference
Reference
Ho
Ho
Figure 4: Illustration of (a) the meaning behind "ooo used in the LOT distance in equation 11 and (b) the idea of the barycenteric projection, which provides fixed-sin representation (i.c., of size N).
14

---

Published as a conference paper at ICLR 2021
[)]
R-11
Optimal Transport
Approsimated Menge Map
Embedded Distribution, Kill
Embodding ng
Retertron
f
Banycenteric Projection
CCC))
Eucidean Goodasic the Embudding Space fromz,
(d)
Fuciddean Geodesic Fimberding Spucee from In
k=0.2
um0.30
x-1 0o
0.80
(osiall olaing)) (o
Nopipp
Figure 5: An experiment demonstrating the capability of the linear Wasserstein embedding, (a) A simple dataset consisting of shifted and sealed noisy ring distributions J. where we only observe samples
Z. 1&t puliz1 from each distribution, together with the process of obtaining the linear Wasserstein embedding with respect to reference distribution. short, for each distribution ps. the embedding approximates the Monge-map (i.c., a vector field) from the reference samples Zo to the target samples Z: by a barycentric projection of the optimal transport plan. Adding samples in the embedding space corresponds to adding their vector fields, which can be used to calculate (b) the mean distribution in the embedding space, i.e., $ M IN ((p.) and (c)-(d) the Euclidean geodesics in this space, i.c., & (aa((p))) + (1 - a)o(p))) for E [0, 1]. As can be seen, the calculated mean is the Wasserstein barycenter of the datasct, and the Euclidean geodesies in the embedding space follow the Wasserstein geodesies in the original space.
Next, to demonstrate the capability of the linear Wasserstein embedding, we present the following experiment. Consider a set of distributions {pr } where each p. is translated and dilated ring distribution in R², and N, samples
observed from pt, where N, and Nj could be different for i * j. We then consider a normal distribution as the reference distribution and calculate the linear Wasserstein embedding with respect to the reference (See Figure55a, Given the -ever-iieee nature of the embedding, to demonstrate the modeling capability of the framework, we calculate the mean in the embedding space (i.e., the vector fields), and invert obtain the mean distribution p. Figure 5b shows the calculated mean, indicating that the linear Wasserstein embedding framework has successfully retrieved a ring distribution as the mean. Finally, we calculate Euclidean geodesics in the embedding space (i.c., the convex combination of the vector fields) between p; and po, as well as between p; and Pj. and show the inverted geodesics Figures 5 and respectively. As the figures demonstrate, the calculated geodesics follow the Wasserstein geodesics.
APPROXIMATION ERROR OF THE EMBEDDING
Given the Euclidean distance in the embedding space is a transport-based distance (i.e., the so- called LOT distance) that approximates the Wasserstein distance, natural question arises about the approximation error. Here we point the reader to the recent work of Moosmüllerr & Cloninger (2020) in which the authors provide bounds on how well the Euclidean distance in the embedding space approximates the 2- Wasserstein distance. In particular, the authors show that:
W2(p,, pb)) & [((())) 0(1)) S W2(4s, H4) + 11.f/ ffj o 1
where fH: is the optimal transport map from μ to μ.. This inequality simply indicates that the approximation enror is caused by conditioning the transport map be obtained by composition of the optimal transport maps from 14 to Po, and then from (An to μıy More importantly, it can be shown that if μ4 and μ4 are shifted and scaled versions of the reference measure, Mu, then the embedding is isometric (See Figure and 2 in Moosmüller & Cloninger (2020).
Maybe less interesting upper bound can also be obtained by the triangle inequality:
W2 (b) < [(²(44) ((1)))12 & Wa(ps Wa(a, which ensures some regularity of the embedding.
15

---

Published as a conference paper at ICLR 2021
REGULARITY THE EMBEDDING
good question was raised during the feedback period, about the regularity of the proposed embed- ding. The regularity of the graph embedding will depend on both the regularity of node-embedding and the Wasserstein embedding. In the following, we avoid the discussion on regularity of the node- embedding (as it is not the main focus of our work), and focus on the regularity of the Wasserstein embedding. To that end, we first point out several regularity characteristics pointed out in Appendix A of Moosmüllee & Cloninger (2020). Most notably Theorem 4.2 in their paper, shows an almost isometric property when the distortions are within an e-tube around the set of shifts and scalings. In short, let (2( R > 0, > 0, E be the set of all shifts and scalings, and
En.n = {h E & || - R),
and let
- (2, 3h E Ep.R l1g h|lp S ct, which is the c-tube around set of shifts and scalings. Now, assume 140 E P2(22 is the reference measure and both μ and μn satisfy Caffarelli's regularity theorem. Then for 91-92 E 9u, R.e we have
0 S (1 (( 11)22 W2(91 #/: 92 ()) S Cpo.p. + nE.
where C. and Cho.p. are constants depending on Nu, ,4, and R.
The results shown above can be used to derive regularity results for the linear Wasserstein embedding with respect to additive noise. We know that the addition of two random variables leads to new random variable with its PDF being the convolution of the original PDFs. Therefore, for features
Zi = [.,, JN, let 2: = 12 NN:
T denote the noisy samples for e; ~ 7), where
EN.
EN1
the noise distribution. Then the noisy samples, 2, will be distributed according to i. μ4 * 17. For instance, for the Gaussian additive noise, fa is the smoothed version of μ4. Therefore, there exists a transport map in g Iy,,n. for which, jág I# b4 and [gg id lp; -W2(it,)s)) S and have:
[(δ(p)) p(p)) ll2 S (Cpa,,ts + 1)e + Cpo.n.Ree
A.2 SENSITIVITY THE CHOICE OF REFERENCE DISTRIBUTION
To measure the dependency of WEGL on the reference distribution choice, we changed the reference normal distribution (i.e., data-independent). We compared the results of WEGL using the new reference distribution to that using reference distribution calculated via k-means on the training set. We used the ogbg-mm mo hiv dataset with initial node embedding of size 300 and diffusion layers. We ran the experiment with 100 different random seeds, and measured the test ROC-AUC of WEGL calculated with the two aforemeentioned reference distributions. Figure shows the results this experiment, indicating that the choice of reference distribution is statistically insignificant.
p-value: 0.05
k-means, 76.86 = 1.44
12
Normal, 76.45 1.47
10
72
74
76
78
80
ROC-AUC (%)
Figure 6: ROC-AUC (%) results on sobii datasct, when the reference distribution is calculated by k-means (Section 4.2) on the training dataset (denoted k--maass, compared to when it is fixed to be normal distribution (denoted Normal). With p-value= 0.05, the choice of the template is statistically insignificant.
16

---

Published as a conference paper at ICLR 2021
OGBG-molhiv, (Avg. #Nodes 25.5)
120
Relative Error
Csinkhom/tis
80
100
60
75
50
25
-0
10
2.5
3.0
35
4,0
Log of Regularizer
Figure 7: Comparing the performance of the linear programming (LP) solver with the Sinkhorn algorithm (on entropy regularized OT problem) on the ogbg-molht holiiv dataset for various regularization parameters,
A.3 LINEAR PROGRAMMING VS. ENTROPY REGULARIZATION
In this paper, used the Python Optimal Transport (Flamary & Courty 2017) for the calculation of the optimal transport plans. During the feedback period, a point came regarding the entropy- regularized version of the OT problem (Cuturi & Doucet 2014), which reduces the complexity of the linear programming problem from being cubic, in the number of nodes, to being quadratic, using the Sinkhorn algorithm. Given the Sinkhorn algorithm's 's iterative nature, the computational gain of the method is prominent when calculating the transportation problem between graphs with large number of nodes (e.g., larger than 10a). However, the graph datasets used in this paper often have small number of nodes (e.g., < 50). In these settings, linear programming is efficient. To obtain any computational gain using the Sinkhorn algorithm, one would need to use large regularization coefficient, which reduces the Sinkhorn algorithm's 's precision.
Here we ran an experiment on the ogoog-mo ivv dataset. We obtain the transport plans between the node embeddings and the reference distribution using the linear programming solver (using ct emd2 from FFlamary& & Courty 2017)) and the Sinkhorn algorithm for the entropy-regularizer problem (using ot s sinkhorn2 from (Flamary & Courty| 2017)). We measure the calculation time well as the calculated distances for both algorithms. For the Sinkhorn algorithm, we used four different regularization values. We report the mean and standard deviation of calculation time ratio, i.e., tainkhern and the relative error of calculating the 2- Wasserstein distance, i.e., Worthhhroo--W.... W2.2 in Figure 7 As the figure shows, due to the small graph size, the LP solver is efficient and little gain be obtained using the Sinkhorn algorithm. Nevertheless, in the case of dealing with larger graph sizes, the entropy regularized formulation should be the definite choice. Finally, for the eegulpe-re OT problem, more efficient solvers have been proposed that outperform the Sinkhorn algorithm (Dvurechensky et al.) 2018). However, given the acceptable performance of the linear programming solver (at least for the graph datasets in this paper), we did not find it necessary to seek more efficient solvers.
fge
A.4 INNER WORKING OF GNNS
In its most general form, GNN consists of hidden layers, where at the 4th layer, each node v E aggregates and combines messages from its 1-hop neighboring nodes No, resulting in the feature vector
V combine
(o-l- (-1))
vI E (1...., LJ, YOEV, (xt (af((-1), Www) )
)
(12)
where V (.m denotes parametrized and differentiable combining function.
At the input layer, each node v E starts with its initial feature vector zQ. = z. E RF, and the sequential application of GNN layers, as in equation 12 computes intermediate feature vectors
(1) 1 H1 At the GNN output, the feature vectors of all nodes from all layers go through global pooling (i.e., readout) function V readour) (.), resulting in the final graph embedding
(b( (G) = readout (fa5))
JoEV.E(O.....))
(13)
17

---

Published as a conference paper at ICLR 2021
A.5 IMPLEMENTATION DETAILS
To derive the node embeddings, we use the diffusion process in equation for the datasets without edge features/labels, i.e., all the social network datasets REDDIT-BINARE
REDD IT MT-MULTT- 5K, and SEDDIITMUULTI-222 IMDB-MULTH COLLAB, and four of the molecule and bioinformaties datasets (ENZYMES, PROTEINS, D&D, and NCI1). We specifically set war = 1 for any (u, v) E & and also for all self-connections, i... Wiv = 1, Ve e V.
The remaining datasets contain edge labels that cannot be directly used with equation Specifically, each edge in the MUTAG and PTC-MR datasets has a categorical label, encoded as one-hot vector of dimension four. Moreover, in the indicating bond type (five categories), bond stereochemistry (six categories) and whether the bond is conjugated (two categories). We first convert each categorical feature to its one-hot representation, and then concatenate them together, resulting in binary 13-dimensional feature vector for each edge.
dataset, each edge has three categorical features
In each of the three aforementioned datasets, for any edge (u, v) E &, let us denote its binary feature vector by Whe E {0, 1 where E'is is equal to 4, 4, and 13 for MUTAG, PTC-MP and ogbo a--moole molhiv, respectively. We then use the following extension of the diffusion process in equation 7
E
Wuo.e
zu (1-1), V (1,..,,,, L}, Vo EV,
(14)
deg. (u)deg.
WEV
where for any e E {I,..., Wav.e denotes the ethh element of Wan, and for any node E V, we define (v) as its degree over the ethh elements of the edge features; i.e., deg..r) Σ Eney Wun.es We assign vectors of all-one features to the self-connections in the graph; i.e., Wyz.c 1, Ve E V, Ve E {l...., E` Note that the formulation of the diffusion process in equation 14 can be seen as an extension of equation 7 where the underlying graph with multi-dimensiona edge features is broken into E' parallel graphs with non-negative migle-dimensionall edge features, and the parallel graphs perform message passing at each round/layer of the diffusion process.
For the ogbg-molhing experiments in which virtual nodes were appended to the original molecule graphs, we set the initial feature vectors of all virtual nodes to all-zero vectors. Moreover, for any graph G; in the dataset with [V]] nodes, we set the edge features for the edge between the virtual node Uvirtual and each of the original graph nodes u E V/ as Wan
[ Ve E (1...., F). The normalization by the number of graph nodes is included so as to regulate the degree of the virtual node used equation [4] We also include the resultant embedding of the virtual node at the end of the diffusion process in the calculation of the graph embedding (G.).
In the experiments conducted on each dataset, once the node embeddings are derived from the diffusion process, we standardize them by subtracting the mean embedding and dividing by the standard deviation the embeddings, where the statistics are calculated based on all the graphs in the dataset. Moreover, to reduce the computational complexity of estimating the graph embeddings for the ogog-mo 1 hiv dataset, we further apply a 20-dimensional PCA on the node embeddings.
HYPERPARAMETERS
We use the following set hyperparameters perform a grid search over in each of the experiments:
Random Forest: min sai__amples__ee leaf E (1,2,5}, 5). dir_samplees_ppl5 {2 10), and n n_estimators E {25, 50, 100, 150, 200).
Gradient Boosted Decision Tree (GBDT): min samples leaf 5]. min samp_saplees.. ssplitttttt____ssep_peeepppsittt E {2, 52 10), n dd_estimator E [25, 50, 100, 150, 200}. and max_depth (1,3,5).
SVM-Limear and SVM-RBF: C (10-2 ..... 109).
Multi-Layer (MLP):
sizes {(12)), (256), (128, 64), (256, 128)].
Auto-ML:
2.0 searches over a space of 42 hyperparameters using Bayesian optimization techniques, as mentioned in Feurer et al. (2020).
Number of Diffusion Layers in equation and equation 14 I. E (3.....8).
18

---

Published as a conference paper at ICLR 2021
Initial Node Feature Dimensionality (for ogbg-molhiv only): {100, 300, 500).
Node Embedding Type: For a graph with F-dimensional initial node features, we consider using the following three types of node embedding:
Concat:
Ze
xv 10)
2 ... x000 (R(+I)F where denotes concatenation.
Average: = I+T Ei=o Tt E ER..
1
Final: 2u = Zt E I.F.
(L)
A.6 TUD BENCHMARK - COMPLETE RESULTS
Here, we report the comprehensive of classification results of our proposed method, WEGL, for each of the node embedding types mentioned in Section A.5
using five different classifiers: Linear
SVM, Kernel-SVM (SVM-RBF), Gradient Boosted Decision Trees (GBDT), Multi-Layer Perceptron
(MLP), and Random Forest (RF). The results are shown in Table 3
19

---

Published as a conference paper at ICLR 2021
NCI
58.9 - 7.2
7.2
5
174+50 5.0
67.8 1 2.1
65.1
64.9 - 2.3 2.3
78.4 1.6
78.2 2.9
7:
4:4:2.2 2.5
61.3 4.6
60.0
59.9 3.7
7..8
76.6
75.9
Tahle 3: Graph classification uccuracy (%) of the Wasserstein embedding hased on different classifiers and the three types, of node embeddings, on vatiios TUD graph classification tasks.
57.4
76.4
61.3
MUTAG
74.1 2..1 83.6 9.5
8.
5.4
34.0
83.0
82.4
89.3
10.0
8.2
8..
81.9
88.315.1
86.245..8
86.249.5
87.8
87.2
84.0
PROTEINS
7.52
73.1 2.11
1.4
7.47
75.8
7..3
2.7
2.
72.5
3.0
4,1
76.113..3
75.8+3..6
76.5+4.2
72.8
76.0
72.8
ENZYMES
50.0 5.8
39.3 7.7
37.5 5.5
4.2
5.0
5.
6.3
5.6
5.4
55.7
5.2
60.5
58.7+6.9
57.5=5.5
58.5
60.0
57.7
57.3
50.8
57.8
55.7
55.5
PTC-MR
61.6
62.84.99
62.2
6.6
1
5.8
5.3
6288
5.4
63.4
65.2
66.2
6.9
60.2
6.11
61.6 6.3
64.6.17.4
67.517.7
66.31..5
RE-M5K
2.5
92.1 19 56.1 2.3
5 1.6
1.6
55.1
53.342.0
13.
91.3 1.8
3.1
80.5 2.6
92.9 1.9
88.8 1.9
3.69
78.6 2.
92.011.88
87.8+3.0
87.5+20
RE-B
COLLAB
72.7
1.3
2.6
78.6
80.6 220
73.1
79.8
78.8+144
79.1+1.1
75.8
IMDB-M
69.1 1.1 41.
1.8
5.
6.3
3.1
2.7
3.11
51.9
2.9
4.0
3.1
3.7
3.8
50.814.0
52.0+4.1
51.744.6
IMDB-B
72.7 5.5
72.5 3.7
71.9 1 2.4
73.4 73.4 2.5
72.8 3.8
74.5 # 4.2
75.2 5.0
75.1 2.0
72.0 2 3.5
72.0 4.0
72.6 72.6 3.5
74.916.3
74.415.6
75.4+5.0
Classifier Embedding Type
Concat
Avg
Final
Concat
Avg
Final
Concat
Avg
Final
Concat
Avg
Final
Concat
Avg
Final
SVM Linear
SVM RBF
GBDT
MLP
RF
VEGI.
20


Reviews:
[{"review_id": "wTJx0HGUW-l", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "##########################################################################\n\nSummary:\n\nThis paper proposes to use a Wasserstein embedding to compare graph using also node embedding methods to convert graphs into vectors. While the elements are not new, the proposed method is new and is competitive with other classical methods (graph kernel and graph neural networks). \n\n##########################################################################\n\nReasons for score:\n\nOverall, I vote for mild acceptation. The method is clearly new, but the ingredients are already new and the results are not godd but not exiting.\n\n##########################################################################\n\nPros:\n\n1. The combination of node embedding with Wasserstein distance is new and show pretty good results.\n\n2. The method itself seems very flexible and effective compare to other state of art methods.\n\n3. The presentation of the methods is rather clear and most of the technical details are available. As such the paper is self-contained.\n\n##########################################################################\n\nCons:\n\n1. The framework relies on a node embedding method. Seems the choice of the method could be crucial, we need some discussion on this point.\n\n2. The proposed node embedding formula (equation 9) assumes that the information on the edge is a scalar. However previously in the paper the edge attributes are a vector, please give some slues on how to deal with such edges.\n\n3. The section 4.2 is not clear for me. Is the zero distribution compute on the whole database? Or is it compute class by class? Both ways seems possible please clarify?\n \n##########################################################################\n\nQuestions during rebuttal period:\n\nPlease address and clarify the cons above.\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "Ha5S7As4WDV", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "\nThe paper presents an embedding method for graphs, denoted as WEGL. The method operates by first representing each graph as a set of vector representations of its nodes, secondly, an optimal transport map is computed with respect to a finite reference set of points and, finally, a fixed-size vector representation is computed from that map.\nThe result is an embedding into a Hilbert space, where the distance between two graph representations approximates the 2-Wasserstein distance between their respective node distributions.\n\nThe paper is partly inspired by WWL (Togninalli et al. 2019) and provides an explicit embedding map, hence improving in terms of computational performance. I have some comments and concerns about the experimental analysis, which does not seem to directly validate the paper claims. Overall, it is well written and clear in almost all its parts.\n\nComments.\n1. The paper considers only the smallest of the datasets in OGB (Hu et al., 2020) for graph-level prediction, specifically ogbg-molhiv. I wonder if the proposed method, which is claimed to be computationally efficient, can handle also the other three.\n2. Choice of the baseline methods. I expected comparisons with other (explicit) embedding methods (eg, Hu* et al. 2020, Kriege et al. 2014) both computationally and at the task. Since WEGL is somehow a variation of WWL, I would like to see its performance compared also in Section 5.1.\n3. How performance has been assessed. The paper mention in the captions of Tables 1 and 2 that some of the results have been reported from already published papers. However, I couldn't find any details about the experimental settings of the other methods. Is the architecture of the GIN in Figure 3 the same as that in Tables 1 and 2? Are all the baseline methods operating on graphs with the \"virtual node\" variant (which is important for a fair comparison of the embedding capability)?\n\nTo further improve the paper, I suggest spending a few more words to\n- clarify the use of the Jacobian and the approximation in points 1 and 3, Section 3.1;\n- define the barycentric project and clarify why it is not invertible.\n\n(Kriege et al. 2014) Kriege, Nils, et al. \"Explicit versus implicit graph feature maps: A computational phase transition for walk kernels.\" 2014 IEEE international conference on data mining. IEEE, 2014.", "extra_scores": {}, "extra_text": {}}, {"review_id": "1LoYVY_q6ch", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Strengths:\n* **Computational advantage**. The quadratic-to-linear complexity reduction for pairwise computation (each of them an LP) of LOT is clearly the main selling point of this paper, which makes it an appealing approach for comparing multiple graphs. \n* **Writing**. The paper is overall very clear, well-written, and engaging. \n\nWeaknesses:\n* **Novelty**. The paper proposes a farily trivial combination of Wang et al's 'linear' OT with simple message passing on graphs. The main aspect where there could have been an interesting novel contribution, the computation of the reference embedding, is solved somewhat unsatisfactorily with a k-means embedding\n* **Unexplained design choices**. The use of the 'virtual node' is not well motivated, nor its contribution evaluatued experimentally. In addition, it's not clear why it was decided to not implement/comparse the entropy-regularized version of the proposed method, given that this could hold the key to truly scalable graph comparison. In addition, entropy regularization has consistently shown to even *improve* performance compared to exact OT in many applications. Most OT libraries already include this variant, so there really seems to be no excuse not to incoporate it. \n* **Some important trade-offs brushed under the rug**. Equation (13) seems to suggest that a NxNxN tensor has to be stored in memory, which could be prohibitive in many settings. Is it fair to say that LOT is implicitly trading time by space complexity? In fact, one could argue that this not even entirely correct, as the LP solved in (13) is now cubic rather than quadratic on N, so it seems there is a trade-off in computation too: reducing the number of pairwise comparisons but making each individual LP problem larger. \n* **Experimental results raise questions**. There are many moving parts in the evaluation (e.g., the addition of the virtual node, different types of classifiers), so it's hard to disentangle the core effect of the proposed method. \n    * The results in Table 1 seem to indicate that the classifier at the end of the pipeline has a much stronger effect on performance than the proposed method. \n    * What kind of classifiers are used in the other GNN methods? Where these also optimized for performance? Why not using the same type of classifier on all methods for fair comparison?\n    * All the baseline methods have a very strong degradation in performance between validation and test sets. For the proposed method, the degradation is much less signigicant, especially for the +AutoML classifier. How do the authors explain this? This seems to indicate that there is significant overfitting in all methods, and that proper regularization of the classifiers (e.g., via AutoML) is mitigating this for WEGL. This casts significant doubts on the results of this section in my opinion.\n    * Again, for the results in Section 2, the classifier seems to have an important effect in the performance of the proposed methods, yet there is no discussion on the classifiers used in all the baselines. In addition, the standard errors are so large that pretty much all methods have overlapping confiedence intervals, so it's hard to draw a statistically significant conclusion from these results. \n    * The runtime comparison is also limited. First, it's not clear whether the runtimes shown for WEGL include the full pipeline (e.g., including the reference embedding computation). Second, one of the methods is run on GPU while the other two are run on CPU, so it's hard to put those two sets of results in direct comparison. \n\n* **Misses importart related work**. E.g., \n    * [1] attempts to realize Characteristic (4) mentioned in Section 3\n    * [2] also leverages generalized geodesic for comparison of multiple measures\n    * [3] combines OT and WL Kernels and applies it to settings very similar to the ones tackled here. \n\nOther comments questions:\n* It would be informative to see a discussion on how well LOT enfoces Characteristic 4, that is, what guarantees does one have on how close the l2 embedded distances is to the true wasserstein distance?\n\nMissing references:\n* [1] Courty et al., \"Learning Wasserstein Embeddings\", 2017\n* [2] Seguy and Cuturi, \"Principal Geodesic Analysis for Probability Measures under the Optimal Transport Metric\", NeurIPS 2015. \n* [3] Bécigneul et al., \"Optimal Transport Graph Neural Networks\", 2020", "extra_scores": {}, "extra_text": {}}, {"review_id": "-PlwgALr2tj", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "# Synopsis of the paper\n\nThis paper presents WEGL (\"Wasserstein Embedding for Graph Learning\"),\na technique for embedding graphs for graph classification. The primary\nnovelty of the paper lies in their smart choice of embedding, which\ncombines the expressivity of optimal transport paradigms (i.e. the\nWasserstein distance, in this case) with improved computational\nperformance; more precisely, the embeddings permit the use of effective\nclassification algorithms, as opposed to scaling quadratically with the\nnumber of samples.\n\n# Summary of the review\n\nThis is a well-written paper with a powerful algorithm and a strong\nexperimental section. It will make an excellent contribution to the\nconference, and I am excited to endorse it!\n\nThere are some issues with the current version of the write-up, though,\nwhich, if fixed, will make this an even stronger publication. My primary\nconcerns at present are:\n\n1. The paper spends a lot of space in outlining background information\n   that is not pertinent to the method. I appreciate the amount of\n   details that are being packed into Section 2.1, but as a reader of\n   this paper, I would prefer a more in-depth discussion of the method\n   rather than a discussion of GNNs, which are used as mere comparison\n   partners here. My suggestion would be to put some of these\n   information into the appendix.\n\n2. The section on Wasserstein distances could be streamlined. At\n   present, too many different concepts are introduced; I feel that\n   a non-expert reader will just be deterred, even though the remainder\n   of the paper is very hands-on. Maybe some additional intuition could\n   be provided here?\n\n3. On the other hand, Section 3 is *missing* important information; in\n   this section, I would be very much interested in knowing more about\n   the theoretical properties of the method (and its empirical\n   performance). The characteristics mentioned on p. 4 are intriguing,\n   but it would improve the paper if a more detailed write-up would be\n   provided; a reference with proofs for these properties would be\n   equally helpful. I am fully aware that it is hard to satisfy both\n   theory and experiments in a paper; I have some more detailed comments\n   about what could be added (potentially in the supplementary materials,\n   if the authors think that it detracts from the flow).\n\n# Detailed Comments\n\n- The point about the 'true metric' in the introduction is slightly\n  ambiguous. My understanding is that one obtains a metric between the\n  embedded feature vectors. This does *not* constitute a metric on the\n  graphs, though, unless the embedding is injective. I think the\n  sentence means to say that one obtains a metric in an embedding space\n  and this metric serves as a proxy for the Wasserstein distance.\n\n- In the related work section, I would suggest citing other variants of\n  the WL algorithm that have recently emerged:\n\n    - Morris et al.: *Weisfeiler and Leman Go Neural: Higher-Order Graph Neural Networks* (https://aaai.org/ojs/index.php/AAAI/article/view/4384) \n    - Morris et al.: *Weisfeiler and Leman go sparse: Towards scalable higher-order graph embeddings* (https://arxiv.org/abs/1904.01543)\n    - Rieck et al.: *A Persistent Weisfeiler-Lehman Procedure for Graph Classification* (http://proceedings.mlr.press/v97/rieck19a.html)\n\n- The way a graph is defined in this paper could be misconstrued as\n  *directed* upon first reading. Why not define edges in both directions\n  by using a subset notation instead of a tuple notation?\n\n- How are categorical attributes treated? In the supp. mat., the usual\n  one-hot encoding is mentioned. This could be discussed earlier (i.e.\n  on p. 2).\n\n- When discussing the embedding, I would suggest briefly mentioning the\n  concept of reproducing kernel Hilbert spaces (RKHS). Such a discussion\n  will provide the relevant backdrop for this publication.\n\n- As suggested above, I would shorten Section 2.1 somewhat (or put some\n  of the text into the appendix). While it is good to provide some more\n  details about the inner workings of GNNS, this is not required for the\n  paper.\n\n- Section 2.2 would benefit from more intuition; a lot of the\n  terminology is introduced too tersely and not re-used. I fully\n  appreciate that the paper is mathematically precise here, but at the\n  same time, I would suggest a more streamlined introduction of the\n  concepts that are necessary to define the embedding in the end.\n\n- Section 3.1 and Figure 1 differ slightly in terms of their notation.\n  I would suggest to harmonise the description here in order to be more\n  consistent.\n\n- For property 4 on p. 4, is it possible to quantify the strength of the\n  approximation? How good is this approximation in practice, and what\n  are the factors influencing it? I would suggest adding some more\n  details here (or citations, if appropriate).\n\n- To add to the previous point, I would in general like to know more\n  about the stability properties of the full embedding. Can this easily\n  be quantified? For example, what happens if I add some noise to the\n  attributes? I would assume that stability is a function of the\n  selected aggregation function *and* the Wasserstein embedding itself.\n  Assuming that the former is fixed to be the function described by\n  Togninalli et al., does the latter satisfy, for example, Lipschitz\n  continuity? (I am asking this out of professional curiosity;\n  understanding the inherent properties of the embedding seems key to me\n  for us to understand better ways to generate such embeddings; the\n  empirical performance of the proposed method speaks for itself, of\n  course!)\n\n- Section 3.2 is rather technical at present, making use of\n  hitherto-undefined concepts. I would suggest improving this section\n  by providing more intuition, relegating some of the more technical\n  results to the appendix.\n\n- Does the pseudo-invertibility of $\\phi$ cause any problems in\n  practice? It is my understanding that the embedding will not be\n  injective in general anyway, or am I mistaken?\n\n- How stable is the reference embedding? In Section 4.2, it is my\n  understanding that the reference embedding is by default obtained from\n  $k$-means, with $k$ being the average number of nodes. The appendix\n  depicts the results for another data-independent reference embedding\n  and states that there are no differences. Would this not suggest that\n  the way the reference distribution is obtained is irrelevant? If so,\n  why not use a normal distribution (which I would expect to be simpler\n  to calculate than a $k$-means embedding) for all experiments?\n\n  This point should be addressed somewhat better.\n\n- Why is the virtual node inclusion necessary? Only for the\n  simplification of the message passing itself? I am aware of this\n  standard modification, but I do not see why the proposed algorithm\n  could equally well work with the original graph.\n\n- Why is the variance of the proposed method high for some of the data\n  sets of the 'TUD' repository? Is this a consequence of the model that\n  was picked for working with the embeddings?\n\n# Style\n\nThis paper is well written; it was a pleasure to read and review. Here\nare some minor suggestions to improve style/clarity:\n\n- I would suggest so sort citations by some criterion (alphabetically,\n  for example) when citing multiple authors.\n\n- \"See the recent survey\" --> \"see the recent survey\"\n\n- \"such transport plan\" --> \"such a transport plan\"", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper proposes  a novel and interesting embedding of graphs emulating the Wasserstein distance. The experiments are good and the authors did a detailed answer taking into account the comments of the reviewer. The responses were appreciated and the AC recommends the paper to be accepted.
Decision: Accept (Poster)