A Practical, Progressively-Expressive GNN
Anonymous Author(s)
AfﬁliationAddress email
Abstract
Message passing neural networks (MPNNs) have become a dominant ﬂavor ofgraph neural networks (GNNs) in recent years.
Yet, MPNNs come with notablelimitations; namely, they are at most as powerful as the 1-dimensional WeisfeilerLeman (1-WL) test in distinguishing graphs in a graph isomorphism testing framework.
To this end, researchers have drawn inspiration from the k -WL hierarchyto develop more expressive GNNs.
However, current k -WL-equivalent GNNs arenot practical for even small values of k , as k -WL becomes combinatorially morecomplex as k grows.
At the same time, several works have found great empirical success in graph learning tasks without highly expressive models, implyingthat chasing expressiveness with a “coarse-grained ruler” of expressivity like kWL is often unneeded in practical tasks.
To truly understand the expressivenesscomplexity tradeoff, one desires a more “ﬁne-grained ruler,” which can moregradually increase expressiveness.
Our work puts forth such a proposal: Namely,we ﬁrst propose the ( k, c )(  ) -S ET WL hierarchy with greatly reduced complexityfrom k -WL, achieved by moving from k -tuples of nodes to sets with  k nodesdeﬁned over  c connected components in the induced original graph.
We showfavorable theoretical results for this model in relation to k -WL, and concretize itvia ( k, c )(  ) -S ET GNN, which is as expressive as ( k, c )(  ) -S ET WL.
Our modelis practical and progressively-expressive , increasing in power with k and c .
Wedemonstrate effectiveness on several benchmark datasets, achieving several stateof-the-art results with runtime and memory usage applicable to practical graphs.
Introduction
In recent years, graph neural networks (GNNs) have gained considerable attention [46, 48] for theirability to tackle various node-level [25, 44], link-level [53, 39] and graph-level [8, 30] learning tasks,given their ability to learn rich representations for complex graph-structured data.
The common tem- plate for designing GNNs follows the message passing paradigm; these so-called message-passingneural networks (MPNNs) are built by stacking layers which encompass feature transformation andaggregation operations over the input graph [29, 18].
Despite their advantages, MPNNs have several limitations including oversmoothing [54, 11, 36], oversquashing [2, 43], inability to distinguishnode identities [50] and positions [51], and expressive power [49].
Since Xu et al.’s [49] seminal work showed that MPNNs are at most as powerful as the ﬁrst-order
Weisfeiler-Leman (1-WL) test in the graph isomorphism (GI) testing framework, there have beenseveral follow-up works on improving the understanding of GNN expressiveness [3, 13].
In response, the community proposed many GNN models to overcome such limitations [40, 4, 55].
Several of these aim to design powerful higher-order GNNs which are increasingly expressive[34, 31, 24, 8] by inspiration from the k -WL hierarchy [42].
A key limitation of such higher-order models that reach beyond 1-WL is their poor scalability;in fact, these models can only be applied to small graphs with small k in practice [30, 31, 34]
1, 2, 2, 3, 3, 1, 1, 1, 2, 4, k=2sets k=1sets k=3sets 1,2, 1,2,      1,2, 1,3, 1,3, 2,3, 2,3, 3,4, 2,4,5 c=1-connected sets c=2-
[Figure]
sets edges in (3,1) super-graph edges in (3,2) super-graph  3 e.g. graph extract k-node subgraphs 
GNN
GNN GNN
GNN init.
“color”vectors e.g., (k=3, c=2)-bipartite super-graph
B i d i r ec t i ona l b i p a r t i t e G
NN l a y e r …
B i d i r ec t i ona l b i p a r t i t e G
NN l a y e r forward msg.
pass.
backward msg.
pass.
graphemb.
Historically, Babai et al.
[5] showed that almost all graphs on n nodes can be distinguished by 1WL.
In other contexts like node classiﬁcation, researchers have encountered superior generalizationperformance with graph-augmented multi-layer perceptron (GA-MLP) models [38, 27] comparedto MPNNs, despite the former being strictly less expressive than the latter [12].
Considering thatincreased model complexity has negative implications in terms of overﬁtting and generalization [21],it is worth re-evaluating continued efforts to pursue maximally expressive GNNs in practice.
Ideally,one could study these models’ generalization performance on various real-world tasks by increasingk (expressiveness) in the k -WL hierarchy.
Yet, given the impracticality of this too-coarse “ruler”of expressiveness, which becomes infeasible beyond k> 3 , one desires a more ﬁne-grained “ruler”,that is, a new hierarchy whose expressiveness grows more gradually.
Such a hierarchy could enableus to gradually build more expressive models which admit improved scaling, and avoid unnecessaryleaps in model complexity for tasks which do not require them, guarding against overﬁtting.
Present Work.
In this paper, we propose such a hierarchy, and an associated practical progressivelyexpressive GNN model , called ( k, c )(  ) -S ET GNN, whose expressiveness can be modulated throughk and c .
In a nutshell, we take inspiration from k -WL, yet achieve practicality through three design ideas which simplify key bottlenecks of scaling k -WL:
First , we move away from k -tuplesof the original k -WL to k -multisets (unordered), and then to k (  ) -sets.
We demonstrate that thesesteps drastically reduce the number nodes in the k -WL graph while retaining signiﬁcant expressiveness. (Sec.s 3.2  3.3)
Second , by considering the underlying sparsity of an input graph, wereduce scope to k, c (  ) -sets that consist of  k nodes whose induced subgraph is comprised of c  k connected components.
This also yields massive reduction in the number of nodes whileimproving practicality; i.e. small values of c on sparse graphs can allow one to increase k wellbeyond 3 in practice.
(Sec. 3.4) Third , we design a super-graph architecture that consists of asequence of k  1 bipartite graphs over which we learn embeddings for our k, c (  ) -sets, using bidirectional message passing.
These embeddings can be pooled to yield a ﬁnal graph embedding.. (Sec.s 4.1  4.2)
We also speed up initializing “colors” for the k, c (  ) -sets, for c > 1 , substantially reducing computation and memory.
(Sec. 4.3) Fig.
1 overviews of our proposed framework.
Experimentally, our ( k, c )(  ) -S ET GNN outperforms existing state-of-the-art GNNs on simulatedexpressiveness as well as real-world graph-level tasks, achieving new bests on substructure countingand ZINC-12K, respectively.
We show that generalization performance reﬂects increasing expressiveness by k and c .
Our proposed scalable designs allow us to train models with e.g. k
=10 or c =3with practical running time and memory requirements.
Our code and datasets are open-sourced athttps://anonymous.4open.science/r/KCSetGNN.
Related Work
Overcoming MPNN limitations.
Several works aim to improve expressiveness limitations of
MPNNs.
One approach is to inject features into the MPNN aggregation, motivated by Loukas[28] who showed that MPNNs can be universal approximators when nodes are sufﬁciently distin- guishable.
Sato et al.
[41] show that injecting random features can better enable MPNNs to solveproblems like minimum dominating set and maximum matching.
You et al.
[50] inject cycle countsas node features, motivated by the limitation of MPNNs not being able to count cycles.
Others proposed utilizing subgraph isomorphism counting to empower MPNNs with substructure informationthey cannot learn [9, 8].
Earlier work by You et al.
[51] adds positional features to distinguishnode which naïve MPNN embeddings would not.
Recently, several works also propose utilizingsubgraph aggregation; Zhao et al.
[55] and Bevilacqua et al.
[7] propose subgraph variants of the
WL test which are no less powerful than 3-WL.
The community is yet exploring several distinct avenues in overcoming limitations: Murphy et al.
[35] propose a relational pooling mechanism whichsums over all permutations of a permutation-sensitive function to achieve above-1-WL expressiveness.
Balcilar et al.
[6] generalizes spatial and spectral MPNNs and shows that instances of spectral
MPNNs are more powerful than 1-WL.k -WL-inspired GNNs.
The k -WL test captures higher-order interactions in graph data by considering all k -tuples, i.e., size k ordered multisets deﬁned over the set of nodes.
While highly expressive,it does not scale to practical graphs beyond a very small k ( k = 3 pushes the practical limit evenfor small graphs).
However, several works propose designs which can achieve k -WL in theory andstrive to make them practical.
Maron et al.
[31] proposes a general class of invariant graph networksk -IGNs having exact k -WL expressivity, while being not scalable.
Morris et al.
have several workon k -WL and its variants like k -LWL [32], k -GNN [34], and  - k -WL-GNN [33].
Both k -LWL andk -GNN use a variant of k -WL that only considers k -sets and are strictly weaker than k -WL but muchmore practical.
In our paper we claim that k (  ) -sets should be used with additional designs to keepthe best expressivity while remaining practical.
The  - k -WL-GNN works with k -tuples but sparsiﬁes the connections among tuples by considering locality.
All three variants proposed by Morris etal.
do not have realizations beyond k > 3 in their experiments while we manage to reach k = 10 .
A practical progressively-expressive isomorphism test: ( k, c )(  ) -S ET WL
We ﬁrst motivate our method ( k, c )(  ) -S ET GNN from the GI testing perspective by introducing( k, c )(  ) -S ET WL.
We ﬁrst introduce notation and background of k -WL (Sec. 3.1).
Next, we showhow to increase the practicality of k -WL without reducing much of its expressivity, by removingnode ordering (Sec. 3.2), node repetitions (Sec. 3.3), and leveraging graph sparsity (Sec. 3.4).
Weclose the section with complexity analysis (Sec. 3.5).
All proofs are included in Appendix A.2.
Notation: Let G =
( V ( G ) , E ( G ) , l G ) be an undirected, colored graph with nodes V ( G ) , edges
E ( G ) , and a color-labeling function l G : V ( G ) !
C where C = { c 1 , ..., c d } denotes a set ofd distinct colors.
Let [ n ] = { 1 , 2 , 3 , ..., n } .
Let ( · ) denote a tuple (ordered multiset), {{·}} de- 110 note a multiset (set which allows repetition), and {·} denote a set .
We deﬁne !
v =
( v 1 , ..., v k ) 111 as a k -tuple, ˜ v = {{ v 1 , ..., v k }} as a k -multiset, and ˆ v = { v 1 , ..., v k } as a k -set.
Let !
v [ x/i ] = 112 ( v 1 , ..., v i  1 , x, v i +1 , ..., v k ) denote replacing the i -th element in !
v with x , and analogously for˜ v [ x/i ] and ˆ v [ x/i ] .
When v i 2 V ( G ) , let
G [ ! v ] , G [˜ v ] , G [ˆ v ] denote the induced subgraph on G withnodes inside !
v , ˜ v , ˆ v respectively (keeping repetitions).
An isomorphism between two graphs G and
G 0 is a bijective mapping p : V ( G ) !
V ( G 0 ) which satisﬁes 8 u, v 2 V ( G ) , ( u, v ) 2 E ( G ) () 116 ( p ( u ) , p ( v ))
2 E ( G 0 ) and 8 u 2 V ( G ) , l G ( u ) =
l G 0 ( p ( u )) .
Two graphs G, G 0 are isomorphic ifthere exists an isomorphism between them, which we denote as G ⇠ = G 0 .
3.1 Preliminaries: the k -Weisfeiler-Leman ( k -WL) Graph Isomorphism Test
The 1-dimensional Weisfeiler-Leman ( 1 -WL) test, also known as color reﬁnement [37] algorithm, isa widely celebrated approach to (approximately) test GI.
Although extremely fast and effective formost graphs ( 1 -WL can provide canonical forms for all but n  1 / 7 fraction of n -vertex graphs [5]),it fails to distinguish members of many important graph families, such as regular graphs.
The morepowerful k -WL algorithm ﬁrst appeared in [45], and extends coloring of vertices (or 1-tuples) tothat of k -tuples.
k -WL is progressively expressive with increasing k , and can distinguish any ﬁniteset of graphs given a sufﬁciently large k .
k -WL has many interesting connections to logic, games,and linear equations [10, 20].
Another variant of k -WL is called k -FWL (Folklore WL), such that( k + 1) -WL is equally expressive as k -FWL [10, 19].
Our work focuses on k -WL.k -WL iteratively recolors all n k k -tuples deﬁned on a graph G with n nodes.
At iteration 0, eachk -tuple !
v =
( v 1 , ..., v k ) 2 V ( G ) k is initialized with a color as its atomic type at k ( G, ! v ) .
Assume
G has d colors, then at k ( G, ! v ) 2 { 0 , 1 } 2 ( k 2 ) + kd is an ordered vector encoding.
The ﬁrst  k entries indicate whether v i =
v j , 8 i, j, 1  i<j  k , which is the node repetition information.
The second k entries indicate whether ( v i , v j ) 2 E ( G ) .
The last kd entries one-hot encode the initial color ofv i , 8 i, 1  i  k .
Importantly, at k ( G, ! v ) =
at k ( G 0 , !
v 0 )
if and only if v i 7!
v
i is an isomorphismfrom G [ s ( ! v )] to G
[ s ( !
v 0 )] .
Let wl ( t ) k
( G, ! v ) denote the color of k -tuple !
v on graph G at t -thiteration of k -WL, where colors are initialized with wl (0) k
( G, ! v )
=
at k ( G, ! v ) .
At the t -th iteration, k -WL updates the color of each !
v 2 V ( G ) k according towl ( t +1)
k
( G, ! v )
=
HASH ✓ wl ( t ) k
( G, ! v ) , ⇢⇢ wl ( t ) k
( G, !
v [ x/ 1])  x 2 V ( G )  ,..., ⇢⇢ wl ( t )
k
( G, !
v [ x/k ])  x 2 V ( G ) ◆ (1) 138 Let gwl ( t ) k
( G ) denote the encoding of G at t -th iteration of k -WL.
Then,gwl ( t ) k
( G ) =
HASH ✓⇢⇢ [Equation] ( t ) k ( G
[Equation]
!
[Equation]
)  ! v
[Equation]
k ◆ (2)
Two k -tuples !
v and !
u are connected by an edge if |{ i 2 [ k ] | v i =
u i }| = k  1 , or informally 141 if they share ( k  1) entries.
Then, k -WL deﬁnes a super-graph S k -wl ( G ) with its nodes being allk -tuples in G , and edges deﬁned as above.
Eq.
(1) deﬁnes the rule of color reﬁnement on S k -wl ( G ) .
Intuitively, k -WL is akin to (but more powerful as it orders subgroups of neighbors) running 1-WLalgorithm on the supergraph S k -wl ( G ) .
As t ! 1 , the color wl ( t +1) k
( G, ! v ) converges to a stablevalue, denoted as wl ( 1 ) k
( G, ! v ) and the corresponding stable graph color denoted as gwl ( 1 ) k
( G ) .
For two non-isomorphic graphs G, G 0 , k -WL can successfully distinguish them if gwl ( 1 )
k
( G )
=gwl ( 1 ) k
( G 0 ) .
The expressivity of wl ( t ) k can be exactly characterized by ﬁrst-order logic with counting quantiﬁers.
Let C tk denote all ﬁrst-order formulas with at most k variables and t -depth countingquantiﬁers, then wl ( t ) k
( G, ! v
[Equation] = wl
[Equation]
Despite its power, k -WL uses all n k tuples and has O ( kn k ) complexity at each iteration.
3.2 From k -WL to k -M ULTISET WL : Removing Ordering
Our ﬁrst proposed adaptation to k -WL is to remove ordering in each k -tuple, i.e. changing k -tuplesto k -multisets.
This greatly reduces the number of supernodes to consider by O ( k !) times.
Let ˜ v = {{ v 1 , ..., v k }} be the corresponding multiset of tuple !
v =
( v 1 , ..., v k ) .
We introduce acanonical order function on G , o G : V ( G ) !
[ n ] , and a corresponding indexing function o  1 G (˜ v , i ) ,which returns the i -th element of v 1 , ..., v k sorted according to o G .
Let mwl ( t ) k ( G, ˜ v ) denote thecolor of the k -multiset ˜ v at t -th iteration of k -M ULTISET WL, formally deﬁned next.
At t = 0 , we initialize the color of ˜ v as mwl (0) k
( G, ˜ v ) =
HASH  {{ at k ( G, p (˜ v ))
| p 2 perm[k] }}  where perm[k] 1 denotes the set of all permutation mappings of k elements.
It can be shown thatmwl (0) k
( G, ˜ v ) =
mwl (0) k
( G 0 , ˜ v 0 )
if and only if G [˜ v ] and G
[˜ v 0 ] are isomorphic.
At t -th iteration, k -M ULTISET WL updates the color of every k -multiset by
[Equation]
Where ˜ v [ x/o  1 G
(˜ v , i )]) denotes replacing the i -th (ordered by o G ) element of the multiset withx 2 V ( G ) .
Let S k -mwl ( G ) denote the super-graph deﬁned by k -M ULTISET WL.
Similar to Eq.
(2),the graph level encoding is gmwl ( t ) k
( G ) =
HASH ( {{ mwl ( t ) k
[Equation]
. Interestingly, although k -M ULTISET WL has signiﬁcantly fewer number of node groups than k -WL,we show it is equally expressive as k -WL in terms of distinguishing node multisets and graphs.
Theorem 1 ( k -WL ⌘ k -M ULTISET WL ) Let k  1 and wl ( t ) k
( G, ˜ v ) :
= {{ wl ( t ) k
( G, p (˜ v )) | p 2perm[k] }} , then for all t 2 N and all graphs G, G 0 : k -WL and k -M ULTISET WL are equivalent in distinguishing multisets G, ˜ v and G 0 , ˜ v 0 at t -th iteration, i.e. wl ( t ) k
( G, ˜ v ) =
wl ( t ) k ( G 0 , ˜ v 0 ) ()mwl ( t ) k
( G, ˜ v ) =
mwl ( t ) k
( G 0 , ˜ v 0 ) .
3.3 From k -M ULTISET WL to k (  ) -S ET WL : Removing Repetition
Next, we propose further removing repetition inside any k -multiset, i.e. transforming k -multisetsto k (  ) -sets.
We assume elements of k -multiset ˜ v and k (  ) -set
ˆ v in G are sorted based on theordering function o G , and omit o G for clarity.
Let s ( · ) transform a multiset to set by removingrepeats, and let r ( · ) return a tuple with the number of repeats for each distinct element in a multiset.
Speciﬁcally, let ˆ v = s
(˜ v ) and ˆ n = r
(˜ v ) , then m :=
| ˆ v | =
| ˆ n | 2 [ k ] denotes the number ofdistinct elements in k -multiset ˜ v , and 8 i 2 [ m ] , ˆ v
i is the i -th distinct element with ˆ n i repetitions.
Clearly there is an injective mapping between ˜ v and (ˆ v , ˆ n ) ; let f be the inverse mapping such that˜ v = f
( s (˜ v ) , r (˜ v )) .
Equivalently, each m -set ˆ v can be mapped with a multiset of k -multisets:
ˆ v ${{ ˜ v = f
(ˆ v , ˆ n ) |
P mi =1 ˆ n i =
k, 8 i ˆ n i  1 }} .
Based on this relationship, we extend the connection 181 among k -multisets to k (  ) -sets: given m 1 , m 2
[ k ] , a m
-set ˆ v is connected to a m 2 -set ˆ u if andonly if 9 ˆ n v , ˆ n n , f
(ˆ v , ˆ n v ) is connected with f (ˆ u , ˆ n u ) in k -M ULTISET WL.
Let S k -swl ( G ) denotethe deﬁned super-graph on G by k (  ) -S ET WL.
It can be shown that this is equivalent to either(1) ( | m 1  m 2 | = 1) ^ ( | ˆ v \ ˆ u | = min( m 1 , m 2 )) or (2) ( m 1 = m 2 ) ^
( | ˆ v \ ˆ u | = m 1  1) 185 is true.
Notice that S k -swl ( G ) contains a sequence of k  1 bipartite graphs with each reﬂectingthe connections among the ( m  1 )-sets and the m -sets.
It also contains k  1 subgraphs, i.e. theconnections among the m -sets for m = 2 , ..., k .
Later on we will show that these k  1 subgraphscan be ignored without affecting k (  ) -S ET WL.
Let swl ( t ) k
( G, ˆ v ) denote the color of m -set ˆ v at t -th iteration of k (  ) -S ET WL.
Now we formallydeﬁne k (  ) -S ET WL.
At t = 0 , we initialize the color of a m -set ˆ v ( m
[ k ] ) as:
[Equation]
Clearly swl (0) k
( G, ˆ v )
= swl (0) k
( G 0 , ˆ v 0 )
if and only if G [ˆ v ] and G 0 [ˆ v 0 ] are isomorphic.
At t -thiteration, k (  ) -S ET WL updates the color of every m -set ˆ v byswl ( t +1) k
[Equation]
( G, ˆ v [ { x } ) | x 2 V ( G ) \ ˆ v }} , {{ swl ( t ) k ( G, ˆ v \ x ) | x 2 ˆ v }} , ⇢⇢ {{ swl ( t ) k ( G, ˆ v [ x/o  1 G (ˆ v , 1)]) | x 2 V ( G ) \ ˆ v }} , ..., {{ swl ( t ) k
( G, ˆ v [ x/o  1 G (ˆ v , m )]) | x 2 V ( G ) \ ˆ v }} ◆ (
[Equation]
Notice that when m = 1 and m = k , the third and second part of the hashing input is an emptymultiset, respectively.
Similar to Eq.
(2), we formulate the graph level encoding as gswl ( t ) k
( G ) =
HASH [Equation]
[Equation]
To characterize the relationship of their expressivity, we ﬁrst extend k -M ULTISET WLon sets by deﬁning the color of a m -set ˆ v on k -M ULTISET WL as mwl ( t ) k
( G, ˆ v ) :
={{ mwl ( t ) k
( G, f ( ˆ v , ˆ n ))
P mi =1 ˆ v
i =
k, 8 i ˆ n i  1 }} .
We prove that k -M ULTISET WL is at leastas expressive as k (  ) -S ET WL in terms of separating node sets and graphs.
Theorem 2 Let k  1 , then 8 t 2 N and all graphs G, G 0 : mwl ( t ) k ( G, ˆ v ) =
mwl ( t ) k
( G 0 , ˆ v 0 )
= )swl ( t ) k
( G, ˆ v )
=
swl ( t ) k
( G 0 , ˆ v 0 ) .
We also conjecture that k -M ULTISET WL and k (  ) -S ET WL could be equally expressive, andleave it to future work.
As a single m -set corresponds to  k  1 m  k -multisets, moving from kM ULTISET WL to k (  ) -S ET WL further reduces the computational cost greatly. 3.4 From k (  ) -S ET WL to ( k, c )(  ) -S ET WL : Accounting for Sparsity
Notice that for two arbitrary graphs G and G 0 with equal number of nodes, the number of k (  ) -sets and the connections among all k (  ) -sets in k (  ) -S ET WL are exactly the same, regardless ofwhether they are dense or sparse.
We next propose to account for the sparsity of a graph G tofurther reduce the complexity of k (  ) -S ET WL.
As the graph structure is encoded inside every mset, when the graph becomes sparser, there would be more sparse m -sets with a potentially largenumber of disconnected components.
Based on the hypothesis that the induced subgraph over aset (of nodes) with fewer disconnected components naturally contains more structural information, we propose to restrict the k (  ) -sets to be ( k, c )(  ) -sets: all sets with at most k nodes and at mostc connected components in its induced subgraph.
Let S k,c -swl ( G ) denote the super-graph deﬁnedby ( k, c )(  ) -S ET WL, then S k,c -swl ( G )
=
S k -swl
( G )
[ { ˆ v | #components ( G [ˆ v ])  c } ] , which is theinduced subgraph on the super-graph deﬁned by k (  ) -S ET WL.
Fortunately, S k,c -swl ( G ) can beefﬁciently and recursively constructed based on S ( k  1 ,c ) -swl ( G ) , and we include the constructionalgorithm in the Appendix A.3.
( k, c )(  ) -S ET WL can be deﬁned similarly to k (  ) -S ET WL (Eq.(4) and Eq.
(5)), however while removing all colors of sets that do not exist on S k,c -swl ( G ) .
( k, c )(  ) -S ET WL is progressively expressive with increasing k and c , and when c = k , ( k, c )(  ) -
S ET WL becomes the same as k (  ) -S ET WL, as all k (  ) -sets are then considered.
Let swl ( t ) k,c ( G, ˆ v )denote the color of a ( k, c )(  ) -set
ˆ v on t -th iteration of ( k, c )(  ) -S ET WL, then gswl ( t ) k,c
( G ) =
HASH ( {{ swl ( t ) k,c  G, ˆ v  8 ˆ v 2 S k,c -swl ( G )  }} .
Theorem 3 Let k  1 , then 8 t 2 N and all graphs G, G 0 :
(1) when 1  c 1 < c 2  k , if G, G 0 cannot be distinguished by ( k, c 2 ) (  ) -S ET WL , theycannot be distinguished by ( k, c 1 )(  ) -S ET WL(2) when k 1 < k 2 , 8 c  k 1 , if G, G 0 cannot be distinguished by ( k 2 , c )(  ) -S ET WL , theycannot be distinguished by ( k 1 , c )(  ) -S ET WL
3.5 Complexity Analysis
All color reﬁnement algorithms described above run on a super-graph; thus, their complexity ateach iteration is linear to the number of supernodes and number of edges of the super-graph.
Insteadof using big O notation that ignores constant factors, we compare the exact number of supernodesand edges.
Let G be the input graph with n nodes and average degree d .
For k -WL, there are n ksupernodes and each has n ⇤ k number of neighbors, hence S k -wl
( G ) has n k supernodes and kn k +1 / 2edges.
For m 2 [ k ] , there are  nm  m -sets and each connects to m number of ( m  1) -sets.
So S k -swlhas
[Equation]
and
[Equation]
237 (derivation in Appendix).
Here we ignore edges within m -sets for any m
[ k ] as they can bereconstructed from the bipartite connections among ( m  1) -sets and m -sets, described in detail in
Sec. 4.1.
Consider e.g., n = 30 , k = 5 ; we get | V ( S k -wl ( G ))
| | V ( S k - [Equation] G ))
| = 139 , | E [Equation] S [Equation]
[Equation]
G
[Equation]
S
[Equation]
| = 2182 .
Directlyanalyzing the savings by restricting number of components is not possible without assuming a graphfamily.
In Sec.
5.3 we measured the scalability for ( k, c )(  ) -S ET GNN with different number ofcomponents directly on sparse graphs, where ( k, c )(  ) -S ET WL shares similar scalability.
A practical progressively-expressive GNN: ( k, c )(  ) -S ET GNN
In this section we transform ( k, c )(  ) -S ET WL to a GNN model by replacing the HASH function in
Eq.
(5) with a combination of MLP and DeepSet [52], given they are universal function approximators for vectors and sets, respectively [22, 52].
After the transformation, we propose two additionalimprovements (Sec. 4.2 and Sec. 4.3) to further improve scalability.
We work on vector-attributedgraphs.
Let G =
( V ( G ) , E ( G ) , X ) be an undirected graph with node features x i 2 R d ,
i 2 V ( G ) .
From ( k, c )(  ) -SetWL to ( k, c )(  ) -SetGNN
( k, c )(  ) -S ET WL deﬁnes a super-graph S k,c -swl , which aligns with Eq.
We ﬁrst rewrite Eq.(5) to reﬂect its connection to S k,c -swl .
For a supernode ˆ v in S k,c -swl , let N G left (ˆ v )
= { ˆ u | ˆ u
S k,c -swl , ˆ u $ ˆ v and | ˆ u | = | ˆ v |  1 } , and N G right (ˆ v ) =
{ ˆ u | ˆ u 2 S k,c -swl , ˆ u $ ˆ v
and | ˆ u | = | ˆ v | + 1 } .
Then we can rewrite Eq.
(5) for ( k, c )(  ) -S ET WL as
[Equation]
[Equation]
where swl ( t + 12 ) k,c ( G, ˆ u ) :
= {{ swl ( t ) k,c
( G, ˆ u ) | ˆ u 2 N G right (ˆ v ) }} .
Notice that we omit HASH and applyit implicitly.
Eq.
(6) essentially splits the computation of Eq.
(5) into two steps and avoids repeatedcomputation via caching the explicit t + 12 step.
It also implies that the connection among m -sets forany m 2 [ k ] can be reconstructed from the bipartite graph among m -sets and ( m  1) -sets.
Next we formulate ( k, c )(  ) -S ET GNN formally.
Let h ( t ) ( ˆ v ) 2 R d t denote the vector representation of supernode ˆ v on the t -th iteration of ( k, c )(  ) -S ET GNN.
For any input graph G , it initializesrepresentations of supernodes by
[Equation]
where the BaseGNN can be any GNN model.
Theoretically the BaseGNN should be chosen toencode non-isomorphic induced subgraphs distinctly, mimicking HASH.
Based on empirical testsof several GNNs on all (11,117) possible 8 -node non-isomorphic graphs [6], and given GIN [49] issimple and fast with nearly 100% separation rate, we use GIN as our BaseGNN in experiments.
( k, c )(  ) -S ET GNN iteratively updates representations of all supernodes by
[Equation]
Then after T iterations, we compute the graph level encoding as
[Equation]
where POOL can be chosen as summation.
We visualize the steps in Figure 1.
Under mild conditions, ( k, c )(  ) -S ET GNN has the same expressivity as ( k, c )(  ) -S ET WL.
Theorem 4 When ( i ) BaseGNN can distinguish any non-isomorhpic graphs with at most k nodes,( ii ) all MLPs have sufﬁcient depth and width, and ( iii ) POOL is an injective function, then for anyt 2 N , t -layer ( k, c )(  ) -S ET GNN is as expressive as ( k, c )(  ) -S ET WL at the t -th iteration.
The following facts can be derived easily from Theorem 3 and Theorem 4.
Corollary 4.1 ( k, c )(  ) -S ET GNN is progressively-expressive with increasing k and c , that is,(1) when c 1 > c 2 , ( k, c 1 )(  ) -S ET GNN is more expressive than ( k, c 2 )(  ) -S ET GNN , and(2) when k 1 > k 2 , ( k 1 , c )(  ) -S ET GNN is more expressive than ( k 2 , c )(  ) -S ET GNN .
4.2 Bidirectional Sequential Message Passing
The t -th layer of ( k, c )(  ) -S ET GNN (Eq.
(8) and Eq.
(9)) are essentially propagating informationback and forth on the super-graph S k,c -swl ( G ) , which is a sequence of k  1 bipartite graphs, inparallel for all supernodes.
We propose to change it to bidirectional blockwise sequential messagepassing, which we call ( k, c )(  ) -S ET GNN ⇤ , deﬁned as follows.
[Equation]
282 Notice that ( k, c )(  ) -S ET GNN ⇤ has lower memory usage, as ( k, c )(  ) -S ET GNN load the complete supergraph ( k  1 bipartites) while ( k, c )(  ) -S ET GNN ⇤ loads 1 out of k  1 bipartites at atime, which is beneﬁcial for limited-size GPUs.
What is more, for a small, ﬁnite t it is even moreexpressive than ( k, c )(  ) -S ET GNN.
Theorem 5 For any t 2 N , the t -layer ( k, c )(  ) -S ET GNN ⇤ is more expressive than the t -layer( k, c )(  ) -S ET GNN .
As lim t !1 , ( k, c )(  ) -S ET GNN is as expressive as ( k, c )(  ) -S ET GNN ⇤ .
4.3 Improving Supernode Initialization
Next we describe how to improve the supernode initialization (Eq.
(7)) and extensively reducecomputational and memory overhead for c > 1 without losing any expressivity.
We achieve this bythe fact that a graph with c components can be viewed as a set of c connected components.
Formally,
Theorem 6 Let G be a graph with c connected components C 1 , ..., C c , and G 0 be a graph also withc connected components C 0 1 , ..., C 0 c , then G and G 0 are isomorphic if and only if 9 p : [ c ] !
[ c ] , s.t.8 i 2 [ c ] , C i and C 0 p ( i ) are isomorphic.
The theorem implies that we only need to apply Eq.
(7) to all supernodes with a single component, and for any c -components ˆ v with components ˆ u 1 , ..., ˆ u c , we can get the encoding byh (0) ( ˆ v ) =
DeepSet ( { h (0) ( ˆ u 1 ) , ..., h (0) ( ˆ u c ) } ) without passing its induced subgraph to BaseGNN.
This eliminates the heavy computation of passing a large number of induced subgraphs.
We give thealgorithm of building connection between ˆ v to its single components in Appendix A.4.
[Table] [Table]
Experiments
[Table]
We design experiments to answer the following questions.
Q1.
Performance:
How does ( k, c )(  ) -
S ET GNN ⇤ compare to SOTA expressive GNNs?
Q2.
Varying k and c : How does the progressivelyincreasing expressiveness reﬂect on generalization performance?
Q3.
Computational requirements: Is ( k, c )(  ) -S
ET GNN feasible on practical graphs w.r.t.
running time and memory usage?
5.1 Setup
Datasets.
To inspect the expressive power, we use four different types of simulation datasets:1) EXP [1] contains 600 pairs of 1&2 -WL failed graphs which we split into two where graphsin each pair is assigned to two different classes; 2) SR25 [6] has 15 strongly regular graphs ( 3WL failed) with 25 nodes each, which we transform to a 15-way classiﬁcation task; 3) Substructurecounting (i.e. triangle, tailed triangle, star and 4-cycle) tasks on random graph dataset [13]; 4) Graphproperty regression (i.e. connectedness, diameter, radius) tasks on random graph dataset [15].
Wealso evaluate performance on two real world graph learning tasks: 5) ZINC-12K [16], and 6) QM9[47] for molecular property regression.
See Table 6 in Appendix A.5 for detailed dataset statistics.
Baselines.
We use GCN [26], GIN [49], PNA ⇤ [15], PPGN [30], and GNN-AK [55] as baselineson the simulation datasets.
On ZINC-12K we also reference CIN [8] directly from literature.
Mostbaselines results are taken from [55].
Finally, we compare to GINE [23] on QM9.
Hyperparameterand model conﬁgurations are described in Appendix A.6.
5.2 Results
Table 1 shows the performance results on all the simulation datasets.
The low expressive modelssuch as GCN, GIN and PNA ⇤ underperform on EXP and SR25, while 3-WL equivalent PPGN excelsonly on EXP.
Notably, ( k, c )(  ) -S ET GNN ⇤ achieves 100% discriminating power with any k largerthan 2 and 3, and any c larger than 1 and 0, resp.
for EXP and SR25.
( k, c )(  ) -S ET GNN ⇤ alsooutperforms the baselines on all substructure counting tasks, as well as on two out of three graphproperty prediction tasks (except Radius) with signiﬁcant gap, for relatively small values of k and c .
In Table 2, we show the train and test MAEs on substructure counting tasks for individual values ofk and c .
As expected, performances improve for increasing k when c is ﬁxed, and vice versa.
It isnotable that orders of magnitude improvements on test error occur moving from k =2 to 3 for thetriangle tasks as well as the star task, while a similarly large magnitude drop is obtained at k
=4 forthe 4-cycle task, which is expected as triangle and 4-cycle have 3 and 4 nodes respectively.
Similarobservations hold for graph property tasks as well.
(See Appendix A.7)
[Table]
In addition to simulation datasets, we evaluate our ( k, c )(  ) -
S ET GNN ⇤ on real-world data; Table 3 shows our performance on
ZINC-12K, standardly used as a benchmark in expressive GNN literature, compared to several strong baselines.
Our method achieves anew state-of-the-art performance, with a mean absolute error (MAE)of 0.0750, using k =5 and c =2 .
In Table 4 we show the test and validation MAE along with trainingloss for varying k and c for ZINC-12K.
For a ﬁxed c , validation MAEand training loss both follow a ﬁrst decaying and later increasingtrend with increasing k , potentially owing to the difﬁculty in ﬁttingwith too many sets (i.e. supernodes) and edges in the super-graph.
Similar results are shown for QM9 in Table 5.
For comparison wealso show GINE performances using both 4 and 6 layers, both ofwhich are signiﬁcantly lower than ( k, c )(  )
-S
ET GNN ⇤ .
[Table]
[Table]
5.3 Computational requirements
[Figure]
We next investigate how increasing k and c change thecomputational footprint of ( k, c )(  ) -S ET GNN ⇤ in practice.
Fig.
2 shows that increasing these parameters expectedly increases both memory consumption (in MB in(a)) as well as runtime (in seconds per epoch in (b)).
Notably, since larger k and c increases our model’s expressivity, we observe that suitable choices for k and c allow us to practically realize these increasingly expressivemodels on commodity hardware.
With conservative values of c (e.g. c =1 ), we are able to consider passing messages between sets of k (e.g. 10) nodes far larger thank -WL-style higher order models can achieve (  3).
Conclusion
Our work is motivated by the impracticality of higher-order GNN models based on the k -WL hierarchy, which make it challenging to study how much expressiveness real-world tasks truly necessitate.
To this end, we proposed ( k, c )(  ) -SetWL, a more practical and progressively-expressive hierarchywith theoretical connections to k -WL and drastically lowered complexity.
We also designed and implemented a practical model ( k, c )(  ) -SetGNN
( ⇤ ) , expressiveness of which is gradually increasedby larger k and c .
Our model achieves strong performance, including several new best results ongraph-level tasks like ZINC-12K and expressiveness-relevant tasks like substructure counting, whilebeing practically trainable on commodity hardware.
Checklist
The checklist follows the references.
Please read the checklist guidelines carefully for informationon how to answer these questions.
For each question, change the default [TODO] to [Yes] , [No] , or[N/A] .
You are strongly encouraged to include a justiﬁcation to your answer , either by referencingthe appropriate section of your paper or providing a brief inline description.
For example:
• Did you include the license to the code and datasets?
[Yes]
• Did you include the license to the code and datasets?
[No] The code and the data areproprietary.•
Did you include the license to the code and datasets?
[N/A]
Please do not modify the questions and only use the provided macros for your answers.
Note that the
Checklist section does not count towards the page limit.
In your paper, please delete this instructionsblock and only keep the Checklist section heading above along with the questions/answers below.
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’scontributions and scope?
[Yes]
(b) Did you describe the limitations of your work?
[No](c) Did you discuss any potential negative societal impacts of your work?
[N/A](d)
Have you read the ethics review guidelines and ensured that your paper conforms tothem?
[Yes]2.
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[Yes](b) Did you include complete proofs of all theoretical results?
[Yes] In Appendix3.
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
[Yes] Will add laterinto the URL(b)
Did you specify all the training details (e.g., data splits, hyperparameters, how theywere chosen)?
[Yes] In Appendix
(c) Did you report error bars (e.g., with respect to the random seed after running experi-
ments multiple times)?
[Yes]
(d) Did you include the total amount of compute and the type of resources used (e.g., type
of GPUs, internal cluster, or cloud provider)?
[Yes] In Appendix
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators?
[Yes]
(b) Did you mention the license of the assets?
[No]
(c) Did you include any new assets either in the supplemental material or as a URL?
[Yes]
(d) Did you discuss whether and how consent was obtained from people whose data
you’re using/curating?
[Yes]
(e) Did you discuss whether the data you are using/curating contains personally identiﬁ-
able information or offensive content?
[N/A]
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if
applicable?
[N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable?
[N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount
spent on participant compensation?
[N/A]