WILTing Trees: Interpreting the Distance Between MPNN Embeddings
Authors: Masahiro Negishi
Thomas Gärtner
Pascal Welke
Venue: ICML 2025 poster (2025)
Keywords: Weisfeiler Leman test
Graph Neural Networks
Interpretability
Graph metric
Graph distance
Primary area: deep_learning->graph_neural_networks
TL;DR: We show that message passing neural networks (MPNNs) are implicitly trained to respect graph functional distances, and introduce the weighted Weisfeiler Leman Labeling Tree (WILT) to identify subgraphs that MPNNs consider functionally important.

Abstract: We investigate the distance function learned by message passing neural networks (MPNNs) in specific tasks, aiming to capture the _functional_ distance between prediction targets that MPNNs implicitly learn. 
This contrasts with previous work, which links MPNN distances on arbitrary tasks to _structural_ distances on graphs that ignore task-specific information. 
To address this gap, we distill the distance between MPNN embeddings into an interpretable graph distance. 
Our method uses optimal transport on the Weisfeiler Leman Labeling Tree (WILT), where the edge weights reveal subgraphs that strongly influence the distance between embeddings. 
This approach generalizes two well-known graph kernels and can be computed in linear time. 
Through extensive experiments, we demonstrate that MPNNs define the relative position of embeddings by focusing on a small set of subgraphs that are known to be functionally important in the domain.

WILTing Trees: Interpreting the Distance Between MPNN Embeddings
Masahiro Negishi Thomas Gärtner' Pascal Welke
Abstract
learned. That line of work has derived generalization bounds under strong assumptions the margin between classes on the Lipschitz constants of MPNNs, which often do not hold in practice. In this work, we instead investigate the distance dMNNN implicitly obtained from a popular MPNN trained on a practical dataset.
We investigate the distance function learned by message passing neural networks (MPNNs) in specific tasks, aiming to capture the functional distance between prediction targets that MPNNS implicitly learn. This contrasts with previous work, which links MPNN distances on arbitrary tasks structural distances on graphs that ignore task-specifie information. To address this gap, distill the distance between MPNN embeddings into an interpretable graph distance. Our method uses optimal transport on the Weisfeiler Leman Labeling Tree (WILT), where the edge weights reveal subgraphs that strongly influence the dis- tance between embeddings. This approach gen- eralizes two well-known graph kernels and can be computed in linear time. Through extensive experiments, we demonstrate that MPNNs define the relative position of embeddings by focusing on a small set of subgraphs that are known to be functionally important in the domain.
Specifically, we ask: What properties does the distance duPwN learned by well-performing MPNN have in practice that can explain its high performance? ? While previous stud- ies (Chuang & Jegelka, 2022; Böker et al., 2024) focused on the alignment between dMPNN and non-task- taillrrd structural graph distance datnee we have found that not critical to the predictive performance of MPNNs. Rather, even if an MPNN was trained with classical cross-entropy loss, dMPNN respects the task-relevant functional distance diune. Furthermore, the alignment between dMPNN and drunc is highly correlated with the predictive performance of the MPNN. Then, we move to our second question: How do MPNNs learn such a metric structure? Since dMPNN is es- sentially a distance between multisets Weisfeiler Leman (WI.) subgraphs, we distill dMPNN into more interpretable distance between these multisets. Our distance dwnL an optimal transport distance on a weighted Weisfeiler Leman Labeling Tree (WILTT), which is a trainable generalization of the graph distances of existing high-performance ker- nels (Kriege et al., 2016; Togninalli et al., 2019). It allows us to identify WL subgraphs whose presence or absence significantly affects the relative position of graphs in the MPNN embedding space. In addition, dwIr is efficiently computable the ground metric is a path length and is at least as expressive dMPNN in terms of binary expressive power We show experimentally that the WILT- ing tree distances fit MPNN distances well. Examination of the resulting edge parameters on a WILT after distillation shows that only a small number of WL subgraphs deter- mine dMPNN. In qualitative experiment, the subgraphs that strongly influence dMPNN those that known to be functionally important by domain knowledge. In short, our contributions are as follows:
1. Introduction
Message passing graph neural networks (MPNNs) have achieved high predictive performance in various domains (Zhou et al., 2020). To understand these performance gains, researchers have focused on the expressive power of MPNNs (Morris al., 2019; Xu et al., 2019; Maron et al., 2019). However, the binary nature of expressive power excludes any analysis of the distance between graph embeddings, which is considered to be key to the predic- tive power of MPNNs (Liu et al., 2022b; Li & Leskovec, 2022; Morris et al., 2024). Recently, there has been growing interest the analysis of MPNN (generalization) perfor- mance using structural distances between graphs (Chuang & Jegelka, 2022; Böker et al., 2024; Franks et al., 2024) that consider graph topology but ignore the target function to be
tree,
'TU Wien, Vienna, Austria Lancaster University Leipzig, Leipzig, Germany. Correspondence to: Masahiro Negishi <m.negishi25 @ imperial.ac.kk>.
We show that MPNN distances after training are aligned with the task-relevant functional distance of the graphs and that this is key to the high predictive performance of MPNNs.
Proceedings of the 4nnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

WILTing Trees:
Interpreting the Distance Between MPNN Embeddings
We propose a trainable graph distance on weighted Weisfeiler-Lchman Labeling Tree (WILT) that gener- alizes Weisfeiler Leman-based distances and is effi- ciently computable.
G
H
20
WILT allows straightforward definition of relevant subgraphs. Thus, distilling an MPNN into WILT enables us to identify subgraphs that strongly influence the distance between MPNN embeddings, allowing an interpretation of the MPNN embedding space.
Figure 1: Example of two iterations of the Weisfeiler Leman algorithm. and @ are colors corresponding to initial node attributes, while and represent edges with different attributes. Node colors in iterations one and two are shown in the small circles next to the nodes. For example, @ = HASH( (((o -), (o. -)D) and e = ( -). ( ". >))) where HASH is a per- feet hash function.
2. Related Work
Recently, it has become increasingly recognized that the geometry of the MPNN embedding space, not just its binary expressiveness, is crucial to its performance (Li & Leskovec, 2022; Morris et al., 2024). For instance, graph contrastive learning methods implicitly assume that good metric struc- ture in the embedding space leads high performance (Liu al., 2022b). Chuang & Jegelka (2022); Böker et al. (2024); Franks et al. (2024) analyzed MPNN's s (generalization) per- formance using structural distances between graphs, which consider graph topology but ignore the learning task at hand. Despite their theoretically sound analyses of generalization bounds, they made strong assumptions about the Lipschitz constants of MPNNs (Chuang & Jegelka, 2022; Böker et al., 2024) or the margin between classes (Franks et al., 2024). In addition. Böker et al. (2024) dealt only with dense graphs and required the consideration of all MPNNs with some Lipschitz constant. Our study also focuses on the geome- try of the embedding space, but we empirically investigate practical MPNNs trained on real, sparse graphs, mainly using task mask-dependent functional distances without limiting Lipschitz constants or the margin between classes.
our method can be applied to graph regression tasks, while previous studies are restricted to classification problems.
3. Preliminaries
We define graph as a tuple G (V,E, E, Inode, ledge). where V and E are the set of nodes and edges. Each node and each edge have an attribute defined by (mode V + Enode and ladge E - Eedge, where Enade and Eedge are finite sets. We restrict them to finite sets because our method is based on the Weisfeiler Leman test described below, which is discrete in nature. We denote the set of all graphs up to isomorphism as G and the set of neighbors of node as ((v). We only consider undirected graphs, but the extension to directed graphs is easy by employing an appropriate version of the Weisfeiler Leman test.
Message Passing Algorithms (Gilmer et al., 2017) in- clude popular GNNs such as Graph Convolutional Networks (GCN, Kipf & Welling, 2017), and Graph Isomorphism Net- works (GIN, Xu et al., 2019). At each iteration, a message passing al algorithm updates the embeddings of all nodes by aggregating the embeddings of themselves and their neigh- bors in the previous iteration. After iterations, the node embeddings are aggregated into the graph embedding hat
This study is also related to GNN interpretability (Liu et al., 2022a; Yuan et al., 2022). Higher interpretability of well- performing models may lead to a new understanding of sci- entific phenomena when applied to scientific domains such as chemistry or biology. In addition, may be a requirement in real-world problems where safety or privacy is critical. Most of the existing interpretation methods are instance- level, identifying input features in a given input graph that are important for its prediction. However, instance-level methods cannot explain the global behavior of GNNs. Re- cently, some studies have proposed a way to understand the global behavior of GNNs by distilling them into highly interpretable models. The resulting model can be GNN with higher interpretability (Müller al., 2024), or a logical formula (Azzolin et al., 2023; Köhler & Heindorf, 2024; Pluska et al., 2024). Our study also distills an MPNN into a highly interpretable WILT for global-level understand- ing. This allows to interpret the metric structure of the MPNN embedding space, while previous studies focused on generating explanations for each label class. In addition,
d = (GG ((() ledgge E N(v)B)
h UPD (
hg READ (EN (UEV))
Here, B.B denotes a multiset, and 0 < < L with h (w)d ! Rd and he Rdd are the embedding of node after the I-th layer and the graph embedding, respectively. AGG("), (PD()) and READ are functions. Message Passing Neural Networks (MPNNs) implement (PD( and (G using multilayer perceptrons (MLPs). Sum and mean pooling are popular for READ.
2

---

WILTing Trees: Interpreting the Distance Between MPNN Embeddings
The Weisfeiler Leman (WL) Algorithm is a message passing algorithm, where UPD(") is an injective function. AGG and READ are the identity function multisets. A node embedding of the WL algorithm is called color. We use de instead of to to refer to it. Figure I shows the progress of the WL algorithm on two graphs: G and H start with the same colors, but no longer share colors after two iterations, i... {-.0) Vee (c.²) e
functional pseudometric space (9, drune) is obtained from drune 9 9 - [0, defined as:
yo lve Vul supyr iutur
(classification)
diunc (G, H) =
(regression),
where yyw is the indicator function that returns if yG # yır, otherwise 0.
0. Message Passing Pseudometrics The WI. algorrthm cannot distinguish some nonisomorphic graphs (Cai et al., 1992) and all MPNNs are bounded by its expressiveness (Xu et al., 2019). Hence, any MPNN yields a pseudometric on the set of pairwise nonisomorphic graphs G.
See Appendix A.2 for a proof that (g, drune) is a pseudomet- ric space. If the sup/inf of yG in 9 are unknown, they can be approximated by the max/min in a training dataset.
The Expressive Power of message passing algorithm is defined based on its ability to distinguish non-isomorphic graphs. Formally, a message passing graph embedding function f is said to be at least as expressive as another one g if the following holds:
Definition 3.1 (Graph Pseudometric). A graph pseudomet- ric space (G, d) is given by non-negative real valued func- tion d : G x G - Rz that satisfies for all F, G, H E g:
G) =0
(Identity) (Symmetry)
vG, G f(G) = f(H) = g()))
d(G, = d) d(G, F) S d(G, + d( F) (Triangle inequality)
where G is the set of all pairwise non-isomorphic graphs. We extend the above to pseudometrics on graphs. Specifically, graph pseudometric d is said to be at least as expressive as d' (d 2 d') iff
Given an MPNN, we obtain a pseudometric space (9, dMPNN) by setting dG H) == d(ha. hu), where and lig and hи
d: Rdd x Rd + R a (pseudo) are graph embeddings. Note that (G, dMPNN) is not metric space since there are nonisomorphic graphs G, H with iden- tical representations and hence dMPNN (G, H) O. For the rest of this paper, we will use dMPNN H) = |lkg -hr l|2. but other distances between embeddings can also be used. Note that dMPNN depends not only on the input graphs but also on the task on which the MPNN is trained. For exam- ple, dMPNN of an MPNN trained to predict the toxicity of molecules will be different from dMPNN of another MPNN trained to predict the solubility of the same molecules.
VG, H E :: : =0 = d'(G, H) = 0.
d and d' are equally expressive (d = d') iff d z d' and d' 2 d. Furthermore, is said be more expressive than d' (d d') iff d 2 d' and there exists G. H E 9 s.t. d(G. H) # H) = O.
4. Is the MPNN Embedding Distance Critical to Performance?
Our first question is what properties dsINN of well perform- ing MPNNs have in practice that can explain their high performance. This section investigates whether the align- ment between dMPNN and the relk-relevant pseudometric duac is such a property. Specifically, we address the ques- tions below:
Structural Pseudometrics To date, many different graph kernels have been proposed (see Kriege et al., 2020). Each positive semidefinite graph kernel ls : G x g + R cor- responds to a pseudometric between graphs. Please see Appendix A. for the definitions of structural pseudometrics from previous studies that are used in this article. We will refer to these pscudometrics structural pseudometries and write d as they only consider the structural and node/edge attribute information of graphs, without being trained using the target label information.
Q1.1 Does training an MPNN increase the alignment be- tween dMPNN and the task-relevant diune?
Q1.2 Does a strong alignment between dMPNN and drunc indicate high performance of the MPNN?
Functional Pseudometrics To formally define the func- tional distance between graphs, we introduce another pseu- dometric on G that is based on the target labels of the graphs.
Note that the alignment between dMPNN and mask-irrelevant struetural graph pseudometrics darue has been considered key to MPNN performance in previous studies (Chuang & Jegelka, 2022; Böker et al., 2024; Franks et al., 2024). However, we found that this property is not consistently im- proved by training and does not correlate with performance. (See Appendix D for detailed analyses).
Definition 3.2 (Functional Pseudometric). Let yG be the target label of graph G in a given task. In classification, ya is a categorical class, while yG is a numerical value in re- gression. We assume the space for yG is bounded. Then, the
3

---

WILTing Trees: Interpreting the Distance Between MPNN Embeddings
To answer Q1.1 and Q1.2, we will first define measure of the alignment between dMPNN and diune- Note that it is inappropriate to adopt a typical min/max of dun (G.H). dhann (G. measure the alignment. This is because dame is a binary function for classification tasks, and expecting the exact match of the two distances is unreasonable. Thus, we define evaluation criterion follows.
Mutagenicity
ENZYMES
0.4
0.73
ALLx
ALL
0.50
0.2
0.25
0.00
Definition 4.1 (Evaluation Criterion for Alignment Between dMPNN and ). Let 'D be a graph dataset, (/2 1) be an integer hyperparameter, and (G) CD) \ {G} be set of graphs that are closest to G under dMPNN. Let
10
20
10
20
k
k
Lipophilicity
711
GG)
1
Σ
drune ( (G, H),
0.15
K HEN
ALL
untroined
0.10
trained
Bx( = D
1
Σ
drun H).
0.05
Then, dMPNN is aligned with diune
1
10
20
k
ALIg (dMENNNN difun)
Σ
[~A+(G)+
BB(G)
(D)
Figure 2: The distribution of ALI (dMPNN, drune) under dif- ferent lt and datasets.
GED
is positive. In addition, the larger ALI; is, the more we say dMPNN is aligned with drune.
Table 1: Correlation (SRC) between (LMENNNN drane) and the performance on Dirain and Diest under different k. We use accuracy for Mutagenicity and ENZYMES, and RMSE for Lipophilicity to measure performance.
Here, As (G) and (G) are the average functional dis- tances between G and its neighbors and non-neighbors, re- spectively. If (G) < B,(()), then G and closeby graphs in MPNN space have lower average functional distance than G and far away graphs in MPNN space.
Mutagenicity
ENZYMES
Lipophilicity
k
train
test
train
test
train
test
We show the distribution of (dMPNN, dlune) for 48 dif- ferent MPNNs on different datasets and varying & in Fig- Each model was trained with a standard loss function (cross entropy loss for classification and RMSE for regres- sion). We did not explicitly optimize ALLg. We also include the results for untrained MPNNs to see the effect of training. We can see that there is little overlap between the distribu- tions of the untrained and trained MPNNs. This means that ALI: consistently improves through training, implying a positive answer to Q1.1. Next, we compute Spearman's 's rank correlation coefficient (SRC) between ALIg (dMPNN; drune) of trained MPNNs and their predictive performance. We use accuracy and RMSE between the ground truth target and predicted values to measure classification and regres- sion performance, respectively. Table shows that SRC for Mutagenicity and ENZYMES always positive, indicating that the higher the ALL,, the higher the accuracy. Similarly, the higher the ALIg, the lower the RMSE for Lipophilic- ity. The correlations sets. These results suggest that the degree of alignment be- tween dMPNN and drune is a crucial factor contributing to the high performance of MPNNs, answering Q1.2 positively. See Appendix C for more details and additional results on non-molecular datasets.
1
0.66
0.70
0.89
0.49
-0.65
-0.57
5
0.65
0.69
0.87
0.50
-0.63
-0.56
10
0.63
0.68
0.87
0.47
-0.62
-0.56
20
0.61
0.67
0.85
0.46
-0.60
-0.55
5. WILTing Pseudometrics
Section 4 confirms that MPNNs are implicitly trained so that dMPNN aligns with drune. which turns out be crucial for MPNN's performance. Then, second research question is: how do MPNNs learn dMPNN that respects drun?? To this end, we note that MPNN embeddings are aggregations of WL color embeddings, i.e., there exists function f with hy = f(ffc, E Voll.. Hence, dMPNN defines a pseudometric dy between multisets of WL. colors
AMPPN(G. H) = 1 - hnll2
consistent across training and test
(1) = [If + (1) 1 E Vo f(fef) l0e Va
=: df(c v Vol. (c(L) (uEVH))
Therefore, we aim to understand dMPNN by distilling it into our more interpretable and equally expressive pseudometric dwnr on the same multisets. dwar is an optimal transport

---

WILTing Trees: Interpreting the Distance Between MPNN Embeddings
distance on the weighted Weisfeiler Leman Labeling Tree (WILT) and generalizes two existing pseudometrics of high- performing graph kernels (Kriege et al., 2016; Togninalli et al., 2019). After distillation, the edge parameters of dwirT allow us to identify WI. colors whose presence or absence significantly affects the relative position of graphs in the MPNN embedding space. addition, dwur is efficient to compute as the ground metric a path length on the tree.
r
5.1. Weisfeiler Leman Labeling Tree (WILT)
The Weisfeiler Leman Labeling Tree (WILT) TD is a rooted weighted tree built from the set of colors obtained by the WL test graph dataset 'D _ G. Given D, we define (TT)) as the set of colors that appear on any node during the WL test plus the root node r, that is, V(Tb) = (ce) | o E Vc, G.G. D, / [4]] J (r). Colors 2, y E ((TD) (r] are adjacent if and only if there exists node in some graph (1) in D and an iteration with x = cs and y = c d r is connected to all æ = ) Due to the injectivity of the AGG and UPD functions in the WL algorithm, it follows that TD is a tree. Figure (top) shows the WILT built from the graphs G and H in Figure 1. See Appendix for detailed algorithm to build a WILT from D.
2/
3
2
1
0
2
0
2
2
0
0
2
0
2
0
2
8G
g0
2
2 0
2
2
2
0 2
Figure: 3: (top): The Weisfeiler Leman Labelling Tree (WILT) built from D = {G, from Figure 1. (middle): The WILT built from D {G, with dummy nodes. (bottom): The WILT embeddings v, ú with size normaliza- tion, and & with dummy node normalization.
We consider edge weights E(Tp) + Rz0 WILT. We only allow non-negative weights so that the WILTing distance in Definition 5.1 will be non-negative. Given a WILT To with weights w, the shortest path length doan (ar yi w) == EeEPath(a..) w(e) is the sum of edge weights of the unique shortest path Path( x. y) between x and y. Note that dpath is a pseudometric on V V(Tp). i.e., the set of WL colors in D. Intuitively, dpan (z. y; w) is large if Path y) is long, but w allows us to tune this pseudometric according to the needs of the learning task.
Definition 5.1 (WILTing Distance). Let G, H E D be
graphs with |Val = [Vr|- Then
|hennThen
dwiLr(G. H:w) = min
Σ Σ
Pu,dpant
(ct)) (oatt( ,Cu
(L))
PEr
E.VV NyVV
where r= (P Pio 2 0, P1 11 1}.
5.2. The WILTing Distance
WILT To with edge weights w yields pseudometric dwIT on the graph set D. This section shows two equivalent characterizations of dwILT an optimal transport distance and as a weighted Manhattan distance. The latter allows us to define the importance of specific WL colors and to com- pute our proposed pseudometric efficiently. For simplicity, we define dwILr for graphs with the same number of nodes. In the next section, we will discuss the extension to graphs with different numbers of nodes. For two distributions with identical mass on the same pseudometric space, optimal transport distances such as the Wasserstein distance (Villani, 2009) measure the minimum effort of shifting probability mass from one distribution to the other. Each unit of shifted mass is weighted by the distance it is shifted. We define our pseudometric (W, H; w) as the optimal transport between Ve and VH. where the ground pseudometric is the shortest path metric on the WILT Tp.
Note that dwILT is not metric but a pseudometric on the set of pairwise nonisomorphic graphs G. This is because there are nonisomorphic graphs G and H whose colors are the same after L iterations, i.e., (L.) E Va (cb E Vn B.
(L)
Generic algorithms to compute Wasserstein distances re- quire cubic runtime. In our case, however, there exists a lineur time algorithm to compute dwn. as shown below, since the ground pseudometric dpah is the shortest path met- ric on a tree (Le al., 2019).
Definition 5.2 (WILT Embedding). The WILT embedding of a graph G E D a vector, where each dimension counts how many times corresponding WL color appears during the WL test on G, i.c., ve := He e Vo 31 E [L] ci c)| for e E ((Tb) \ (r). (see the bottom of Figure 3).
5

---

WILTing Trees: Interpreting the Distance Between MPNN Embeddings
Proposition 5.3 (Equivalent Definition of WILTing Dis-
Algorithm Optimizing edge weights of WILT
tance). dwur in Definition 5.1 is equivalent to:
Input:
graph dataset D, an MPNN J with L message
dwnr(G, H) E
(eta.peet) 2e v²! (EV(To))
passing layers trained on D, and WILT T' built from the results of L-iteration WL test on D
Parameter: batch size, number of epochs E, and learn- ing rate Ir
where e(c.(())) is the edge connecting c and its parent plc) TD.
Output: learned edge weights w of WILT TD ne < [E(Tb)] <- R
This equivalence allows efficient computation of dwur given the WILT embeddings of graphs, which can be com- puted by the WL algorithm in ( EGl) time, where is the number of WL iterations. Using sparse vectors for D and ,", dwarrGG H) can be computed in O([Val + [V+l))
optimizer +- Adam(params=u, Ir=(r)
for to E do for batch B in DD2 do Σ
(dwnr(G, H) dMMPNN H))
(G,H)EB
5.3. Normalization and Special Cases
1. 1.backwardg
The definition of (WWLT(G,a H) as an optimal transport dis- tance requires |Vll Vx However, |Vel and [Vrl are usually different, so propose solutions. Interest- ingly, the two modified WILTing distances generalize two pseudometrics corresponding to well-known graph kernels.
optimizer.stepO
1 Ensuring that edge weights w are non-negative */
w +- max(w,
return w
Size Normalization Straightforwardly, we can restrict the mass of each node to |Vol when calculating the Wasserstein distance in Definition 5.1. In other words, we replace I' with f = {P E TR. Vaalx(Vsi P., 2 0, P1 Val Val P Vzl 1). Similarly, n in Proposition 5.3 is changed to v GG == \Vol The resulting pseudometric dwur effectively ignores differences in the number of nodes of G and H, generally assigning fractions of colors in G to colors in H. In Figure 3 (bottom), we show ú of G and H from Figure dwur(G, H) is calculated as:
where N maxced Val (Figure 3 bottom). Then, the re- sulting pseudometric dwir((, H) for the graphs in Figure I is:
dWILT H) =w ( + (
w(e 1/.
Similar to size normalization, dwnr includes the pseudomer- ric of Weisfeiler Leman optimal assignment kernel (Kriege et al., 2016) as special case. When }.2 dWILT is equivalent to their pseudometric. See Appendix A.3 for details.
3
2
2
2
dwar(G, H) =w(e
+ (((
5
4
5
2
5.4. Edge Weight Learning and Identification of Important WL Colors
w(((o,))))
An interesting property of dwur is that it generalizes the pseudometric corresponding to the Wasserstein Weis- feiler Leman graph kernel (Togninalli et al., 2019): when w = -[1+Ij dwiLT matches their pseudometric. See Ap- pendix A.3 for technical details.
Now, we have a graph pseudometric on WILT defined for any pairs of graphs in D. Next, we show how to optimize the edge weights w. Proposition 5.3 allows us to learn the edge weights w, given training data. Specifically, given a target pseudometric drarger we adapt dwnr by minimizing
Dummy Node Normalization We can also add isolated nodes with a special attribute, called dummy nodes, to graphs so that all the graphs have the same number nodes. The WILT will be built in the same way as described in Section 5.1 after dummy nodes
[(w) =
Σ (G.
(dwir(G,, H;w) - charga (G, H))
with respect to w. Note that dWILT can refer both dwIr and dwIr.. In this work, we focus on diarget dMPNN. That is, train dwır to mimic the distances between the graph embeddings of a given MPNN, as shown in Algorithm 1. Once we have trained w by minimizing C, we can gain insight into dMPNN via dwI..- WL colors with large edge weights are those whose presence or absence in a graph significantly affects dMPNN between the graph and other
added all graphs in D. The resulting WILT has new colors eº, cd...... c that arise from the WI. iteration on the isolated dummy nodes (Figure 3 middle), The WILT embedding will be slightly changed to
N Volifce ct,,c. G otherwise
6

---

WILTing Trees: Interpreting the Distance Between MPNN Embeddings
graphs. Specifically, we can derive the following reasoning.
Table 2: The mtd of RMSE(dMPNN) d) [x [x10- over five different seeds. Each row corresponds to a GCN with a given graph pooling method, trained on a given dataset.
Large difference between G and H in the number or ratio of WL colors with a large ((((er)()))
Large (Wwr(G, H) (: Proposition 5.3)
dwwi.
dwioA
dwnr
dwn.
Large dMPNN(G, H) (-" dwur approximates dAPPN)
Mutagenicity
meam 9.25
18.7443.36
3.34 1.7410.52
sum ENZYMES
12.25:+0.54
1.98
1.22++0.31
0.82+0.17
5.5. Expressiveness of Pseudometrics WILT
mean 12.18+00.23
16.79+2.33
2.71+0038
4.64
sum
11.28:10.65
6.83:10.41
9.15=0.47
1.43:0.10
Here, we discuss which of the two normalizations is pre- ferred for a given MPNN based on the expressive power. Below are the relationships between the expressiveness of AMPNN and dWILT-
Lipophilicity
mean 10.92+00.42
13.97+0.97
0.11+0.54
1.35
sum
10.83:10.73
1..00::1.34
2.50-0067
2.64110.7
Theorem 5.4 (Expressive Power of the Pseudometrics on WILT), Let MPPN and dMPNN be dMPNN of MPNNs with mean and sum graph poolings, respectively. We also define a pseudometric based on the L-iteration WL test:
distribution of learned edge weights of WILT is skewed towards 0, and a large part of them can be removed with L.I regularization. Finally, we investigate the WL colors that influence dMPNN most. Due space limitations, we report results only for selection of MPNNs and datasets. Code is available online, and experimental settings and additional results are in Appendix E.
dwI(G, H) 160 (cd
Then, for WILT with positive edge weights, the following relationships hold between the expressive power of each pseudometric.
We trained 3 3-layer GCNs with mean or sum pooling on the three datasets with five different seeds. We then distilled each into two WILTs, one with size normalization and one with dummy node normalization. To evaluate how well distance d approximates dMPNN, we used variant of RMSE:
dWIr < dwir = dWL:
gmean S dWILT(< :wur)):
MPNN I dwur, daMPNN SdwwE.
RMSE(dMPN)) d)
Proof. See Appendix A.4.
min NER
(GM H) (G.H)EDZ
Σ
a -
-d. 2
DI2
Since dwur is more expressive than dwar, one might think
that dwur is always preferable to approximating dMPNN- However, dwnr is expected to be better at approximating Imean <MPNN. since it provides a tighter bound. Intuitively, this follows from the fact that mean pooling and the size nor- malization are essentially the same procedure: they both ignore the number of nodes. In contrast, dwirr is expected to work well on MPNN. which retains the information about the number of nodes and thus cannot be bounded by dwIrT- We will experimentally confirm these analyses in Section 6. Note that Theorem 5.4 considers only the binary expressive- ness of pseudometrics. Regarding the size of the family of functions that each pseudometric can represent, dMPNN is expected to be superior to dWIL., because dwIT is restricted to an optimal transport on the tree for faster computation and better interpretability. Still, in Section 6, we empirically show that dwnr can approximate dMPNN well.
where dwnr and d means they are normalized to [0, 1]. Intu- itively, the closer the RMSE is to zero, the better the align- ment is, and zero RMSE means perfect alignment. We do not use the correlation coefficient hecause it can be one even if dMPNN is not a constant multiple of d: it allows a non-zero intercept. Note that the minimization over a can be solved analytically. In practice,
compute the RMSE using only 1000 pairs from D2 for speed. Table 2 shows the RMSE between dMPNN and dwnr or dwIr. We also include results for dwwL and dwLOA. which are special cases of dwur and dwnrr with fixed edge weights, respectively. It is obvious that dwılr aligns with dMPNS much better than dwwL and dwr Oww.O Interestingly, dwn. approximates dMMPNN (mean) bet- ter, while dwILT approximates dMPNN(sum) better, except for dsass(sum) trained on Lipophilicity, where their per- formance is close. This observation is consistent with the theoretical analysis in Section 5.5.
6. Experiments
Next, we look into the distribution of the learned edge weights of WILT. Figure 4 shows the histogram of the edge weights of the WILT with dummy node normalization after
In this section, we confirm that our proposed dwiLr can successfully approximate dMPNN- Then, we show that the
7

---

WILTing Trees: Interpreting the Distance Between MPNN Embeddings
Weight distribution
Results with different X
10a
14
0.07
Rotio of 2000-200 weights
(3)
10a
0.06
Prequency
0.05
RMSE
422
10.
0.04
0.03-
I155 of WELD
(
10'
I1ME of WETT
0.02
Rallo of ausoer walla
0.01
10
0.0001 0.001 0.01
Lo
Cl2y
Weight value
Ll coefficient
Figure 4: The distribution of the edge weights of WILT after distillation.
Figure 5: The RMSE and the ratio of non-zero edge weights after distillation under different L1 coefficients (mean and std over five different seeds).
Figure 6: Example graphs with high- lighted significant subgraphs corre- sponding to colors with the largest weights.
distillation from a 3-layer GCN with sum pooling trained on Mutagenicity. The distribution is heavily skewed towards zero. This plot, together with Proposition 5.3, suggests that the relative position of MPNN graph embeddings is deter- mined based on only small subset of WL colors. To further verify this idea, we added an LI regularization term to the objective function L and minimized it so that (((((e)()))) would be set to zero for some colors. Figure 5 shows the RMSE between dMPNN and the resulting dWT.r. as well as the ratio of non-zero edge weights, under different LI coeffi- cient A. As expected, the larger À is, the more edge weights are set to zero and the larger the RMSE. However, it is worth noting that dwILT is much better aligned with dMPNN than dwLOA even when trained with À = 1.0 and about 95% of the edge weights are zero. This good approximation with only 5% non-zero edges implies that MPNNs rely on only a few important WL colors to define dyPNN.
7. Conclusions
We analyzed the metric properties of the embedding space of MPNNs. We found that the alignment with the functional pseudometric improves during training and is a key to high predictive performance. In contrast, the alignment with the structural psudometrics does not improve and is not correlated with performance. To understand how MPNNs learn and reflect the functional distance between graphs, we propose a theoretically sound and efficiently computable new pseudometric on graphs using WILT. By examining the edge weights of the distilled WILT, we found that only a tiny fraction of the entire WI. colors influence dMPNN- The identified colors correspond subgraphs that are known to be functionally important from domain knowledge.
One limitation of our study is that we only distilled two GNN architectures (GCN and GIN) with fixed hyperparam- eters to WILT. Thus, it remains to be seen how different architectures and hyperparameters affect the edge weights of WILT. We also limited our analysis to the final embeddings of MPNNs, but in principle, WILT can be trained to approx- imate the MPNN embedding distance at internal layers. It would be interesting to investigate the embedding distances at different layers, and how they relate to the performance. We expect results similar to I Liu et al. (2024), which showed that consistency between distances at different iterations is key to high performance. While we investigated MPNNs specifically, there a hierarchy of more and more expressive
Finally, we show the subgraphs corresponding the colors with the largest weights, thus influencing dMPNN the most. Again, we only show results for the 3-layer GCN with sum pooling trained on the Mutagenicity dataset. To avoid iden- tifying colors that are too rare, we only consider colors that appear in at least I % of the entire graphs. Figure 6 shows ex- ample graphs with subgraphs corresponding colors with the four largest weights. The identified subgraphs in (1) and (4) are known to be characteristic of mutagenic molecules (Kazius et al., 2005). In fact, (1) and (4) are classified as "epoxide" and "aliphatic halide" based on the highlighted subgraphs. Given that only a tiny fraction of the entire WL colors correspond to the subgraphs reported in (Kazius et al., 2005), this result suggests that MPNNs learn the relative po- sition of graph embeddings based on WI colors that are also known to be functionally important by domain knowledge.
GNNs that are bounded in expressiveness by corresponding WL test variants. In this paper, we have defined WILT the hierarchy of 1-WL labels. Still, it is straightforward to ex- tend the proposed WILT metric to color hierarchies obtained from higher-order WL variants (Morris et al., 2023; Geerts & Reutter, 2022) or extended message passing schemes (Frasca al., 2022; Graziani et al., 2024). While beyond
8

---

WILTing Trees: Interpreting the Distance Between MPNN Embeddings
the scope of this work, higher-order WILTing trees may prove useful in interprering a range of GNNs. However, as the number of trainable WILT weights scales with the number of colors, the practical relevance of higher-order WILTSs remains an open question. It is also worth exploring extending WILT to the analysis of GNNs for node classifi- cation on large graphs such as social networks or citation networks. Kothapalli al. (2023) showed that GNNs are trained such that their node embeddings respect the func- tional alignment, but their analysis was limited to unrealistic graphs generated from a stochastic block model. We believe that the node embedding space of GNNs can also be dis- tilled to WILT in a similar way, i.e., by tuning the weights to approximate the node embedding distance with a path distance on WILT. Using WILT for purpose other than understanding GNNs is also interesting. For example, by training WILT's edge parameters from scratch, might be able to build high-performance, interpretable graph learning method.
Chuang, C.-Y. and Jegelka, S. Tree mover's distance: Bridging graph metrics and stability of graph neural net- works. Advances in Neural Information Processing Sys- tems, 2022.
Franks, B. J., Morris, C., Velingker. A., and Geerts, F. Weisfeiler-leman at the margin: When more expressivity matters. In Forty-first International Conference on Ma- chine Learning, ICML 2024, Vienna, Austria, July 21-27, 2024, 2024.
Frasca, F., Bevilacqua, B., Bronstein, M. M., and Maron, H. Understanding and extending subgraph GNNs by rethink- ing their symmetries. In Advances in Neural Information Processing Systems, 2022.
Geerts, F. and Reutter, J. L. Expressiveness and approxima- tion properties of graph neural networks. In International Conference on Learning Representations, 2022.
Gilmer, J., Schoenholz, S. S., Riley, P. F., Vinyals, O., and Dahl, G. E. Neural message passing for quantum chem- istry. In International Conference on Machine Learning, 2017.
Acknowledgements
PW acknowledges TGs ability to devise great paper titles involving puns on some authors' last names. This work was supported by the Vienna Science and Technology Fund (WWTF) and the City of Vienna project StruDL (10.47379/ ICT22059) and by the Austrian Science Fund (FWF) project NanOX-ML (6728). We express our gratitude to the Japan Student Services Organization for financially supporting the first author's five-month research stay in Vienna. Finally, we thank the reviewers for their insightful comments.
Graziani, C., Drucks, T., Jogl, F., Bianchini, M., Scarselli, F., and Gärtner, T. Expressiveness and approximation proper- ties of graph neural network. In International Conference on Machine Learning. 2024.
Kazius, J., McGuire, R., and Bursi, R. Derivation and validation of toxicophores for mutagenicity prediction. Journal of Medicinal Chemistry, 48 48(1):312-322 2005.
Kipf. T. N. and Welling, M. Semi-supervised classifica- tion with graph convolutional networks. In International Conference on Learning Representations, 2017.
Impact Statement
This paper advances our understanding of what and how competitive MPNNs learn. We hope that this will contribute to safe and robust machine learning. We are not aware of any immediate negative consequences of our work.
Kothapalli, V., Tirer, T. and Bruna, J. A neural collapse perspective on feature evolution in graph neural networks. Thirty-seventh Conference on Neural Information Pro- cessing Systems, 2023.
References
Kriege, N. M., Giscard, P., and Wilson, R. C. On valid opti- mal assignment kernels and applications to graph classi- fication. In Advances in Neural Information Processing Systems, 2016
Azzolin, S., Longa, A., Barbiero, P., Liò, P., and Passerini, A. Global explainability of GNNs via logic combina- tion of learned concepts. In International Conference Learning Representations, 2023.
Kriege, N. M., Johansson, F. D., and Morris, C. A survey on graph kemels. Applied Network Science, 5:1-42, 2020.
Böker, J., Levie, R., Huang, N., Villar, S., and Morris, C. Fine- -rainer expressivity of graph neural networks. Ad- vances in Neural Information Processing Systems, 36, 2024.
Köhler, D. and Heindorf, S. Utilizing description logics for global explanations of heterogeneous graph neural networks. arXiv 2reprimt:2405.12654 2024.
Cai, J., Fürer, M., and Immerman, N. An optimal lower bound on the number of variables for graph identification. Combinatorica, 12(4):389-41) 1992.
Le, T., Yamada, M., Fukumizu, K., and Cuturi, M. Tree- sliced variants of Wasserstein distances. Advances in Neural Information Processing Systems, 2019.
9

---

WILTing Trees: Interpreting the Distance Between MPNN Embeddings
Li, P. and Leskovec, J. The expressive power of graph neu- ral networks. In Graph Neural Networks: Foundations, Frontiers, and Applications. 2022.
Togninalli, M., Ghisu, E., Llinares-López, F,, Rieck, B., and Borgwardt, K. Wasserstein Weisfeiler-Lehman graph kernels. Advances in Neural Information Processing Systems, 2019.
Liu, N., Feng, Q. and Hu, X. Interpretability in graph neu- ral networks. In Graph Neural Networks: Foundations, Frontiers, and Applications. 2022a.
Villani, C. Optimal transport: old and new, volume 338 of Grundlehren der mathematischen Wissenschaften. Springer, 2009.
Liu, X., Cai, Y., Yang. Q. and Yan, Y. Exploring consistency in graph representations: from graph kernels to graph neural networks. In The ehihty- Annual Conference on Neural Information Processing Systems, 2024.
Wu, Z.,, Ramsundar, B., Feinberg, E. N., Gomes, J., Ge- niesse, C., Pappu, A. S., Leswing, K., and Pande, Moleculenet: benchmark for molecular machine learn- ing. Chemical Science, 9(2):113-530 2018.
Liu, Y., Jin, M., Pan, S., Zhou, C., Zheng, Y., Xia, F., and Yu, P. S. Graph supf-suuee learning: A survey. IEEE Transactions on Knowledge and Data Engineering, (6):5879-5900, 2022b.
Xu, K., Hu, W., Leskovec, J., and Jegelka, S. How powerful are graph neural networks? In International Conference on Learning Representations, 2019.
Yuan, H., Yu, H., Gui, S., and Ji, S. Explainability in graph neural networks: A taxonomic survey. IEEE Transac- tions on Pattern Analysis and Machine Intelligence, 45 (5):5722-579,, 2022.
Maron, H., Ben-Hamu Hamu, H., Serviansky, H., and Lipman, Y. Provably powerful graph networks. Advances in Neural Information Processing Systems, 2019.
Morris, C., Ritzert, M., Fey, M., Hamilton, W. L., Lenssen, E., Rattan, G., and Grohe, M. Weisfeiler and Leman go neural: Higher-order graph neural networks. In AAAI Conference on Artificial Intelligence, 2019.
Zhou, J., Cui, G., Hu, S., Zhang, 7., Yang, C., Liu, Z., Wang, L., Li, C., and Sun, M. Graph neural networks: A review of methods and applications. Al open, 1:57-81, 2020.
Morris, C., Kriege, N. M., Bause, F., Kersting, K., Mutzel, P., and Neumann, M. Tudataset: A collection of benchmark datasets for learning with graphs. In ICML Workshop on Graph Representation Learning and Beyond, 2020.
Morris, C., Lipman, Y., Maron, H., Rieck, B., Kriege, N. M., Grohe, M., Fey, M., and Borgwardt, K. M. Weisfeiler and leman go machine learning: The story far. Journal of Machine Learning Research, 24:333:1-333:3 2023.
Morris, C., Frasca, F., Dym, N., Maron, H., Ceylan, i., Levie, R., Lim, D., Bronstein, M. M., Grohe, M., and Jegelka, Position: Future directions in the theory of graph machine learning. In International Conference on Machine Learning, 2024.
Müller, P., Faber, L., Martinkus, K., and Wattenhofer, R. GraphChef: Decision-tree tree recipes to explain graph neu- ral networks. In International Conference on Learning Representations, 2024.
Pluska, A., Welke, P., Gärtner, T., and Malhotra, S. Logical distillation of graph neural networks. In International Conference Knowledge Representation and Reasoning, 2024.
Sanfeliu, A. and Fu, K. A distance measure between at- tributed relational graphs for pattern recognition. IEEE Transactions on Systems, Man, and Cybernetics: Systems, 13(3):353-62, 1983.
10

---

WiLTing Trees: Interpreting the Distance Between MPNN Embeddings
A. Theoretical Analysis
In this section, we present the proofs of our theoretical results in the main paper. We formally define related structural pseudometrics and show that dwn. generalizes them.
A.1. Structural Pseudometries
Here we introduce the definitions of the graph edit distance (dGED, Sanfeliu & Fu, 1983), Weisfeiler Leman optimal assignment distance (dwLOs, Kriege et al., 2016), and Wasserstein Weisfeiler Leman graph distance (dwwL, Togninalli et al., 2019). For the definition of tree mover's distance, please refer to the original paper (Chuang & Jegelka, 2022).
Definition A.1 (Graph Edit Distance (Sanfeliu & Fu, 1983)). Let & be the set of graph edit operations, and c & + Rzo be function that assigns cost to each operation. Then, the graph edit distance (GED) between G and H is defined as the minimum cost of a sequence of edit operations that transform G into H. Formally,
dGED (G, H) =
min *ES(G.H) CES
Edee),
where S(G, H) is a set of sequences of graph edit operations that transform G into H.
In the experiments in Appendix D, & consists of insertion and deletion of single nodes and single edges, as well as substitution single node or edge attributes. We set the cost of each operation to I, i.c. c(e) = 1. Next, we move on to the Weisfeiler Leman optimal assignment (WLOA) kernel.
Definition A.2 (Weisfeiler Leman Optimal Assignment Kernel (Kriege et al., 2016)). Consider G (Va, Ea) and H (VH, Ea). Let Va and VH be the extended node sets resulting from adding special nodes G or H so that G and H have the same number of nodes. Let the base kernel & be defined as:
Elzo 1 100=0
(v 2 u
k(V,u)
(u (v= V 2),
where & and eu represent the colors of vertices v and u at iteration of the WL algorithm (see Section 3). Then, the Weisfeiler Leman optimal assignment (WLOA) kernel is defined as;
kWLOA(G. H)
max BEB(V
Σ
k(vG, uz),
Vi)
(EG
where B(V,, Vi²) denotes the of all possible bijections between V/2 and Viz-
Kriege et al. (2016) proved that kwr.ou o^ is positive semidefinite kernel function. While they focused only on the kernel, corresponding graph pseudometric can be defined in the following way:
Definition A.3 (Weisfeiler Leman Optimal Assignment (WLOA) Distance). The function dWLOA below is a pseudometric on the set of pairwise nonisomorphic graphs g:
dWLOA (G, H) = (I+ (aax(VVal, [vl)) -.,LLA H)
Proof. Theorem A.8 shows that dwLOA defined as above is a special case of dwu.r. Since dwur is a pseudometric on the set of pairwise nonisomorphic graphs G. so is dwLoA.
We will show later that the above WLOA distance a special case of our WILT distance with dummy node normalization (Theorem A.8). Togninalli et al. (2019) proposed another graph pseudometric based on the WL algorithm, called Wasserstein Weisfeiler Leman graph distance.
Definition A.4 (Wasserstein Weisfeiler Leman (WWL) Distance (Togninalli et al., 2019)). Let d (v, u) be the hamming
distance between ( ,00)) ((),....(()
(1)
co. (L) and Cu .... Cu (L) where ( is the color of node v at iteration of the WL
(1)
algorithm (see Section 3). Specifically,
dham (v, u) =
GV 40
L
1=0

---

WILTing Trees: Interpreting the Distance Between MPNN Embeddings
Then the WWL distance is defined
dwwL(G, = min Σ
Σ
Paycham (U; Uz),
PeTwwL v. MvEV
where Twwr. (P distributions.
R Valx x Val
P1
1/2 Val 1, 11
[V 1} is of valid transports between two uniform discrete
Togninalli et al. (2019) have shown that dwwL is a pseudometric. In addition, they proposed corresponding kernel (Ww H) e (G.H) and showed that it is positive semidefinite. We will prove later in Theorem A.7 that WILT distance with size normalization includes this WWL distance as special case.
A.2. Functional Pseudometries
Here, show that diue is pseudometric.
Definition 3.2 (Functional Pseudometric). Let yo be the target label of graph G in given task. In classification, yG is categorical class, while is a numerical value in regression. We assume the space for yo is bounded. Then, the functional pseudometric space (9, drune) is obtained from drunc G x G + [0, 1] defined as:
1 ya Fux Iuee-Vi
(classification)
diunc (G, H)
(regression),
inl I89
Feb
where 1 voru the indicator function that returns 1 if yG F yH, otherwise
Proof. We start with the classification case.
diunc (G, G) Lyc
diunc (G, H)
yn $yo drune H, G)
diune (G, F)
Lyctur
UG uG + /
drac (G, H) + dranc (H, F)
For the regression case we can assume that supyr * inf yr as the regression problem is trivial otherwise. Then we can
reg TEg
12

---

WILTing Trees: Interpreting the Distance Between MPNN Embeddings
show:
diunc (G, G)
y/G
yG
supyj 1e9
inf y/
IE9
= 0
B VH
drunc (G, H)
SuPys
inf y!
1EG
IE9
ycl
supyr IEG
inf y! lE9
drune [H, G)
1uG yF
dianc (G,
SuPYл inf yr
TEG
IE9
yc
yH
IY yF
supys
inf yr IE9
supyr TEG
inf yr 1E9
TEG
drune (G, H) diunc (H F))
In both cases, identity, symmetry, and triangle inequality are satisfied.
the supremum (resp. infimum) of ye in G are unknown, they can be approximated by the maximum (resp. minimum) distance in a training dataset D, and we can similarly prove that diun is pseudometric.
A.3. Normalized WILTing Distances and Relationship to Existing Distances
We present the formal definitions of the size normalization and dummy node normalization. We then show that dwir with size normalization generalizes the WWL distance dwwr. and that dwur with dummy node normalization generalizes the WLOA distance dwLos-
Definition A.5 (WILTing Distance with Size Normalization). We define the WILTing distance with size normalization as:
dwnr (G, Hyw)
min Pet
Σ
Σ
Pasthaat (CL ((L))
V.VV uyEV
where I' (P [VVolllVvll Paj 2 0,
PT
Vel 1). It equivalent to:
Vol
dwnr(G: (G, H)
Σ
w(e(ep(o)) 10c 1
(EV(TO)\
where GG
[Vll DG
The only difference between Definition 5.1 and Definition A.5 is the mass assigned to each node. The equivalence between the two definitions of dwur a straightforward consequence of (Le et al., 2019). The other normalization is defined as follows.
Definition A.6 (WILTing Distance with Dummy Node Normalization). Let Ve be an extension of Vo with additional N V isolated dummy nodes with a special attribute, where N == maxged VVl.- Let To WILT built from the extended graphs {(Vc, EGD] Note that To just slight modification of Tp (see Figure 3). We define the WILTing distance with dummy node normalization
dwur (G, H;w) = min
Σ
Σ
Proc (put alll))
PET v.VVc Vv
13

---

WILTing Trees: Interpreting the Distance Between MPNN Embeddings
where := (P E IR Volx/Vel Pi, 2 20, 1, 1}, and de is the color of node v on To after I. iterations, equivalent definition is:
Σ
w(e(ep(()) c -
(EV(Ta))r)
where &G is the WILT embedding of G using TD-
Intuitively speaking, add dummy nodes all the graphs so that they have the same number of nodes', and compute the WILTing distance in exactly the same way shown in Section 5.2.
Next, we show that dwITT includes the Wasserstein Weisfeiler Leman distance and dwiLr includes the Weisfeiler Leman optimal assignment distance as a special case, respectively.
Theorem A.7 (dwwl, as A Special Case of dwur)). The WWL distance in Definition A.4 is equal to the WILTing distance with size normalization with all WILT edge weights set to Z[+1J- Specifically,
dwwr (G, H) dwnT G,HH: 2(T 1(T++))
Proof.
dwwL (G, H)
min Σ
Σ
Pr (Uis Uz))
PEl'ww
onE Va us EV
I. Σ 1=0
min PETwWL veVa ayEVH
ΣE
P L
l scl
min Pel'ww. v,VVo U,EVH
Σ
Σ
Pirchaahh
((LL) (L),
2(I + 1)
min
Σ
Σ Pagdhaa (L) (L)
(l))cll. 2(L
Pef weVe uyEVa
= dwnr G, H; 2(I+1))
Theorem A.8 (dwLos as A Special Case of dwur). The WLOA distance in Definition A.3 is equal to the WILTing distance with dummy node normalization with all WILT edge weights set to 1 2- Specifically.
dwLoa (G, -dwur G 2
Proof. First, dwLoA (G, H) can be transformed as follows.
dwLoa HI) = (L + max (VVbl - KWLOA (G, , H!)
= (L + 1) max Vol. VHl)
max
Σ
k(vG,
BEB(V) V/A)
(VG
min BEB(VL
Σ
(L + k(ve, uH))
(EG
Next, introduce equivalent definition of k(o, u). In Definition A.2, the WL algorithm is applied only Vo and VH, not on special nodes. Assume w.l.o.g. that [V(G) & (V(H),, i.e., V(G) is extended with [VH)) (G)
D G, it be defined without explicit of N. To this end, that s C - ! = [V(G) - V(H)) for any dummy node color ch.
In fact, dwur remains pseudometric
14

---

WILTing Trees: Interpreting the Distance Between MPNN Embeddings
dummy nodes. By treating the special nodes in Va dummy nodes, we can define WL colors for the special nodes (o2 Co (1)) (ch, (..........)) c²). Then, as only Va contains special nodes, k(u, u) can be simplified to:
k(,,u) = Σ 00 CE
where e0 is the color of node the WILT T' with dummy node normalization after / iterations. Then, L + 1 k(v, u)
equivalent to deatt ( (L), Cu w 1):
L
1+1-k(u, = L+1
Σ
@ -0 !
1=0
Σ
1
(
1=0
(0 (d
dnah
2
Therefore, dwT.OA is a special case of dwnr:
dwT.OA (G, H)
min
Σ
(L+ - (L+1-kll
k(va, uH))
BEB(VA Vi)
(vG.M))EB
min BEB(V Vii) (BG.H)))B
Σ
dpath
( Cux w 2
Lb), 1
(LL) ا(), 1
- min Per
Σ
Σ
Pisdppahh
2
veVa u,EVn
1
dwwL
G, H; 2
Note that * holds since adding the same number of dummy nodes to both G and H does not change the left side, and the optimal transport on WILT always delivers a mass on a node to only one node.
A.4. Expressiveness Graph Pseudometrics
We now discuss in detail the expressiveness graph pseudometrics, which was summarized in Section 5.5. We split Theorem 5.4 in Section 5.5 into three theorems below, and prove each one separately. The discussion below provides a possible explanation for some results in Section and Appendices and First, we introduce pseudometric defined by the WL test:
dwr. (on.(G, H) = loVe De
where L is the number of WL iterations. In other words, (G, H) = if the L-iteration WL test can distinguish G' and H, otherwise With this definition, we start with the comparison of dwILT and dwt, for a better understanding of dwur-
Theorem A.9 (Expressiveness of the WILTing Distance), Suppose dwur and dwnr are pseudometrics defined with WILT with some edge weight functions. We assume that all edge weights are positive for dwur. Then,
dwur < dwur 2t dwL-
15

---

WILTing Trees: Interpreting the Distance Between MPNN Embeddings
Proof. We first show dWLr= S dwirr.
dw H) =0 = ba of
A
VVl = Vvnl
A leaf color c:
[r E Va c = c = 1 E VH I 1g, = ccH
A
|Val = [Val
H E Va d = eHl l Vn 44) = =c]
A leaf color e:
[Vll
i =i
dwnr(G. H) =0.
Note that leaf color e means that e is a leaf of the WILT. The first implication follows from the fact that dummy node normalization implies that only graphs with identical numbers of nodes can have a distance of zero if the weights are positive. To see that dwur(G. (G, H) is more expressive than dwwr(G, H). note that there are G and H s.t. dwar H) & 0 A (C, H)- O: for example, let G and H be k-regular graphs (such as cycles) with different numbers of nodes and identical node and edge attributes. Next, we show dwnr= dwL.
dw () = 0 + D D
A leaf color e: K & E Val (0) = = = He e Vv c0 = ccH
(cf0) vc) = (df Vn
dWl (G, H) =0
The first equivalence again follows from the fact that weights are positive.
Since dMPNN < dwL holds for any MPNN (Xu et al., 2019), the above theorem implies that dMPPN S dwnr if all edge weights are positive. At first glance, this seems to suggest that dwur can better align with dMPNN of any MPNN than dwır because of its high expressiveness. However, the results in Section 6 and Appendix E show that dwar is suitable for MPNNs with mean pooling, while dwnr is suitable for MPNNs with sum pooling. Next, we compare dMPNN and dwnr in more detail to interpret these results. We start with MPNNs with mean pooling, whose pseudometrics we will call MPNN-
Theorem A.10 (Expressiveness of the Pseudometric of MPNN with Mean Pooling). Suppose dwir and dwnr are pseudo- metrics defined with WILT with some edge weight functions. We assume that all edge weights are positive. Then,
meean MPNN S dwarl<< dwar).
Proof. We first show the left inequality.
dwur((, H) 0
ic 3
K E Vo | le = =c]l
[iv Vn
Co = c)l
(T)
V leaf color e
[Val
[n
leaf color e:
Σ
(L
Σ
(
VVal
[HH
eEVi:el
(4) =c
h (1)
Σ
(L))
VVl
VH
vEVG
VEVE
demun (G, (G,
= 0.
The first implication follows from the fact that ((((c.p(o)) for all colors. The third implication follows from Xu et al. (2019) by noting that c (L) co t ni (L) for any MPNN. dwur < dwurT follows from Theorem A.9.
In Section 6 and Appendix E, we show that RMSE mean dwnr) smaller than RMSE( MPNN- dwnr). The above theorem and the proof yield an interpretation of the result. terms of expressiveness, dwar is a stricter upper bound on d'pNN than dwurT since the mean pooling and the size normalization are essentially the same procedure. Both ignore the information
16

---

WILTing Trees: Interpreting the Distance Between MPNN Embeddings
about the number of nodes in graphs. When we try to fit dwIrr to MPNNN it is difficult tune edge parameters that dwir can ignore the number of nodes in graphs, but dwnr satisfies this property by definition. This may be the reason why dwILT can be trained be better aligned with MPNN than dwırr. A similar discussion can be applied to dwwl and dWLOA, which are special cases of dwILr and dwIrT. respectively. Next, we analyze MPNNs with sum pooling.
Theorem A.11 (Expressiveness of the Pseudometric of MPNN with Sum Pooling). Suppose dwur is defined with WILT with an edge function that assigns a positive value to all edges. Then,
Join MPNNN-
S dwn...
In addition, 3f E G s.t. Enevo 1 + then
MPNN dWILT
Proof. We begin with MPNN S dwn.T.
dWLT(G. (G, H) = = G T
A leaf color e:
Ho Vo I (10) = c}l = Ke E Vn I ( = cc
A leaf color e:
Σ
h (1)
Σ
h
V =c
w VEV2 (4) =c
h0 = Σ
(())
VE Va
VEVH
MPNN
(G, = O.
Next, we show dumm XdwITT. Let G be a graph that satisfies E A * O. We can consider graph H that consists of two copies of G. Then, G iH, since 2 n and 2|Val = Virl. Therefore, dwnr(G, H) = 0. On the other hand,
MPNN (G,
Σ h
Σ (
veVV
VEVI
Σ
h
2 Σ
VEVO
VoVo
(LL
VeVo
* O.
In terms of expressiveness, damn is almost always not bounded by dwur except for the trivial MPNN which embeds all graphs to zero. In fact, the opposite dwur S drum holds if the MPNN is sufficiently expressive, e.g. GIN. These analyses may explain why RMSE(duPyy- dwirr) is generally smaller than RMSE(CHPNN dwar) in Section 6 and Appendix E. No matter how much it is trained, dwIrT cannot capture the information about the number of nodes that dMPNN can. On the other hand, dwaT expressive enough to capture the information, and thus has a chance of aligning well with MPNN Again, a similar reasoning can be applied to dwwL and dwLOA.
B. Algorithm to Construct the WILT
Algorithm 2 shows how to build the WILT of a graph dataset D. It starts with the initialization of a root node of the WILT and adds each node color appearing in D as child of the root. The algorithm then runs L iterations of the Weisfeiler Leman algorithm. Whenever a new vertex color c, is encountered at some node v in iteration it is added to the WILT Tp as a child of the color of node e in iteration 1
17

---

WILTing Trees: Interpreting the Distance Between MPNN Embeddings
Algorithm 2 Building WILT
Input: Graph dataset D Parameter: T. 2
Output: WILT TD
TD Initial tree with only the root
for G in D do
Opre 6-
# Keeping colors in the previous iteration # Keeping colors in the current iteration
Cnow 6
1 Add initial colors children of root */
for in Va do
if mode (v) $ (TTp) then
V(Tb) < ((Tb) U (lode(
E(TD) < E(TD) U [(r, Inodee(M) (v))
Cpre [0] f (v)
1* L-iteration WL test on $/
for 1 = 1 to T. do
for in Va do
Cp < (<re((( { (Crre [u], ledge (evo)) | )))))))
# Compute iteration / WL color
/$ Add new colors to WILT #/
Cu & (TD)) then
((TD) + V(TD) U [en)
E(Tp) E(Tp) [(ppe(e) [e), cm)3
Cnow [c] + Cu
Cure f- Caow
Cnow +-
return TDD
C. Experimental Details for Section 4
Here, we present the detailed experimental setup for the results in Figure 2 and Table 1. We conduct experiments on three different datasets: Mutagenicity and ENZYMES (Morris et al., 2020), and Lipophilicity (Wu et al., 2018). We chose these datasets to represent binary classification, multiclass classification, and regression tasks, respectively. For the models, we adopt two popular MPNN architectures: GCN and GIN. For each model architecture, we vary the number of message passing layers (1, 2, 4), the embedding dimensions (32, 64, 128), and the graph pooling methods (mean, sum). This results in total of 2 x 4 x 3 x 2 = 48 different MPNNs for each dataset. In each setting, we split the dataset into Diainn Deval and Duss (8:1:1). We train the model for 100 epochs and record the performance on Deval after each epoch. We set the batch size to 32, and use the Adam optimizer with learning rate of 10-3 AL (dMPNN, drune) and the performance metric (accuracy for Mutagenicity and ENZYMES, RMSE for Lipophilicity) are calculated with the model at the epoch that performed best on Deval- The code run our experiments is available https: github. com/ masahiro- negi shi /wilt.
and COLLAB (obtained from Morris al., 2020). Figure visualizes the distribution of ALL (dMPNN, diune) on these datasets and varying k. Similar to Figure 2, ALI, consistently improves with training. Table 3 also offers results similar to Table 1, showing that there is positive correlation between ALI; of trained MPNNs and their accuracy in general. Figure shows the data used to compute the Spearman's rank correlation coefficient (SRC) in Table and Table for better understanding. Each blue dot
Next, offer additional experimental results non-molecular datasets:
Table Correlation (SRC) between ALI (dMPNN, d and accuracy on Duain and Drest under different k:
IMDB-BINARY
COLLAB
5
10
20
1
5
10
20
train
0.36
0.57
0.55
0.54
0.95
0.94
0.92
0.91
test
-0.43
0.10
0.14
0.16
0.81
0.79
0.77
0.76
18

---

WILTing Trees: Interpreting the Distance Between MPNN Embeddings
IMDB-BINARY
COLLAB
0.4
0.20
ALIK
0.3
untrained
0.15
trained
0.2
0.10
10
20
5
10
20
k
k
Figure 7: The distribution of ALIg(dMPNN) drune) under different & and datasets.
Mutagenicity
ENZYMES
Lipophilicity
IMDE-BINARY
COLLAB
6.00-
SRC
SICO
Sat .051
0.57
SOLL
6.90-
380
A1.
13c
ALL
3.30
0.10
034
0.30
A1C.
6.500
SRC DuEb
SBG 0.50 0.41
SRU--O01 -0.0
SRU 0.10
S011 0,70
POC:
XCC
Fod mose
4.77
6.750
A..
A.L.
ALL.
A11
ABC
Figure 8: Scatter plots between ALIs (dMPNN, d and the performance on the train/test set. SRC and PCC stand for the Spearman's ' rank correlation coefficient and the Pearson's correlation coefficient, respectively. In general, higher ALIs (dMPNN, drune), i.e., higher alignment between dMPNN and d'unes indicates higher performance.
represents one of the 48 different models. We also fit linear function and show Pearson's correlation coefficient (PCC). For ALIg with + 5, similar plots were observed.
D. MPNN Pseudometric and Structural Pseudometrics
There has been intensive research graph kernels, which essentially aims to manually design graph pseudometrics dyarue that lead to good prediction performance. Recent studies have theoretically analyzed the relationship between dMPNN and such darmes but they only upper-bounded dNPNN with damue (Chuang Jegelka, 2022), or showed the equivalence for untrained MPNNs on dense graphs (Böker et al., 2024). Therefore, this section examines dMPNN really aligns with dame in practice, and if the alignment explains the high performance of MPNNs. Specifically, we address the following questions:
Q1.3 What kind of denve is dMPNN best aligned with?
Q1.4 Does training MPNN increase the alignment?
Q1.5 Does a strong alignment between dMPNN and d indicate high performance of the MPNN?
19

---

WILTing Trees: Interpreting the Distance Between MPNN Embeddings
We first define an evaluation criterion for the alignment between dMPNN and darue answer them, which is the same as the one used in Section 6 and Appendix E.
Definition D.1 (Evaluation Criterion for Alignment Between dMPNN and dsauu). Consider graph dataset denoted by D. Let AMPNN and datree be normalized versions of dMPNN and ditree respectively:
AMPNN (G, H) max dMPNN (G', H') (G',H"ED:
datuue (G, H) datrae (G', H")
AMPNN (G H)
dsnuc (G,nu (G
max [G' (G'.H))DS
We measure the alignment between dMPNN and datrne by the RMSE after fitting a linear model with the intercept fixed at zero to the normalized pseudometries:
RMSE (AMEE,, daaue) ==
min aER DD22
Σ (G,H)E
(G. H) -x- dsauue (G, H))
The closer the RMSE is to zero, the better the alignment is. Zero RMSE means perfect alignment. That is, dMPNN is a constant multiple of damuee Note that the evaluation criterion is different from ALI, (Definition 4.1) for measuring the alignment between dMPNN and drune- There are multiple reasons for this. First, the RMSE in principle designed for non-binary dsuuu. Therefore, RMSE(dMPNN drune) is not a meaningful value when diunc a binary function, which is the case when the task is classification. Second, the computation of ALI (dMPNN, d is computationally too expensive. We explain this in terms of how many graph pairs we need to compute the distance for. Both the RMSE and ALI, require the calculation of the distance between 2D)2 pairs in the original definition. This is too demanding, especially when damu is dGED, which is NP-hard to compute. Therefore, in practice, we approximate the RMSE with 1000 randomly selected pairs from D2_ This kind of approximation is difficult for ALI.- To approximate ALI,, we first choose a subset Duub of D, and then compute danne of all pairs in D2 sub- Even if we set [Duubl = 100, which is quite small, we still need about 10 times more computation than the RMSE.
We evaluate four structural pseudometriiss graph edit distance (dGED, Sanfeliu & Fu, 1983), tree mover's distance (dIM Chuang & Jegelka, 2022), Weisfeiler Leman optimal assignment distance (dwLOA, Kriege et al., 2016), and Wasserstein Weisfeiler Leman graph distance (dwwL. Togninalli et al., 2019). See Appendix A. for detailed definitions. dIMD. dwLOA- and dwwr. are pseudometrics on the set of pairwise nonisomorphic graphs G. Only dGED for strictly positive edit costs is a metric. i.c., dG H) = 0 and only if G and H are isomorphic. We will also call dGED a pseudometric for simplicity. We chose dGED because is a popular graph pseudometric. The others were chosen because they are based on the message passing algorithm, like MPNNs, and classifiers based on their corresponding kernels were reported achieve high accuracy. In addition, dTMD has been theoretically proven be an upper bound of dMPNN (Chuang & Jegelka, 2022). Note that the exact calculation of dGED is in general Na-hard due to the combinatorial optimization over the set of valid transformation sequences (see Definition A. Therefore, in our experiment, we limit the computation time of dOED of each graph pair (G, H) to a maximum of 30 seconds. If this time limit is exceeded, we consider the lowest total cost at that point to be dEED (G, H). When we compute the RMSE between dMPNN and any of dTMD, dwLox. and dwWL. we set the depth of the computational trees used to compute these dstrue as the number of message passing layers in the MPNN for a fair comparison.
Figure 9 presents the distributions of the RMSE in different datasets (Morris al., 2020; Wu et al., 2018), darues and the readout functions used in MPNN. We followed exactly the same procedure for training and evaluating MPNNs as shown in Appendix C. Each distribution consists of RMSE( dsmme) of 24 MPNNs with different architectures and hyperparameters. We also provide results for untrained MPNNs to see the effect of training. As can be seen from the plots, the distributions of the untrained and trained MPNNs overlap, and there is no strong and consistent improvement in RMSE after training (answer to Q1.4). Regarding Q1.3, none of the four d performs best in all cases. The best one depends on the choice of dataset and pooling. One intersting observation is that dMPNN with sum pooling is more aligned with dwLOA than dWwL, while the reverse is true for dMPNN with mean pooling. This difference between pooling methods can be explained by different normalizations of the structural pseudometrics (see Section 5.5 and Appendix A.4).
Another insight from Figure 9 is that the degree of alignment between dMPNN and dame varies by model. To see if the alignment is crucial for the high predictive performance of MPNNs, we examined the SRC between RMSE(dMPNN. dsaree) of trained models and their performance on the training and test sets. We used accuracy and RMSE as performance criteria. Table 4 shows that the correlation is neither strong nor consistent across settings. Thus the alignment between dMPNN and damc not a key to high MPNN performance. This answers Q1.5 negatively.
20

---

WILTing Trees: Interpreting the Distance Between MPNN Embeddings
Mutagenicity
ENZYMES
Lipophilicity
0.30
0.25
0.20
0.25
0.20
R.M.
0.20
0.15
0.15
0.15
0.10
0.10
0.10
0.05
0.05
0.05
0.00
dern
dran
dwroa
dww1.
dGFn
drum
dW1.A
dwwT.
dGBD
dram
dwioA
dwwi.
IMDB-BINARY
COLLAB
0.6
0.4
untrained/man /men
RMSE
0.4
trained /mean untrained/sum
0.2
0.2
trained/sum
0.0
0.0
dGen
drap
dwnoa
dwwn.
dGen
dram
dwioA
dwwn.
Figure 9: The distributions of RMSE( dMPNN, dstrr) under different dstrue and datasets. Each color represents whether the MPNNs are trained not and which graph pooling function they use.
E. Experimental Details for Section
For the experiments in Section trained 3-layer GCN and GIN with embedding dimensions of 64 the three datasets. We explored both mean and pooling. Each model was trained on the full dataset for 100 epochs using the Adam optimizer with learning rate of 10-3 Then, each model was distilled to WILT by minimizing the loss L defined in Section 5.4. We used the entire data set for 'D in L. The distillation was done using gradient descent optimization with the Adam optimizer for 10 epochs. The learning rate and batch size were set to 10-2 and 256, respectively. See Algorithm for details.
In Table 2, we only show the results for GCN. Here, we show results for GIN in Table 5, The overall trend is the same between Tables and 5: dwur and dwnr much better aligned with dMPNN than dwwr. and dWLOx- In addition, dwwl. and dwnr approximate AMPNN(mean) better, while the opposite is true for dsm)) We also observed the same trend in the IMDB-BII BIDB-BINARY and COLLAB datasets (see Table 6).
Next, we plot the distribution of WILT edge weights after distillation in Figure 10. While the range of edge weights varies by model and dataset, all the distributions are skewed to zero (note that the y-axis is log scale). This suggests that only a small fraction all WL influence dMPNN. In other words, MPNNs build up their embedding space based on a small subset of entire WL colors, regardless of model and dataset.
Finally, visualize the WL colors with the largest weights, i.e., whose presence absence influence dwirr and therefore by approximation d'PNN the most. We use the Mutagenicity dataset as functionally important substructures are known from domain knowledge (Kazius et al., 2005). It should be noted that we only consider colors that appear in at least 1% of all graphs in the dataset. Table and show graphs with substructures corresponding to the WL colors with the top ten largest weights. Table is the result for GCN with sum pooling, while Table 8 is for GCN with mean pooling. the highlighted subgraph matches one the seven toxicophore substructures listed in Table Kazius al. (2005), we show
21

---

WILTing Trees: Interpreting the Distance Between MPNN Embeddings
Mutagenicity
ENZYMES
Lipophilleity
IMDB-BINARY
COLLAB
Frequency
GTNjmm
10
CBN
11
10
12
Weight
value
Figure 10: The distribution of edge weights of WILT after distillation from varying models trained on different datasets. The models with sum pooling were distilled into WILT with dummy normalization, while the models with mean pooling were distilled into WILT with size normalization. The log scale y-axis shared across all plots.
22

---

WILTing Trees: Interpreting the Distance Between MPNN Embeddings
Table 4: Correlation (SRC) between RMSE(dMPNN, deme) and the performance on the training and test sets. Performance was measured based on aceuracy for Mutagenicity and ENZYMES, and based on RMSE for Lipophilicity.
Train
Test
GED
TMD
WLOA
WWL
GED
TMD
WLOA
WWL
Mutagenicity
mean
0.33
0.19
-0.09
0.46
0.36
0.36
0.05
0.66
sum
0.09
0.25
0.38
0.20
0.10
0.31
0.14
0.09
mean
0.32
0.27
0.30
0.15
0.33
0.53
0.44
0.12
ENZYMES
sum
-0.36
0.71
0.44
0.35
-0.60
0.14
-0.25
-0.22
Lipophilicity
mean
-0.57
-0.61
-0.58
-0.46
-0.55
-0.77
-0.65
-0.65
sum
-0.12
-0.62
-0.54
-0.32
-0.53
-0.85
-0.85
-0.64
mean
0.04
0.10
-0.53
-0.20
-0.30
-0.31
-0.39
0.41
IMDB-BINARY
sum 0.37
0.65
-0.54
-0.51
0.01
0.22
-0.23
0.02
mean
0.60
0.58
0.57
-0.39
0.64
0.56
0.61
-0.42
COLLAB
sum
-0.40
0.63
-0.50
-0.54
-0.31
0.53
-0.39
-0.48
Table 5: The std of RMSE (dMPNN, d) [x 1002] over five different seeds. Each corresponds to GIN with a given graph pooling method, trained on a given dataset.
dwwr.
dwToOA
dWILT
dWILT
Mutagenicity
mean
11.47=0.24
17.99-2.79
3.70:10.57
4.98
sum
14.08=00.7
13.05+1.44
3.86++0.40
3.36+10.36
ENZYMES
mean
11.54-10.30
23.71+0081
5.32++0.20
7.55+
sum
12.10=0.84
9.94
8.60
3.86
Lipophilicity
mean 14.12-
16.55
6.31++0.66
9.52++0.70
sum
14.97=0.58
13.9710.75
0.49:10.50
6.5910.51
the toxicophore name as well. Four out of ten WL colors in Table correspond toxicophore substructures, while three out of ten in Table 8. These are quite lot considering that only seven toxicophore substructures are listed in Table Kazius al. (2005). Furthermore, there (2005). For instance, (6) and (9) in Table 7 and (8) in Table 8 partially match "aromatic nitro", while (7) in Table 8 is part of "polyeyclic aromatic system". Note that it is impossible to identify subgraphs that perfectly match these toxicophore substructures, since our method can only identify subgraphs corresponding to a region reachable within fixed steps from a root node. For example, the subgraph (1) of Table is a region reachable in 2 steps from the oxygen O. This limiation may seem to be a drawback of our proposed method, but fact it is not. natural identify only subgraphs corresponding WL colors to interpret dMPNN, because MPNNs can only see input graphs as multiset of WL colors.
colors that not fully but partially match of the substructures in Kazius et al.
23

---

WILTing Trees: Interpreting the Distance Between MPNN Embeddings
Table 6: The std of RMSE(d d) [x 1x10-2] over five different seeds. Each row corresponds to GCN or GIN with given graph pooling method, trained on IMDB-BINARY or COLLAB dataset.
dWWL
dWLOA
dWILT
dWILT
GCN mean 16.98+2.00
19.04:4.39
6.19=1124
7.62+1127 4.69+
sum
16.21+2.45
12.01+33.81
9.08++4.31
IMDB-BINARY
GIN
21.32::0.25 19.65:0.4
2.61
3.09
23.49#0.42
21.23+0.39
8.09=0.89
0.85
GCN
mean
16.7441.83
32.60+3.42
8.34
19.49
16.97+40.77
6.17+11226
2.21
2.03++0.41
COLLAB
GIN
mean
20.37:L0.73
13.31
3.56
10.54
sum
25.61
14.38#1.19
2.38
1.11#0.21
Table Example graphs with highlighted significant subgraphs corresponding colors with top largest weights. GCN with pooling was used. The toxicophore name is shown if the highlighted subgraph matches toxicophore substructures reported in Table 1 of Kazius et al. (2005)
(I) three-membered
heterocycle
(2)
(3)
(4) alphatic halide
(5)
(epoxide)
NO
C3
(6) nitroso
(9)
(10) alphatic halide
096
(36
Grownn
C3
CH,43
C1
GB
24

---

WILTing Trees: Interpreting the Distance Between MPNN Embeddings
Table Example graphs with highlighted significant subgraphs corresponding to colors with top 10 largest weights. GCN with mean pooling was used. The toxicophore name is shown if the highlighted subgraph matches toxicophore substructures reported in Table 1 of Kazius et al. (2005)
(2) three- membeed
heterocycle
(3) alphatic halide
(4)
(5)
(epoxide)
D6C
CHN 066
C3
NO.
(6)
(7)
(8) nitroso
(9)
(10)
CB
Ch
25


Reviews:
[{"review_id": "aA9FuCBjYf", "reviewer": "Reviewer_n2hL", "summary": "The paper investigates (i) the properties of the distances defined for the MPNN based on structural and functional pseudometrics to find the one that explains the high performance of MPNN, and (ii) how MPNNs learn such a structure. The main contribution is the new graph distance based on the Weisfeiler Leman labeling tree (WILT) that is weighted. The nodes in WILT are the colors of the Weisfeiler Leman (WL) test, and the edges connect the parent color with the color that it changes to in the iterations of WL. \nThe edge weights in WILT are learned such that the graph distance is close to the graph embedding of the message passing neural networks (MPNN). Using this, the authors claim that the edge weights identify the subgraphs that strongly influence the distance between the embeddings of MPNN and use this to interpret the embedding space.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "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": "Thank you for reviewing our paper, including the supplementary materials, and acknowledging our contributions. We address each weakness and question individually below.\n\n---\n> Complexity of building the tree\n\nSuppose the dataset consists of $|D|$ graphs with at most $|E|$ edges. Then, building WILT with depth $L$ takes $O(|D| \\times L \\times |E|)$ time. Thus, the complexity is linear in $|D|$ and $|E|$. If the graphs are sparse, i.e. $|V| \\propto |E|$, the complexity is also linear in $|V|$. For dense graphs, the complexity is quadratic in $|V|$.\n\n> The extracted subgraph needs to be compared to the existing methods such as GNNExplainer [1], SubgraphX [2], etc.\n\nMost interpretation methods, including [1] and [2], aim to find a subgraph that GNN considers important for a prediction for ONE given graph (instance level). On the other hand, our method tries to explain the GLOBAL behavior of GNN by analyzing $d_\\mathrm{MPNN}$. The highlighted subgraphs in Figure 6 do not mean important subgraphs for the four specific graphs. They are the ones that strongly influence $d_\\mathrm{MPNN}$ in general. Because of the fundamental difference in purpose, it is difficult to compare our method with [1] and [2]. There are also global explanation methods as mentioned in Section 2, but they use decision trees or logical formulas as a language, which are difficult to compare with the subgraphs extracted by our method.\n\n> Essential References Not Discussed: [1] studies empirically and theoretically the alignment of the embedding to the task related functional metric.\n\nThank you for bringing this important prior study to our attention. We will add it to the Related Work section in our final version. At this point, we would like to clarify the relationship between our study and [1]. While both investigate the alignment between the embedding distance and the functional distance, there are fundamental differences: Our study examines the graph embedding space of GNNs trained on practical graphs, while [1] analyzes the node embedding space of GNNs trained on graphs generated from a stochastic block model. Nevertheless, both studies reach a similar conclusion: GNNs are trained such that the embeddings respect the functional alignment. It would be interesting to investigate how the degree of alignment between the graph embedding distance and the functional distance changes with layer depth, as [1] did for node embedding.\n\n> Weakness 1: Unclear caption for Figure 1\n\nWe apologize for the confusion. $-$ and $\\cdots$ represent edges with different labels. We consider a variant of the WL algorithm that takes edge labels into account. This allows WILT to handle graphs with edge labels, such as the molecules in the Mutagenicity dataset. We will clarify the caption in the final pdf.\n\n> Weakness 2: The alignment to structural pseudometric is evaluated in the appendix.\n\nWe will include the alignment to $d_\\mathrm{struc}$ in the main text in the final pdf.\n\n> Weakness 3: Discussion on the limitations of the method is missing.\n\nOne limitation of our study is that we only distilled two GNN architectures (GCN and GIN) with fixed hyperparameters to WILT. Thus, it remains to be seen how different architectures and hyperparameters affect the edge weights of WILT. Another limitation is that our method cannot be directly applied to the analysis of GNNs for node classification on large graphs such as social networks or citation networks. This is because we only deal with graph-level tasks. However, we believe that the node embedding space of GNNs can also be distilled to WILT in a similar way, i.e., by tuning the weights to approximate the node embedding distance with a path distance on WILT. We will include these limitations and possible future work in the final pdf.\n\n> 'Figure 5' in line 382 second column shouldn't it be Figure 4? Similarly 'Figure 4' in line 411 first column should be Figure 5?\n\nWe are sorry that these are mistakes. They will be fixed in the final pdf.\n\n> Why do you have $\\alpha$ in the RMSE definition?\n\nWe use $\\mathrm{RMSE}(d_\\mathrm{MPNN}, d)$ to measure how well $d$ captures the geometric structure of $d_\\mathrm{MPNN}$. We first normalize $d_\\mathrm{MPNN}$ so that we can compare or compute the mean of the RMSEs of different MPNNs. The normalization of $d$ is optional as we optimize $\\alpha$ anyway. $\\alpha$ reflects our belief that the scale of the distance does not affect the importance of the subgraphs. In other words, if $d = \\alpha \\cdot d_\\mathrm{MPNN}$, we consider that $d$ correctly captures $d_\\mathrm{MPNN}$. Since $d_\\mathrm{WILT}$ is optimized to approximate $d_\\mathrm{MPNN}$, $\\alpha$ is close to one. However, for $d_\\mathrm{WWL}$ and $d_\\mathrm{WLOA}$, $\\alpha$ is not necessarily close to one. We introduce $\\alpha$ for a fair comparison between $d_\\mathrm{WILT}$ and $d_\\mathrm{WWL}$, $d_\\mathrm{WLOA}$.\n\n---\nPlease let us know if you have any further questions or suggestions, and we will be happy to answer them.", "claims_and_evidence": "There are two main claims made in the paper based on the analysis: (i) MPNN distance defined on the embedding space is critical to the task performance, and (ii) develops a new distance based on WILT and an algorithm to learn its edge weights to be close to the embedding of MPNN. The edge weights allow to identify the subgraph that yields high performance, thereby providing a way to interpret the embedding space of MPNN.\n\nClaim (i) is expected and intuitive. It is however established clearly using the structural and functional pseudometrics, and evaluation of the alignment to the MPNN embedding on several datasets. \n\nWhile claim (ii) is evaluated, the computational aspect of the algorithm is not discussed in detail. I understand $d_{WILT}$ is linear in time. But, for a dataset with $n$ graphs with at least $m$ nodes, one needs to apply the WL test algorithm for all the graphs and then build the WILT traversing the resultant graphs from the WL test. What is the complexity of building this tree? and what about its scalability, especially with the number of graphs in the dataset and their sizes?", "methods_and_evaluation_criteria": "The proposed method is novel and the evaluation is sensible. That being said, the purpose of the proposed graph distance is not clear. Is interpretability the primary objective? Then, to strengthen the evaluation, I think the extracted subgraph needs to be compared to the existing methods such as GNNExplainer [1], SubgraphX [2], etc.  \n\n[1] Ying et al. GNNExplainer: Generating Explanations for Graph Neural Networks. NeurIPS 2019\n\n[2] Yuan et al. On explainability of graph neural networks via subgraph explorations. ICML 2021.", "theoretical_claims": "The expressivity theorems appear correct. The proofs in the appendix are not checked thoroughly.", "experimental_designs_or_analyses": "The experiments are sound and detailed. The plots are especially helpful in conveying the idea and results.", "relation_to_broader_scientific_literature": "The key contribution is the WILT based graph distance. The subgraph identified from its edge weights is shown to be functionally important, which is interesting for the interpretability in MPNNs.", "essential_references_not_discussed": "Regarding the first investigation on the properties of the MPNN distances that explain its high performance, [1] studies empirically and theoretically the alignment of the embedding to the task related functional metric (adopting the language from the current paper). This work needs to be discussed. \n\n[1] Kothapalli et al. \"A neural collapse perspective on feature evolution in graph neural networks.\" NeurIPS 2023", "other_strengths_and_weaknesses": "**Strengths**\n\nThe paper is clearly written, and the claims are supported by experiments.\n\nI found the discussion on the pseudometrics helpful.\n\nThe expressivity of WILT is also analyzed which further strengthens the contribution. \n\n**Weaknesses**\n\nThe clarity can be improved in some aspects. For instance, \n\n1. The caption of Figure 1 last line is not clear. What is the tuple in the multi-set with $-$ mean? \n\n2. The alignment to structural pseudometric is evaluated in the appendix. It would improve the draft if it is mentioned and referenced in Section 4 as well. Perhaps, the authors can reference Figure 9 and contrast it with Figure 2.\n\n3. Discussion on the limitations of the method is missing.", "other_comments_or_suggestions": "'Figure 5' in line 382 second column shouldn't it be Figure 4? Similarly 'Figure 4' in line 411 first column should be Figure 5?", "questions_for_authors": "Please check the other sections as well for questions. \n\nOne other question: Why do you have $\\alpha$ in the RMSE definition? Is it somehow influencing the lower distance for WILT compared to WWL and WLOA?", "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": "", "extra_scores": {}}, {"review_id": "NohSn2y76b", "reviewer": "Reviewer_PD52", "summary": "The authors of this paper investigate the distance function learned by message-passing neural networks (MPNNs) and introduce a new framework called Weisfeiler Leman Labeling Tree to interpret these distances. Unlike previous work that aligns MPNN embeddings with structural graph distances, the authors focus on task-specific functional distances. Their key contribution is the introduction of the WILT framework, which applies optimal transport techniques to a weighted Weisfeiler-Leman (WL) tree, enabling efficient computation and interpretation of MPNN distances. They demonstrate through experiments in graph classification datasets, that MPNN embeddings capture functionally relevant subgraphs, leading to improved interpretability and performance understanding.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "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": "We would like to thank you for your thorough review and valuable feedback. We hope the points below adequately address your concerns.\n\n---\n> Weakness 1: How different GNN architectural choices affect the learned embedding distances and the performance of WILTing Trees. For example, how does the number of layers affect the results? Is there a connection with over-smoothing?\n\nSince we put more focus on trying out different datasets than on changing GNN architectures, we used only two types of GNNs (GCN and GIN) and poolings (sum and mean). Thus, we admit that the effect of the choice of architecture and hyperparameters is still unclear. At this point, we would like to share some insights on how hyperparameters affect the learned embedding distances. We observed that the deeper the GNN is, the more likely its embedding space is aligned with the functional distance, similar to the observation of [1] Kothapalli et al. mentioned by reviewer n2hL. For example, in Figure 8, deeper models tend to be plotted on the right side of each graph. The connection to over-smoothing is an interesting future work that we did not address during our investigation.\n\n> Weakness 2, Question 1: Does WILT extend to more expressive GNNs? If so, have you conducted experiments on them?\n\nAs briefly mentioned in the conclusions, it is possible to construct WILT for higher-order variants of the WL labeling algorithm. We expect that such a higher-order WILT can be used in a similar way to analyze GNNs whose expressiveness is only bounded by higher-order WL variants. However, we have not yet empirically explored this direction.\n\n> Question 2: Do you see any way to improve GNN performance using insights from WILT, besides the interpretability aspect?\n\nThis is a very important question. At the moment we don't have an effective way to improve GNN's performance using insights from WILT. However, WILT has a nice property that might help us improve GNN's performance: The tree structure of WILT is determined only by the dataset and isn't influenced by the GNN. Therefore, different GNNs can be distilled into the same WILT. By comparing the resulting edge weights, we may be able to gain insight into why some GNNs perform better than others. This is really a promising future direction.\n\nAlthough not fully related to your question, we would like to mention the possibility of using WILT to achieve better performance in graph classification or regression. In this study, we train WILT weights by approximating the GNN embedding distance. However, if we find a way to effectively train the weights from scratch, the resulting $d_\\mathrm{WILT}$ can be used as a kernel function that is more flexible than other WL-based kernels at no additional computational cost.\n\n---\nIf you have any other questions or suggestions, please let us know and we'll be happy to address them.", "claims_and_evidence": "The claims are supported by convincing evidence.", "methods_and_evaluation_criteria": "The proposed methods and evaluation criteria are suitable for the tasks.", "theoretical_claims": "The paper makes several theoretical claims, regarding the properties of the proposed WILT distance, its relationship to existing graph distances, and its alignment with MPNN embeddings. The claims are easy to follow and appear correct.", "experimental_designs_or_analyses": "The experimental design seems valid.", "relation_to_broader_scientific_literature": "The paper shifts focus from the binary expressiveness of MPNNs (i.e., distinguishing non-isomorphic graphs) to the metric structure of their embeddings. It extends previous theoretical analyses by exploring functional distance alignment (rather than using only structural distance alignment) .", "essential_references_not_discussed": "The paper provides a detailed  discussion of related work.", "other_strengths_and_weaknesses": "## Strong Points:\n\n- The paper studies the distance metric learned by MPNNs which is crucial for the interoperability of those models.\n\n- The paper moves beyond binary expressivity and addresses how MPNNs learn distance metrics that impact predictive performance.\n\n- WILT can be computed in linear time, making it practical for large-scale graph datasets.\n\n## Weak Points: \n\n-  The paper provides good motivation and theoretical intuition for understanding the learned distance between MPNN embeddings. However, I feel that the experimental results are limited. One of the key limitations is the lack of a systematic ablation study examining how different GNN architectural choices affect the learned embedding distances and the performance of WILTing Trees. For example, how does the number of layers affect the results? Is there a connection with oversmoothing?\n\n- The authors only use GCN and GIN in the experiments. However, many modern GNNs exceed 1-WL expressivity (e.g., higher-order message passing, subgraph-based architectures, additional positional features, graph transformers). It is unclear whether the current findings extend to these more powerful architectures,", "other_comments_or_suggestions": "NA", "questions_for_authors": "- Does WILT extend to more expressive GNNs? If so, have you conducted experiments on them?\n\n- Do you see any way to improve GNN performance using insights from WILT, besidesπ the interpretability aspect?", "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": "", "extra_scores": {}}, {"review_id": "DlmOBpIhCn", "reviewer": "Reviewer_gzCR", "summary": "This paper investigates how MPNNs learn to embed graphs in a way that captures functional relationships between them. The authors propose a novel interpretable graph distance, the WILTing distance, which effectively approximates the distances between MPNN embeddings. They demonstrate that MPNNs focus on a small subset of WL colors that are functionally important.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "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": "Thank you for your review and valuable feedback, which help us clarify our arguments and consider future directions. Below are our answers to your questions and comments.\n\n---\n> Why some MPNNs achieve better functional alignment than others\n\nSince we focused more on trying out different datasets, we have not conducted in-depth experiments to compare different GNNs. At this point, we share some insights on how hyperparameters affect the functional alignment. We observed that the embedding space of deeper GNNs is more aligned with the functional distance. For example, in Figure 8, deeper models tend to be plotted on the right side of each graph. In future studies, we plan to investigate in more detail what is a key factor contributing to higher functional alignment and thus higher performance. WILT has a nice property that may help in this investigation: The tree structure of WILT is determined only by the dataset and isn't influenced by the GNN. Therefore, different GNNs can be distilled into the same WILT, allowing a direct comparison of the resulting edge weights.\n\n> Comparison against other interpretability methods in GNN\n\nMost interpretation methods aim to find a subgraph that GNN considers important for a prediction for ONE given graph (instance level). On the other hand, our method tries to explain the GLOBAL behavior of GNN by analyzing $d_\\mathrm{MPNN}$. The highlighted subgraphs in Figure 6 do not mean important subgraphs for the four specific graphs. They are the ones that strongly influence $d_\\mathrm{MPNN}$ in general. This fundamental difference in purpose makes it difficult to compare our method with existing methods. There are also global explanation methods as mentioned in Section 2, but they use decision trees or logical formulas as a language, which are difficult to compare with the subgraphs extracted by our method.\n\n> How the approach scales to very large graphs or datasets with many graphs\n\nWe expect WILT to scale to large graphs or datasets, although we have not conducted actual experiments. The number of WILT colors is bounded by the number of nodes in all graphs times the WILT depth. In practice, however, it is much smaller because nodes with the same neighborhood are assigned the same color. As for the computational time, building the WILT and computing the embeddings require linear time in the total number of edges in the dataset. Once the embeddings are prepared as sparse vectors, the computation of $d_\\mathrm{WILT}$ is always linear in the number of nodes of the two graphs involved.\n\nOur method is designed for graph-level tasks and cannot be directly applied to node classification, which is a typical task on very large graphs such as social networks. However, we believe that the node embedding space of GNNs can also be distilled to WILT in a similar way, i.e., by tuning the weights to approximate the node embedding distance with a path distance on WILT.\n\n> Q1: Why the Pearson correlation coefficient is used as a measure of correlation\n\nWe used Pearson's correlation coefficient (PCC) because it is commonly used and preferable in terms of visualization (see red lines in Figure 8). However, you are right in that there is no valid reason to expect a linear relationship, so we have reanalyzed the results using Spearman's rank correlation coefficient (SRCC). The result is similar to when PCC is used: the functional alignment is more consistently and highly correlated with performance than the structural alignment. Here is the result for the Mutagenicity dataset. We will include these results in the final pdf.\n\n- SRCC between the functional alignment and accuracy (c.f. Table 1)\n|| train | test |\n|-|-|-|\n|k=1|0.66|0.70|\n|k=5|0.65|0.69|\n|k=10|0.63|0.68|\n|k=20|0.61|0.67|\n- SRCC between the structural alignment and accuracy (c.f. Table 4)\n||train/GED|train/TMD|train/WLOA|train/WWL|test/GED|test/TMD|test/WLOA|test/WWL|\n|-|-|-|-|-|-|-|-|-|\n|mean|0.33|0.19|-0.09|0.46|0.36|0.36|0.05|0.66|\n|sum|0.09|0.25|0.38|0.20|0.10|0.31|0.14|0.09|\n\n> Q2: I wonder whether this metric is effective for investigating the internal embeddings of MPNNs, rather than just the final embeddings, as previous work[1] (also based on the WL kernel) suggests that consistency in the distance relationship is crucial for MPNN performance.\n\nIn principle, $d_\\mathrm{WILT}$ can be trained to approximate the MPNN embedding distance at internal MPNN layers. Our conjecture is that the deeper the layer is, the more the MPNN embedding distance (and thus $d_\\mathrm{WILT}$) respects the functional distance. $d_\\mathrm{WILT}$ may also keep the consistency proposed in [1], because $d_\\mathrm{WILT}$ includes $d_\\mathrm{WLOA}$ as a special case, which has been shown theoretically to preserve consistency. However, the proofs provided in [1] require the WILT weights to be constant and hence the results do not immediately extend to our case.\n\n---\nPlease feel free to let us know if you have any other questions or suggestions.", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "Yes", "theoretical_claims": "The paper offers a solid theoretical foundation for understanding MPNN embedding spaces through the lens of metric learning. However, it does not fully explain why some MPNNs achieve better functional alignment than others.", "experimental_designs_or_analyses": "The paper provides comprehensive experiments across multiple datasets and a clear demonstration that the alignment between MPNN distances and functional distances correlates with model performance. However, \n1. the authors compare their method only to existing graph kernels and do not compare it against other interpretability methods in GNN, making it difficult to assess the relative advantages of WILT. \n2. the experiments do not address how the approach scales to very large graphs or datasets with many graphs, which is crucial for practical applications.", "relation_to_broader_scientific_literature": "No", "essential_references_not_discussed": "No", "other_strengths_and_weaknesses": "N/A", "other_comments_or_suggestions": "N/A", "questions_for_authors": "Q1: I am curious about why the Pearson correlation coefficient is used as a measure of correlation. Is there a linear relationship, or would another method, such as rank correlation, be better for quantifying this relationship?\n\nQ2: I wonder whether this metric is effective for investigating the internal embeddings of MPNNs, rather than just the final embeddings, as previous work[1] (also based on the WL kernel) suggests that consistency in the distance relationship is crucial for MPNN performance.\n\n\n[1]Liu, et al Exploring Consistency in Graph Representations: from Graph Kernels to Graph Neural Networks. NIPS", "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": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The paper proposes an optimal transport based distance between MPNN embeddings. In contrast to existing "structural" distances between embeddings, the proposed distance is task-specific and hence, it is useful to identify relevant subgraphs for a task. All reviewers agree that the theoretical contribution of the work is sound and interesting. Some reviews note that the experiments could be improved, and the author response also shows that there is possibility of further investigation of the proposed WILT framework. Hence, it is an interesting contribution.
Decision: Accept (poster)