From Relational Pooling to Subgraph GNNs: A Universal Framework for More
  Expressive Graph Neural Networks
Relational pooling is a framework for building more expressive and
permutation-invariant graph neural networks. However, there is limited
understanding of the exact enhancement in the expressivity of RP and its
connection with the Weisfeiler Lehman hierarchy. Starting from RP, we propose
to explicitly assign labels to nodes as additional features to improve
expressive power of message passing neural networks. The method is then
extended to higher dimensional WL, leading to a novel $k,l$-WL algorithm, a
more general framework than $k$-WL. Theoretically, we analyze the expressivity
of $k,l$-WL with respect to $k$ and $l$ and unifies it with a great number of
subgraph GNNs. Complexity reduction methods are also systematically discussed
to build powerful and practical $k,l$-GNN instances. We theoretically and
experimentally prove that our method is universally compatible and capable of
improving the expressivity of any base GNN model. Our $k,l$-GNNs achieve
superior performance on many synthetic and real-world datasets, which verifies
the effectiveness of our framework.
Cai Zhou, Xiyuan Wang, Muhan Zhang
# From Relational Pooling to Subgraph GNNs:

###### Abstract

Relational pooling (RP) is a framework for building more expressive and permutation-invariant graph neural networks (GNN). However, there is limited understanding of the exact enhancement in the expressivity of RP and its connection with the Weisfeiler-Lehman (WL) hierarchy. Starting from RP, we propose to explicitly assign labels to nodes as additional features to improve graph isomorphism distinguishing power of message passing neural networks. The method is then extended to higher-dimensional WL, leading to a novel \(k,l\)-WL algorithm, a more general framework than \(k\)-WL. We further introduce the subgraph concept into our hierarchy and propose a localized \(k,l\)-WL framework, incorporating a wide range of existing work, including many subgraph GNNs. Theoretically, we analyze the expressivity of \(k,l\)-WL w.r.t. \(k\) and \(l\) and compare it with the traditional \(k\)-WL. Complexity reduction methods are also systematically discussed to build powerful and practical \(k,l\)-GNN instances. We theoretically and experimentally prove that our method is universally compatible and capable of improving the expressivity of any base GNN model. Our \(k,l\)-GNNs achieve superior performance on many synthetic and real-world datasets, which verifies the effectiveness of our framework.

Machine Learning, Knowledge-theoretic, Graph Neural Networks, Graph Neural Networks

## 1 Introduction

Graph-structured data has recently revealed a significant importance in many fields, including bio-informatics, combinatorial optimization and social-network analysis, among which graph neural networks (GNNs) achieve great successes (Bronstein et al., 2016; Klicpera et al., 2020; Dai et al., 2017). Message passing neural network (MPNN) is one of the simplest and most commonly used GNNs (Zhou et al., 2018), whereas its expressivity in distinguishing non-isomorphic graphs is bounded by the one-dimensional Weisfeiler-Lehman test (1-WL) (Xu et al., 2018; Morris et al., 2018). Therefore, designing GNNs with stronger expressivity has aroused increasing attention.

Numerous approaches have been proposed to enhance GNN's expressivity. Relational Pooling (RP) (Murphy et al., 2019; Chen et al., 2020) is a framework to build powerful permutation-invariant models by symmetrizing expressive permutation-sensitive base models. Concretely, RP first feed adjacency matrix to a powerful permutation-sensitive model, like Multi-Layer Perceptron (MLP), to achieve strong expressivity. Then the permutation invariance is guaranteed by averaging or summing over representations under all permutations of node IDs (hence all permutations of adjacency matrix). However, RP is impractical for most real-world graphs due to the \(O(n!)\) complexity, where \(n\) is the number of nodes. Based on RP, Chen et al. (2020) further introduce a local version called Local Relational Pooling (LRP), which performs permutation and averaging within an induced sub

Figure 1: The expressivity hierarchy of \(k,l\)-WL. The blue arrows indicate Theorem 5.5 and Theorem 5.6, showing that increasing \(k\) and \(l\) will strictly increase expressivity. The yellow arrows imply Theorem 5.7, which states that \(k+1,l\)-WL is strictly more powerful than \(k,l+1\)-WL when \(k\geq 2\).
graph. LRP's time complexity is reduced to the number of subgraphs \(O(n^{l})\), where \(l\) is the subgraph size. But so far, there is a lack of theoretical analysis on the expressivity of LRP. In this paper, we propose ID-MPNN, a variation of LRP that avoids the time complexity of RP. Instead of using MLP as the base encoder, ID-MPNN runs MPNN on the whole graph. Meanwhile, to improve expressivity, \(l\) nodes in the whole graph are labeled with \(1,2,...,l\). Through the lens of ID-MPNN, we establish a connection between (local) Relational Pooling and subgraph GNNs.

Furthermore, by replacing MPNN with more powerful base encoders \(k\)-WL, we propose \(k,l\)-WL, a universal framework for many expressive GNNs (shown in Figure 2). Intuitively, \(k,l\)-WL can be viewed as running \(k\)-WL on a graph with \(l\) nodes labeled and symmetrizing over \(l\). We theoretically analyze the expressivity of \(k,l\)-WL and build a complete expressivity hierarchy of the algorithms with different \(k,l\) as shown in Figure 1. As a universal framework, \(k,l\)-WL incorporates a wide range of existing algorithms and GNN models, including relational pooling, the original \(k\)-WL, many subgraph GNNs, and some other GNN extensions.

In summary, the organization of this paper and our main contributions are as follows.

1. Section 4 proposes ID-MPNN to improve LRP. ID-MPNN is further extended to a general framework \(k,l\)-WL, which incorporates a majority of existing WL and GNN variations, including RP and many subgraph GNNs.
2. Section 5 theoretically analyzes the algorithm's expressivity and builds a strict \(k,l\)-WL expressivity hierarchy, which is more general than the \(k\)-WL hierarchy.
3. Section 6 discusses practical issues in our \(k,l\)-WL framework and proposes techniques to improve scalability.
4. Section 7 evaluates \(k,l\)-WL with extensive experiments on both synthetic and real-world datasets. Our models achieve state-of-the-art results on several tasks and significantly outperforms previous works based on RP.

## 2 Related Work

Graph Neural Network and Weisfeiler-Lehman test

Weisfeiler-Lehman tests are a classical family of algorithms to distinguish non-isomorphic graphs. Previous works have built connections between the expressivity of GNNs and WL hierarchy (Xu et al., 2018; Frasca et al., 2022; Morris et al., 2019, 2020). We propose \(k,l\)-WL hierarchy, which is finer than \(k\)-WL and covers a wide range of existing models.

Subgraph GNNsSubgraph GNNs encode a set of subgraph instead of the original graph for graph representation learning. Through careful designs, they can have both strong expressivity and good scalability. Many subgraph GNNs sample a subgraph for each node (You et al., 2021; Zhang et al., 2021; Bouritsas et al., 2023; Sun et al., 2021; Zhao et al., 2021). Frasca et al. (2022); Zhang et al. (2023) upper bound the expressivity of these subgraph GNNs by 3-WL. \(I^{2}\)-GNN extracts a subgraph for each connected node pair and boosts the cycle counting power (Huang et al., 2022). Qian et al. (2022) further extract a subgraph for each \(l\)-tuple of nodes and propose \(l\)-OSAN, which is equivalent to our ID-MPNN and \(2,l\)-WL. We propose a more general framework \(k,l\)-WL to incorporate most existing subgraph GNNs.

## 3 Preliminary

Given an undirected graph \(G=(V,E,X)\), where \(V,E\) are the node set and edge set respectively, and \(X_{i}\) is the node feature of node \(i\), let \(N(v,G)=\{u\in V|(u,v)\in E\}\) denote the set of neighbors of node \(v\) in graph \(G\). Let \([n]\) denote the set \(\{1,2,...,n\}\). Given \(k\)-tuple \(\mathbf{a}\in V^{k}\) and \(l\)-tuple \(\mathbf{b}\in V^{l}\), let \(\mathbf{a}|\mathbf{b}\) denote a \(k+l\)-tuple, the concatenation of \(\mathbf{a}\) and \(\mathbf{b}\). Let \(\mathbf{a}_{i}\) denote the \(i\)-th element in tuple \(\mathbf{a}\), \(\psi_{i}(\mathbf{a},u)\) denote a tuple produced by replacing \(\mathbf{a}_{i}\) with \(u\). Let \(\mathbf{a}_{a:b}\) denote the slice of tuple \(\mathbf{a}\) containing \(a,a+1,...,b-1\)-th elements, where \(a\) is omitted if \(a=1\), \(b\) is omitted if \(b=|\mathbf{a}|+1\), and \(|\mathbf{a}|\) is the length of \(\mathbf{a}\).

Weisfeiler-Lehman test (1-WL) is a common graph isomorphism test, which also bounds the expressivity of message passing neural networks (MPNNs) (Xu et al., 2018). It assigns a color \(c_{1}^{0}(v,G)\) to each node \(v\) in graph \(G\) initially according to \(X_{v}\). If the graph has no node feature, the colors of all nodes are the same. Then, 1-WL iteratively updates the node colors. The \(t\)-th iteration is as follows.

\[c_{l}^{t}(v,G)\!=\!\text{Hash}(c_{l}^{t-1}\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\! \!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\! \!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\! \!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\! \!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\! \!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\! \!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\! \!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\! \!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\! \!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\! \!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\! \!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\! \!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\! \!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\! \!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\! \!\!
where \(c_{1}^{t}(v,G)\) is the color of node \(v\) at the \(t\)-th iteration. The color of \(v\) is updated by its original color and the colors of its neighbors. The color of the whole graph is the multiset of the node colors

\[c_{1}^{t}(G)=\text{Hash}(\{\{c_{1}^{t}(v,G)|v\in V(G)\}\}). \tag{2}\]

There still exist non-isomorphic graphs that 1-WL cannot differentiate. \(k\)-dimensional Weisfeiler-Lehman test has stronger expressivity. It assigns colors to all \(k\)-tuples and iteratively updates them. The initial color \(c_{k}^{0}(\mathbf{v},G)\) of tuple \(\mathbf{v}\in V(G)^{k}\) is determined by the isomorphism type of tuple \(\mathbf{v}\)(Maron et al., 2019) (see Appendix E). At the \(t\)-th iteration, the color updating scheme is

\[c_{k}^{t}(\mathbf{v},G)=\text{Hash}(c_{k}^{t-1}(\mathbf{v},G),(\{\{\] \[c_{k}^{t-1}(\psi_{i}(\mathbf{v},u),G)|u\in V(G)\}\}|i\in[k])), \tag{3}\]

where \(\psi_{i}(\mathbf{v},u)\) means replacing the \(i\)-th element in \(\mathbf{v}\) with \(u\). The color of \(\mathbf{v}\) is updated by its original color and the color of its high-order neighbors \(\psi_{i}(\mathbf{v},u)\). The color of the whole graph is the multiset of all tuple colors,

\[c_{k}^{t}(G)=\text{Hash}(\{\{c_{k}^{t}(\mathbf{v},G)|\mathbf{v}\in V(G)^{k}\}\}). \tag{4}\]

Note that \(k\)-WL (\(k\geq 2\)) takes a different form from \(1\)-WL. Our discussion mainly focuses on \(k\geq 2\) cases. Since \(1\)-WL has the same expressivity as \(2\)-WL, we can directly apply the conclusion of \(2\)-WL to \(1\)-WL.

## 4 \(k,l\)-WL: A Universal Framework

### Message passing with labels: enhancement by asymmetry

The expressivity of models built by Relational Pooling (RP) depends on the power of the base encoder before symmetrization. Some previous works use MLP (Chen et al., 2020) and RNN (Huang et al., 2022) to capture relations between nodes. They have high expressivity but little inductive bias for graph data. Moreover, in practical settings, where Local Relational Pooling (LRP) is used, they can only encode induced subgraphs and lose the global information of graph. To solve these problems, we introduce asymmetry to MPNN by assigning nodes unique labels (which are additional features, different from the node indices only to name different nodes) and use _MPNN with labels_ as the base encoder.

Given an input graph \(G\), MPNN with labels first assigns label \(i\) (node ID) to each node \(i\) as an additional feature and then runs standard message passing on the labeled graph. MPNN with full labels is expressive enough to encode the full graph information: MPNN can encode the multiset of neighbors into node representations. With node ID labels, each node's representation can identify the neighboring nodes connected to it. Therefore, the representation of the whole graph can identify the connectivity between nodes in the whole graph and thus enable distinguishing non-isomorphic graphs. Moreover, the standard message passing introduces inductive bias on graph data. MPNN with labels can also be easily adapted to the LRP setting. Instead of assigning all nodes unique labels, MPNN with labels can assign \(1,2,..,l\) to only \(l\) nodes and run message passing on the whole graph. Therefore, MPNN with labels can still capture global graph feature, unlike standard LRP only taking induced subgraphs as input.

Our ID-MPNN combines MPNN with labels with LRP. An ID-MPNN parameterized by \(l\) (called \(l\)-IDMPNN) explicitly assigns \(l\)**unique** labels (IDs) to \(l\) nodes (can be duplicated, thus \(n^{l}\) labeled graphs in total) as an additional feature. Then, a standard message passing is performed on each labeled graph. Finally, the representations of these labeled graphs are aggregated to produce the original graph representation. A contemporary work by Qian et al. (2022) also proposes similar models. However, they neither connect ID-MPNN with LRP nor extend ID-MPNN to a more general framework \(k,l\)-WL as in the following section.

### \(k,l\)-WL: enhancement by higher dimension

So far, it is natural to ask: what if we replace the MPNN with other more powerful GNNs? Equivalently, can 1-WL be replaced by higher-dimensional WL tests on the labeled graphs? When all nodes are assigned with unique labels, even MPNN can distinguish all non-isomorphic graphs and thus using more powerful models is meaningless. However, when the number of nodes labeled is fixed, we give a positive answer: if we run \(k\)-WL (\(k\geq 3\)) with \(l\) labels, it will be more powerful than \(1\)-WL (with \(l\) labels). We name running \(k\)-WL on labeled graphs with \(l\) IDs as \(k,l\)-WL, which is formally defined as follows.

1. Given an \(l\)-tuple of nodes \(\mathbf{v}\) in graph \(G=(V,E,X)\), the labeled graph is \(G^{\mathbf{v}}=(V,E,X^{\mathbf{v}})\), where \(\forall u\in V,X_{u}^{\mathbf{v}}=\text{Hash}\left(X_{u},\{[i|\mathbf{v}_{i}=u,i\in[l] \}\right)\}\). In other words, node \(\mathbf{v}_{i}\) will have an extra label \(i\).
2. \(k,l\)-WL then runs \(k\)-WL on each labeled graph \(G^{\mathbf{v}}\). * \(c_{k}^{0}(\mathbf{u},G^{\mathbf{v}})\), the color of \(k\)-tuple \(\mathbf{u}\) in graph \(G^{\mathbf{v}}\) is initialized by the isomorphism type of \(\mathbf{u}\) in \(G^{\mathbf{v}}\). * The tuple color at the \(t\)-th iteration: \[c_{k}^{t}(\mathbf{u},G^{\mathbf{v}})=\text{Hash}(c_{k}^{t-1}(\mathbf{u},G^{\mathbf{v}}),\] \[(\{c_{k}^{t-1}(\psi_{i}(\mathbf{u},w),G^{\mathbf{v}})|w\in V\})|i\in[k])).\] (5)
* The full graph color at the \(t\)-th iteration: \[c_{k}^{t}(G^{v})=\text{Hash}(\{\{c_{k}^{t}(\mathbf{u},G^{v})|\mathbf{u}\in V^{k}\}\}).\] (6)
3. The color of the whole graph is produced by aggregating the colors of all labeled graphs. \[c_{k,l}^{t}(G)\!=\!\text{Hash}\big{(}\{\{c_{k}^{t}(G^{v})|\mathbf{v}\in V^{l}\}\} \big{)}.\] (7)

Here, we briefly explain the \(k,l\)-WL algorithm, a more general and powerful form of \(k\)-WL. The key difference between \(k,l\)-WL and traditional \(k\)-WL lies in the initialization process. In \(k\)-WL, by applying a hash function, two \(k\)-tuples will get the same initial color if and only if they are from the same isomorphism class. However, this initialization results in limited initial colors due to the limited size of isomorphism classes within \(k\)-tuples. For example, \(1\)-WL assigns all nodes the same color since there is only one isomorphism type of \(1\)-tuple, which further restricts the expressivity of the following steps in the algorithm. In comparison, at the initialization of \(k,l\)-WL, unique labels are assigned to \(l\)-nodes. We then assign colors to these \(k\)-tuples according to their isomorphism types concerning the labeled tuples. See Appendix E for details and the mathematical forms. We will show this initialization makes \(k,l\)-WL more expressive than \(k\)-WL. Then the update scheme in \(k,l\)-WL is the same as that of \(k\)-WL.

The method to explicitly assign IDs to certain nodes aligns well with the original \(k\)-WL hierarchy, and we will show in Section 5 that \(k,l\)-WL is **strictly** more powerful than \(k\)-WL when \(l>0\). We refer our readers to Appendix E for detailed proofs, Appendix G for an example which helps to understand better the effect of explicitly introducing IDs, and Appendix H for more insights.

Additionally, note that any \(k,l\)-WL algorithm has a corresponding GNN implementation, which we name as \(k,l\)-GNN. Theoretically, if the instance contains a base GNN encoder equivalent to \(k\)-WL, explicitly embeds IDs to \(l\)-tuples, and uses an injective pooling function, \(k,l\)-GNN is as powerful as \(k,l\)-WL. To align with existing methods, we design two practical network architectures to implement \(k,l\)-GNN as shown in Figure 3.

### Unifying existing hierarchies

Here we briefly discuss the connection between our framework and previous work, including subgraph GNNs, relational pooling, GNN extensions mentioned in (Papp and Wattenhofer, 2022), and other methods. We refer our readers to Appendix F for more details.

Firstly, \(k,l\)-WL can incorporate all relational pooling (RP) and Local Relational Pooling methods, since node marking is the most general and expressive form and can simulate all other extensions (Papp and Wattenhofer, 2022).

Secondly, \(k,l\)-WL incorporates a wide range of subgraph GNNs. Zhang et al. (2023) shows that all node-based subgraph GNNs fall in one of 6 equivalent classes of Subgraph Weisfeiler-Lehman Tests. Remarkably, SWL is exactly \(1,1\)-WL (and equivalently, \(2,1\)-WL) in our framework, which reveals the connection between our work and many other subgraph GNNs unified by Zhang et al. (2023). Moreover, \(k,l\)-WL also incorporates some subgraph GNNs out of the scope of SWL (Zhang et al., 2023), such as \(I^{2}\)-GNN (Huang et al., 2022). Our \(1,2\)-WL is a slightly more powerful version than \(I^{2}\)-GNN since we consider all \(2\)-tuples to label, while \(I^{2}\)-GNN only considers those connected \(2\)-tuples. \(1,2\)-WL can distinguish some non-isomorphic graph pairs that SWL and 3-WL fail to discriminate, and the algorithm becomes even more powerful as we increase \(k\) or \(l\).

Thirdly, while a number of works such as OSAN (Qian et al., 2022) are a strict subclass of our framework, there are still some works cannot be incorporated directly. For example, \(k,l\)-WL operates independently on different labeled graphs and does not include interaction between labeled (sub)graphs as in Zhao et al. (2021) and those more expressive SWL variants in Zhang et al. (2023). Introducing inter-labeled-graph message passing will increase the expressivity of \(k,l\)-WL, but at the price of additional computation cost. It will also be complicated to analyze its theoretical expressivity if we introduce labeled graph interactions, which we leave for future work.

Finally, our framework incorporates all four kinds of GNN extensions in Papp and Wattenhofer (2022): higher-order WL, counting substructures, injecting local information, and marking nodes. Due to the limited space, the detailed discussion is in Appendix F.

## 5 The Expressivity Hierarchy of \(k,l\)-Wl

In this section, we theoretically analyze the expressivity of \(k,l\)-WL. We first define the comparison between algorithms' expressivity:

For any algorithm \(A\) and \(B\), we denote the final color of graph \(G\) computed by them as \(c_{A}(G)\) and \(c_{B}(G)\), we say:

* \(A\) is **more powerful** than \(B\) (\(B\preceq A\)) if for any pair of graphs \(G\) and \(H\), \(c_{A}(G)=c_{A}(H)\Rightarrow c_{B}(G)=c_{B}(H)\). Otherwise, there exists a pair of graphs that \(B\) can differentiate while \(A\) cannot, denoted as \(B\not\preceq A\).
* \(A\) is **as powerful as**\(B\) (\(A\cong B\)) if \(B\preceq A\wedge A\preceq B\).
* \(A\) is **strictly more powerful** than \(B\) (\(B\prec A\)) if \(B\preceq A\ \wedge\ A\not\cong B\), i.e., for any pair of graphs \(G\) and \(H\), \(c_{A}(G)=c_{A}(H)\Rightarrow c_{B}(G)=c_{B}(H)\), and there exists at least one pair of graphs \(H,G\) s.t. \(c_{B}(G)=c_{B}(H),c_{A}(G)\neq c_{A}(H)\).
* \(A\) and \(B\) are **incomparable** (\(A\nsim B\)) if \(A\npreceq B\ \land B\not\preceq A\). In this case, \(A\) can distinguish a pair of non-isomorphic graphs that cannot be distinguished by \(B\) and vice versa.

### Connection with existing hierarchies

A special case of \(k,l\)-WL is that \(l=0\) and no extra labels is attached to the graph. We have

**Theorem 5.1**.: \(\forall k\geq 2\)_, \(k,0\)-WL\(\cong k\)-WL._

Another special case of \(k,l\)-WL is the \(k=1\) case. With no label, \(2\)-WL is of the same expressivity to \(1\)-WL. We find that with \(l\) labels, the equality still holds.

**Theorem 5.2**.: \(\forall l\geq 0,\,1,l\)-WL\(\cong 2,l\)_-WL_

where \(1,l\)-WL is just \(l\)-OSAN (Qian et al., 2022).

A variant of \(k\)-WL test is \(k\)-Folklore Weisfeiler-Lehman (\(k\)-FWL) test. It is known that \(\forall k\geq 1\), \(k\)-FWL \(\cong k+1\)-WL. With labels, the equality still holds:

**Theorem 5.3**.: \(\forall k\geq 1,l\geq 0,\,k,l\)_-FWL\(\cong k+1,l\)-WL_

where \(k,l\)-FWL runs \(k\)-FWL on \(l\)-labeled graphs. See Appendix A for more details.

\(1,1\)-WL is equivalent to the vanilla subgraph Weisfeiler-Lehman test (SWL(VS)) proposed by Zhang et al. (2023).

**Theorem 5.4**.: \(1,1\)_-WL\(\cong\) SWL(VS)._

SWL(VS) further unifies various subgraph GNNs, like Nested GNN (Zhang and Li, 2021) and ID-GNN (You et al., 2021).

### Expressivity hierarchy of \(k,l\)-WL

Similar to WL tests, we can establish a hierarchy for \(k,l\)-WL in terms of distinguishing non-isomorphic graphs. The full hierarchy is shown in Figure 1.

\(\forall k\geq 2,l\geq 0,\,k,l\)-WL essentially produces colors for \(|V|^{k+l}\) tuples. Intuitively, increasing \(k\) and \(l\) will boost expressivity as more tuple colors will be computed. We show increasing \(k\) and \(l\) both **strictly** increases the expressivity.

**Theorem 5.5**.: \(\forall k\geq 2,l\geq 0,\,k,l\)_-WL\(\prec k+1,l\)-WL_

**Theorem 5.6**.: \(\forall k\geq 1,l\geq 0,\,k,l\)_-WL\(\prec k,l+1\)-WL_

However, with fixed \(k+l\) and number of colors, a larger \(k\) will lead to more message passing processes between tuples and stronger expressivity, shown in the following.

**Theorem 5.7**.: \(\forall k\geq 2,l\geq 0,\,k,l+1\)_-WL\(\prec k+1,l\)-WL_

With these theorems, we can prove a lot of useful corollaries. For example, \(k,l\)-WL is less expressive than \(k+l\)-WL:

**Corollary 5.8**.: \(\forall k\geq 2,l\geq 1\)_, \(k,l\)_-WL\(\prec k+l\)-WL_

Moreover, \(l+1\)-WL is not more powerful than \(2,l\)-WL (Qian et al., 2022).

**Theorem 5.9**.: \(\forall l\geq 1,\,2,l\)_-WL\(\not\preceq l+1\)-WL_

When \(l=2\), the above result recovers the known result of \(I^{2}\)-GNN (Huang et al., 2022).

Besides graph isomorphism power, counting power is also an important measure of the representation capability of GNNs. We conclude that \(k,l\)-WL is able to count all connected substructures within \(k+1\) nodes, see also Yan et al. (2023). This is also verified by our experiments on substructure counting.

## 6 Practical \(k,l\)-GNN: Improving Scalability and Compatibility

In this section, we will discuss \(k,l\)-GNN, the neural implementation of \(k,l\)-WL. \(k,l\)-GNN runs GNNs with the same expressivity as \(k\)-WL (Xu et al., 2018; Maron et al., 2019) on \(l\)-labeled graphs. We will discuss practical issues affecting the performance and scalability as well as our solutions. Our framework can also be applied to any other expressive GNNs to improve expressivity.

### Implementation of \(k,l\)-Gnn

Our \(k,l\)-WL runs \(k\)-WL on \(l\)-labeled graphs. Therefore, in implementation, we run a base encoder on all labeled graphs and pool their representations. Generally speaking, we can adopt any architecture as our base encoder and improve its expressive power via labeling. Specifically, when the base encoder has the same expressivity as \(k\)-WL, it reduces to our standard \(k,l\)-GNN. The base encoder, however, does not need to have exactly the same expressive power as certain \(k\)-WL in practice, since we can always upper bound the expressive power of the lifted model by some \(k,l\)-WL. In other words, our method can lift the expressivity of many existing architectures through labeling.

To make our framework applicable to any model, we propose two architectures to lift the expressive power of a base model. For convenience, we first suppose that the base encoder has \(k\)-WL-equivalent expressivity, and we will lift its expressivity to \(k,l\)-WL. _Architecture (a)_ in Figure 3 (a) exactly simulates the original form of \(k,l\)-WL: The input node feature matrix (represented by blue rectangles) is replicated to the number of labeled graphs and concatenated (or by other combination methods like add) with the corresponding label features. Then the original node features and label features are jointly learned by the same \(k\)-WL-equivalent GNN (base encoder). If the pooling function is injective (like Deepset (Zaheer et al., 2017)), architecture (a) fully preserves the expressivity of \(k,l\)-WL. _Architecture (b)_ in Figure 3 (b) is slightly different from the original
\(k,l\)-WL, which learns original node features and new label features by two models separately. The \(k\)-WL-equivalent base encoder learns the original node features only once. An extra \(k^{\prime}\)-WL (\(k^{\prime}\neq k\) is permitted) equivalent _structure encoder_ learns label features without node features in each \(l\)-labeled graph. The \(l\)-labeled graph representations are then aggregated by an aggregator and concatenated to the representation of the original graph. We call it structure encoder since label features do not introduce information more than graph structure.

Theoretical analysis and comparison between these two architectures are presented in Appendix I. In short, architecture (b) is less expressive but more scalable. As the number of labeled graphs is \(n^{l}\), and the complexity of \(k\)-WL-equivalent GNN can be \(n^{k}\), architecture (a) can preserve the expressivity of \(k,l\)-WL with a complexity of \(O(n^{k+l})\). Architecture (b) can use different structure encoder and base encoder, i.e. \(k^{\prime}<k\). Then the complexity of architecture (b) can be reduced to \(O(n^{k}+n^{k^{\prime}+l})\). Moreover, we experimentally find that architecture (b) tends to outperform (a). One intuition behind is that we may need two different sets of parameters to learn label feature and node features, respectively.

In summary, architecture (b) is designed to reduce the complexity at the cost of losing part of expressive power compared with architecture (a), the original GNN implementation of \(k,l\)-WL. However, the decoupled architecture tends to perform better in real-world tasks. Despite their differences, we emphasize that both architectures can boost the expressivity.

Below, we propose several \(k,l\)-GNN instances parameterized by different \(k\) and \(l\), including ID-MPNN and ID-PPGN. We as well lift some other base models through our framework, such as ID-Transformer. Note that the following instances can all adopt either architecture (a) or (b), depending on application scenarios.

Id-MpnnAn \(l\)-IDMPNN is an instance of \(1,l\)-WL, since Message Passing Neural Network (MPNN) is equivalent to \(1\)-WL (Xu et al., 2018). The model is easy to implement but reveals strong expressivity as \(l\) increases. One can easily verify that ID-MPNN incorporates Identity-aware GNN (\(l=1\)), \(I^{2}\)-GNN (\(l=2\)) and \(l\)-OSAN. With our implementation, ID-MPNN outperforms the above models experimentally.

Id-PpgnFor \(3\)-WL equivalent base encoders, we select PPGN (Maron et al., 2019) as our base encoder. An ID-PPGN with \(l\) labels is as powerful as \(3,l\)-WL, which is strictly more powerful than \(3\)-WL when \(l>0\).

ID-TransformerWe also apply our framework to graph transformers. It is not a \(k,l\)-GNN strictly as graph transformer is not \(k\)-WL equivalent. However, our framework is general enough to take any base graph learning model.

Additionally, many techniques can be applied in the implementation of \(k,l\)-GNN, such as positional encodings (PE) and structure encodings (SE) (Rampasek et al., 2022), GNN as kernel techniques (Zhao et al., 2021), etc. In conclusion, our method is a universal framework to improve the expressivity of base models while being compatible with many other methods and techniques. Please refer to Appendix J for more implementation details. In the next section, we discuss how to reduce the complexity when \(l\) is large.

### Labeled graph selection

As \(k,l\)-GNN runs a base encoder on each labeled graph, the total complexity is proportional to the number of labeled graphs. This section focuses on how to select a subset of labeled graphs. We also discuss how to segregate a subgraph

Figure 3: Two architectures of \(k,l\)-GNN. (a) The input graph feature (represented by blue rectangle) is duplicated for each labeled subgraph, and one base encoder jointly learns ID features (represented by squares of different colors) and graph features. (b) The graph features and ID features are parallelly learned by the base encoder and a structure encoder, respectively, which are then aggregated together and passed to downstream architectures.
from the labeled graph and thus reduce the size of labeled graphs in Appendix B, as subgraph size also affects the scalability and reduced subgraph size is important for the successes of many subgraph GNNs (Zhang and Li, 2021).

We list some labeled graph selection strategies as follows. Among them, random sampling and node-based policies are commonly used. We propose two new sampling policies: constraint-based policy and hierarchical policy.

* Random sampling. This is one of the most common methods in existing subgraph GNNs. Statistically, the graph representation is permutation invariant and unbiased but with a variance among different samplings. We can conduct parallel samplings for variance reduction and use mean/median/voting as the final representation. The complexity of randomly sampling \(l\)-tuples is \(O(\alpha n^{l})\), where \(\alpha\) is the sampling rate.
* Node-based policies. This is another family of invariant sampling methods. For example, we can extract a K-hop ego-net for each root node and select all \(l\)-tuples in the ego-net with the root node always being the first node in each \(l\)-tuple. The complexity is \(O(nm^{l-1})\), where \(m\) is the average size of ego-nets.
* Constraint-based policies. These methods search within all possible \(l\)-tuples and filter out those failing to meet certain constraints. For example, if we upper-bound the shortest path distance \(3\) between any pair of nodes in a \(6\)-tuple, we can sample many \(6\)-rings while excluding any \(6\)-paths. Compared with node-based policy, constraint-based methods do not sample the same induced subgraphs repetitively and enjoy a higher design freedom. In implementation, this can be efficiently implemented by dynamic programming, e.g., Floyd-Warshall algorithms. The number of subgraphs depends on the graph and constraints.
* Hierarchical policies. Theses methods hierarchically select subgraphs. We can use algorithms like min-cut or node clustering (e.g., spectral clustering) to divide the graph into clusters with an average size \(m\). Label nodes are then selected only within each cluster, resulting in an average complexity of \(O(m^{l}.\frac{n}{m})\). Since \(m\ll n\), the hierarchical policy can significantly reduce the number of subgraphs while still being able to encode sufficient local structure information.

In most real-world task experiments, we use constraint-based and hierarchical policies, achieving impressive experimental results at a low computation complexity. See Appendix L for an ablation study on different sampling policies.

Learn to sample and Learn to labelAside from traversing or sampling according to certain rules, Qian et al. (2022) use Implicit-MLE framework (Niepert et al., 2021) (which allows us to back-propagate through continuous-discrete architecture) to sample subgraphs in a data-driven fashion. With this method, \(k,l\)-GNN can also learn to label tuples and sample subgraphs to minimize the target loss function in a data-driven manner. See Appendix L for more details and experimental results.

## 7 Experiments

In this section, we conduct experiments on both synthetic and real-world tasks to verify our models' expressivity and real-world performance.

### Graph isomorphism task

DatasetWe select two synthetic datasets, EXP and SR25, to empirically verify our models' expressivity for graph isomorphism tasks. EXP (Abboud et al., 2020) contains 600 pairs of non-isomorphic graphs that 1-WL and 2-WL fail to distinguish. SR25 (Balcilar et al., 2021) contains 15 non-isomorphic strongly regular graphs (i.e., 105 non-isomorphic pairs) that 3-WL fails to distinguish. An accuracy of \(50\%\) on EXP and \(6.67\%\) on SR25 suggest the model fails to distinguish any non-isomorphic graphs in the dataset.

ModelsFor baseline models, we choose GIN, PNA (Corso et al., 2020), Identity-aware GNN (You et al., 2021), GIN-AK+ (Zhao et al., 2021) and PPGN (Maron et al., 2019). In comparison, we choose our ID-MPNN and ID-PPGN to understand the expressivity hierarchy better.

ResultsThe results are shown in Table 2. When the number of IDs \(l\geq 2\), ID-MPNN and ID-PPGN achieve perfect performance on the two datasets. In comparison, all other models fail on the SR25 dataset. By comparing the results of GIN and \(l\)-IDMPNN as well as PPGN and \(l\)-IDPPGN, we verify that our framework can improve expressivity. Concretely, we have the following conclusions:

* \(1,1\)-WL is more powerful than \(1\)-WL and \(2\)-WL.
* \(1,2\)-WL and \(3,2\)-WL can distinguish some non-isomorphic graph pairs that are indistinguishable by \(3\)-WL.

These results are consistent with our theoretical analysis in Section 5.
### Substructure counting

DatasetTo verify our model's expressivity of counting substructures, we evaluate on random regular graph dataset (Chen et al., 2020). There are four target substructures: triangle, tailed triangle, star and chordal-cycle. Test MAE measures the results.

ModelsWe choose GCN, KP-GIN+ (Feng et al., 2022), GIN-AK+ (Zhao et al., 2021), and DeepLRP (Chen et al., 2020) as the baseline models. We use \(l\)-IDMPNN, and additionally restricts message passing only on the labeled tuples, where \(l\) is the size of the target substructure.

ResultsThe results are shown in Table 3. Our model achieves state-of-the-art performance on all tasks, and the test MAE is nearly \(0\). This verifies our theoretical results that \(1,l\)-GNN can completely count substructures within \(l\) nodes.
follow a commonly used training/validation/test split ratio of 0.8/0.1/0.1, and the results of the first 12 targets are reported. ZINC12k is a subset of ZINC250k containing 12k molecules. The task is also molecular property (constrained solubility) regression. ogbg-molhiv contains 41k molecules for graph binary classification (whether a molecule inhibits HIV virus replication or not). We use the official split for ZINC and ogbg-molhiv.

ModelsFor QM9 dataset, MPNN, DTNN (Wu et al., 2017), DeepLRP (Chen et al., 2020), PPGN (Maron et al., 2019) and Nested GNN (Zhang and Li, 2021) are chosen as baseline models. For ZINC12k, we choose GIN (Xu et al., 2018), PNA (Corso et al., 2020), DeepLRP (Chen et al., 2020), OSAN (Qian et al., 2022), KP-GIN+ (Feng et al., 2022), GNN-AK+ (Zhao et al., 2021), CIN (Bodnar et al., 2021) and GPS (Rampasek et al., 2022) for comparison. For ogbg-molhiv, PNA (Corso et al., 2020), DeepLRP (Chen et al., 2020), NGNN (Zhang and Li, 2021), KP-GIN (Feng et al., 2022), \(I^{2}\)-GNN (Huang et al., 2022), CIN (Bodnar et al., 2021) and SUN(EGO) (Frasca et al., 2022) are selected.

ResultsOur \(4\)-IDMPNN achieves superior performance in 7 out of 12 tasks on QM9 dataset (Table 1), while results for the remaining targets are also highly competitive. On ZINC12k (Table 4) and ogbg-molhiv (Table 5), although IDMPNN does not achieve the best results, it is still comparable to the state-of-the-art models. Moreover, we do not use any additional features or pretraining in any datasets, reflecting the power of our model. This suggests that our method can effectively enhance the performance of base encoders for real-world tasks in addition to increasing the expressivity.While our \(k,l\)-GNN framework captures DeepLRP, GSN and OSAN, we find instances such as \(4\)-IDMPNN that greatly surpass these works in real-world tasks.

## 8 Conclusions

In this work, we establish a novel \(k,l\)-WL framework that explicitly assigns labels to \(l\) nodes while running a \(k\)-WL algorithm. We theoretically analyze the expressivity hierarchy of \(k,l\)-WL, which incorporates many existing relational pooling methods and subgraph GNNs. Due to its strong compatibility, our framework can improve the expressivity of any base model by just augmenting ID features on (sampled) subgraphs. Various acceleration methods are also discussed to build practical, effective models. Some of our \(k,l\)-GNN instancees achieve state-of-the-art performance on several synthetic and real-world tasks, verifying the power of our framework.

## Acknowledge

This project is supported in part by the National Key Research and Development Program of China (No. 2021ZD0114702).
