Exponentially Improving the Complexity of Simulating the Weisfeiler-Lehman Test with Graph Neural Networks
Abstract
Recent work shows that the expressive power of Graph Neural Networks (GNNs) in distinguishing non-isomorphic graphs is exactly the same as that of the
Weisfeiler-Lehman (WL) graph test.
In particular, they show that the WL test can be simulated by GNNs.
However, those simulations involve neural networks for the “combine” function of size polynomial or even exponential in the number of graph nodes n , as well as feature vectors of length linear in n .
We present an improved simulation of the WL test on GNNs with exponentially lower complexity.
In particular, the neural network implementing the combine function in each node has only polylog( n ) parameters, and the feature vectors exchanged by the nodes of GNN consists of only O (log n ) bits.
We also give logarithmic lower bounds for the feature vector length and the size of the neural networks, showing the (near)-optimality of our construction.
Introduction
Graph Neural Networks (GNNs) have become a popular tool for machine learning on graphstructured data, with applications in social network prediction [HYL17], trafﬁc prediction [YYZ18], recommender systems [YHC + 18], drug discovery [WKK + 20], computer vision [LGD + 19,
FLM + 19, MBM + 17, QSMG17] and combinatorial optimization [CCK + 21].
Standard message passing GNNs use the topology of the input graph to deﬁne the network structure: in each step k , a node aggregates messages from each of its neighbors and combines them using a function  ( k ) , computed by a neural network, to determine its message for the next round.
Crucially, the aggregation function must be symmetric , to ensure that the output of GNNs is invariant under node permutation.
This restriction raised questions about how expressive such network architectures are, and in particular what classes of graphs are distinguishable using GNNs.
The seminal works of Xu et al.
[XHLJ19] and Morris et al.
[MRF + 19] (see also [Gro21]) showed that
GNNs are exactly as powerful in distinguishing graphs as the Weisfeiler-Lehman (WL) test [WL68], also known as color reﬁnement .
This combinatorial procedure is a necessary but not sufﬁcient test for graph isomorphism.
It proceeds in repeated rounds: in each round, a node labels itself with the “hash” of the multiset of labels of its neighbors.
The aforementioned papers show that (i) GNNs can simulate the WL test and (ii) GNNs can only distinguish those graphs that the WL test determines to be different.
This provides a complete characterization of the expressive power of GNN architectures.
[Table]
Xu et al.
[XHLJ19]
[Table]
Morris et al.
[MRF + 19]
[Table]
This paper
[Table]
The connection between GNNs and the WL test has spawned a wave of new results studying GNN variants that either match the distinguishing power of the WL test or adopt new methods beyond message passing on the edges of the input graph to overcome this barrier (see the excellent surveys [MFK21, MLM + 21, Gro21] for an overview of this area).
However, the results in the original as well the follow up works mostly focus on qualitative questions (how expressive GNNs are) as opposed to quantitative questions such as the network complexity.
In particular, while Xu et al.
[XHLJ19] show that there exist GNN architectures that can simulate the WL coloring procedure as long as the aggregation step is injective, they rely on the universal approximation theorem to show that there exists a neural network that can simulate the hash function used in WL.
As a result, the size of the network could be exponential in the number of nodes n .
In contrast, the construction of
Morris et al.
[MFK21] uses networks of size polynomial in n .
However, the weights of the network implementing  ( k ) in their construction depend on the structure of the underlying graph, which sufﬁces for node classiﬁcation, but is not sufﬁcient for the context of graph classiﬁcation.
Overall, the quantitative understanding of the complexity of simulating WL remains an open problem.
Indeed, the survey [Gro21] states “The size of the GNNs and related parameters like depth and width, which directly affect the complexity of inference and learning, deﬁnitely require close attention .”
Our Results.
The main question addressed in this work is: what is the simplest (in terms of the number of neural network units and message length) GNN capable of simulating the WL test?
Equivalently, at what point does a GNN become so small that it loses its expressive power?
Our main result is a highly efﬁcient construction of a GNN architecture that is capable of simulating the WL test.
For graphs with n nodes, it can simulate poly ( n ) steps of the WL test, such that the neural network implementing  ( k ) in each round has polylog( n ) parameters, and the messages exchanged by the nodes of the GNN in each round consist of O (log n ) bits.
This offers at least an exponential improvement over the prior bounds obtained in [XHLJ19, MRF + 19] (see Table 1), extending the equivalence between the WL test and the expressive power of GNNs to neural networks of reasonable (in fact, quite small) size.
Furthermore, our architecture is simple, using vector sum for aggregation and ReLU units for the combine function  ( k ) .
Finally, our construction can be generalized to yield a depth-size tradeoff: for any integer t > 0 , we can can construct a neural network of depth O ( t ) and size n O (1 /t ) polylog( n ) .
To achieve this result, our construction is randomized, i.e., some weights of the neural networks are selected at random, and the simulation of the WL test is correct with high probability 1  1 / poly ( n ) .
Thus, our construction can be viewed as creating a distribution over neural networks computing the function  ( k ) .
2 In particular, this implies that, for each graph, there exists a single neural network implementing  ( k ) that accurately simulates WL on that graph.
The size of the network is exponentially smaller than in [MRF + 19], although the construction is probabilistic.
We complement this results with two lower bounds for executing a WL iteration.
Our ﬁrst lower bound addresses the communication complexity of this problem, and demonstrates that to solve it, each node must communicate labels that are at least O (log n ) bits long, matching the upper bound achieved by our construction.
Our second lower bound addresses the computational complexity , namely the parameters of the neural network.
It shows that if the messages sent between nodes are vectors with entries in [ F ] = { 0 , 1 , . .
, F  1 } , then the network implementing  ( k ) must use ⌦ (log F ) ReLU units.
Related work.
The equivalence between the discriminative power of GNNs and the WL test has been shown in the aforementioned works [XHLJ19, MRF + 19].
A strengthened version of the theorem of [XHLJ19], where the same combine function  is used in all iterations (i.e.,  =  ( k ) for all k ) appeared in the survey [Gro21].
Many works since have studied various representational issues in GNNs; we refer the reader to excellent surveys [Gro21, HV21, MLM + 21, Jeg22].
In particular, [Lou19] established connections between GNNs and distributed computing models such as LOCAL and CONGEST, and derived lower bounds for several computational tasks based on this connection. [CVCB19] drew a connection between the expressiveness of GNNs in graph isomorphism testing and in function approximation.
[BKM + 20] studied the expressiveness of GNNs in computing Boolean node classiﬁers, and [GMP21] studied the expressiveness of graph convolutional networks (GCNs).
The emergence of WL as a barrier in GNN expressivity has also led to a ﬂurry of work on enhancing their expressivity by means of more general architectures.
These include higher-order GNNs inspired by higher-dimensional analogs of WL [MRF + 19, MBHSL19], unique node identiﬁers [Lou19, VLF20], random node initializations [ACGL21, SYK21], relational pooling [MSRR19], incorporating additional information on the graph structure [NM20, BGRR21, BFZB22, CMR21,
TRWG21], and more.
We refer to [MLM + 21] for a comprehensive survey of this line of work.
1.1 Preliminaries
Notation.
For the rest of the paper, we use N ( v ) to denote the neighborhood of v in a graph
G ( V, E ) including v itself , and we use {·} to denote multisets rather than sets.
GNNs.
Let G =
( V, E ) be a graph with N nodes.
GNNs use the graph structure of G to learn node embeddings for all the nodes across multiple iterations.
Let h ( k ) v denote the embedding vector of node v 2 V in the k th iteration.
The vectors h (0) v represent the initial node embeddings.
In every iteration k  1 , each node v sends its current embedding h ( k  1) v to all its neighbors, and then computes its new embedding h kv by the equation
[Equation]
where  ( k ) is implemented by a neural network with ReLU activations (note that  ( k ) may differ across different iterations k ).
The function f is called the ‘aggregate’ function, and  ( k ) is called the ‘combine’ function.
The embeddings h ( K ) v at the ﬁnal iteration K can be used for node classiﬁcation.
For graph classiﬁcation, they can be aggregated into a graph embedding h G with a ‘readout’ function,
[Equation]
Weisfeiler-Lehman (WL) Test.
The WL test [WL68] is a popular heuristic for the graph isomorphism problem.
While the exact complexity of this problem remains unknown [Bab16], the WL test is a powerful heuristic capable of distinguishing a large family of graphs [BK79].
The WL test is as follows.
Given a graph G ( V, E ) , initially all nodes are given the same ﬁxed label ` (0) v , say ` (0) v = 1 for all v 2 V .
Then, in every iteration k  1 , each node v is assigned the new label ` ( k ) v = HASH ( { ` ( k  1) u : u 2 N ( v ) } ) , where HASH is a 1-1 mapping (i.e., different multisets of labels are guaranteed to be hashed into distinct new labels).
Performing this procedure on two graphs G, G 0 , the WL test declares them non-isomorphic if the label multiset { ` ( k ) v : v 2 V } differs between the graphs at some iteration k .
Note that the algorithm converges within at most n iterations.
Simulating WL with GNNs.
A GNN simulates WL deterministically if the node embeddings h ( k ) v at each iteration k constitute valid labels ` ( k ) v for the WL test.
We also consider randomized simulation, where the GNN’s weights are selected at random
and we allow some small failure probability where distinct multisets of labels from iteration k  1 are hashed into the same label at iteration k .
This is captured by the next deﬁnition.
Deﬁnition 1.1 (Successful Iteration of the WL Test) .
A WL iteration gets existing labels { h v : v 2 V } for all nodes, and outputs new labels { h 0 v : v 2 V } given by h 0 v
=  ( f ( { h w : w 2 N ( v ) } )) , for an aggregate function f and neural network  with random weights.
We say the iteration is successful if for all v, u 2 V , the following holds:
• If { h w : w 2 N ( v ) } = { h w : w 2 N ( u ) } then h 0 v = h 0 u with probability 1, and •
If { h w : w 2 N ( v ) } {+
+} = { h w : w 2 N ( u ) } then h 0 v 6 = h {+
+} u with probability 1  p ,
where the probability is over the choices of the random weights of  .
To ensure the WL simulation is successful across poly( n ) iterations and for all pairs of nodes, we can set failure probability p to 1 / poly( n ) and apply a union bound (see Prop. A.1 in Appendix A).
1.2 Overview of Our Techniques
In this section we give an overview of our GNN architectures for simulating WL.
To explain our ideas in stages, we begin with a simpler construction of a polynomial size GNN.
It is far larger than the ultimate polylogarithmic size we are aiming for, but forms a useful intermediate step toward our second and ﬁnal construction.
Construction 1 (Section 2).
Recall that the k th WL iteration, for a node v , aggregates the labels of its neighbors from the previous iteration, H ( k  1) v :=
{ h ( k  1) w : w 2 N ( v ) } , and hashes them into a new label h ( k ) v for v .
Our GNNs aggregate by summing, i.e., they sum H ( k  1) v into S ( k ) v :=
P w 2 N ( v ) h ( k  1) w , and then hash the sum into h ( k ) v using a ReLU neural network of our choice.
If two nodes u, v satisfy H ( k  1) u 6 = H ( k  1) v , then WL assigns them distinct labels in iteration k , and therefore we want our construction to satisfy h ( k ) u
6 = h ( k ) v with probability at least 1  p (as per
Deﬁnition 1.1).
This can fail in two places: either due to summing (if S ( k ) u
= S ( k ) v even though
H ( k  1) u
6 = H ( k  1) v ) or due to hashing (
if h ( k ) u = h ( k ) v even though S ( k ) v 6 = S ( k ) v ).
To avoid the ﬁrst failure mode (summing), we use one-hot encodings for the node labels, meaning that for every node v and iteration k , h ( k ) v is a one-hot vector in { 0 , 1 } F (for F to be determined shortly).
This ensures that if H ( k  1) u
[Equation] then S ( k ) u 6 = S ( k ) v due to the linear independence of the one-hot vectors.
To handle the second failure mode (hashing), we use random hashing into F
= O (1 /p ) buckets, rendering the collision probability no more than p .
That is, we let h ( k ) v =
One-Hot ( g ( S ( k ) v ))
where g : Z F ! { 1 , . .
. , F } is chosen at random from an (approximately) universal hash family.
We use the simple hash function g ( x ) =
h
a, x i mod F where a is random vector.
It can be implemented with a ReLU neural network, since the dot product is just a linear layer, and the mod operation can be implemented with constant width and logarithmic depth, by a reduction to the “triangular wave” construction due to Telgarsky [Tel16] (see Appendix B.1).
Finally, we show that turning the index of the hash bucket into a one-hot vector can be done with O ( F ) ReLU units and constant depth.
Since our desired collision probability is p = 1 / poly( n ) , we set F = O (1 /p )
= poly( n ) .
Since our one-hot vectors are of dimension F , the resulting GNN has width F = poly( n ) (and depth O (log n ) due to implementing the mod operation with ReLUs), and thus total size poly( n ) .
Construction 2 (Section 3).
We now proceed to describe our polylogarithmic size construction.
The weak point in the previous construction was the wasteful use of one-hot encoding vectors, which caused the width to be F .
In the current construction, we still wish to hash into F bins — that is, to have F distinct possible labels in each iteration — but we aim to represent them using bitstrings of length O (log F ) , thus exponentially improving the width of the network.
That is, for every node v and iteration k , the label h ( k  1) v would now be a vector in { 0 , 1 } O (log F ) .
The challenge is again to avoid the two failure modes above, ensuring that the failure probability does not exceed p .
Since we cannot use one-hot encoding, we need to devise another method to avoid the ﬁrst failure mode, i.e., ensure that with high probability S ( k ) u 6 = S ( k ) v
[Equation]
To this end, suppose for a moment that we had access to a truly random vector a 2 { 0 , 1 } F .
Then each node v , instead of sending the one-hot encoding of its label h
( k  1) v , we could instead send [Figure] the dot product h a, h ( k  1) v i , which is the single bit a h ( k  1) v .
Each node u thus receives the bits n a h ( k  1) w : w 2 N ( w ) o from its neighbors and aggregates them into the sum
P
[Equation]
i , which, by linearity, is equal to h a, S ( k ) u i (using the notation from Construction 1).
It is easy to observe that if S ( k ) u 6 = S ( k ) v then
h a, S ( k ) u i 6 = h a, S ( k ) v i with probability at least 0 .
Repeating this process log F independent times decreases the collision probability to the requisite 1 /F .
Thus, we can deﬁne a new labeling scheme ¯ h ( k ) v 2 { 0 , 1 } log F that concatenates log F dot products with independent random vectors a 1 . . .
a log F 2 { 0 , 1 } F as just described, failing at the summing operation with probability at most 1 /F .
The second failure mode (hashing) can again be handled as before.
That catch is that, since a has length F , the overall number of parameters in the GNN would again be at least F .
To avoid this, we appeal to the computational theory of pseudorandomness.
The idea is to replace the random vector a with an efﬁcient pseudorandom analog.
Note that the above approach goes through even if the probability that h a, S ( k ) u i 6 = h a, S ( k ) v i is slightly larger than 0 .
5 , say 0 .
5 + ✏ for a small constant ✏ .
It is well-known in complexity theory that there exist pseudorandom generators, called ✏ -biased spaces , that generate vectors a satisfying this property given only
O (log F + log(1 / ✏ )) truly random bits.
3 Crucially, each bit of a can be computed using a threshold circuit with size polynomial in O (log F + log(1 / ✏ )) and constant depth (Theorem C.1), which translates to a ReLU neural network with the same parameters (Lemma C.2).
Using these generators in our GNN to implement a pseudorandom ( ✏ -biased) analog of a yields our ﬁnal construction.
First Construction: Polynomial-size GNN
Our ﬁrst construction towards Deﬁnition 1.1 is exponentially larger compared to our ﬁnal optimized construction of Section 3.
Nevertheless, it is instructive and motivates our optimized construction.
Let h ( k ) u denote the label of a vertex u in the k th iteration.
For our ﬁrst construction, we will always maintain the invariant that h ( k ) u will be a one-hot encoded vector in { 0 , 1 } F for all u 2 V and all iterations k .
F > 2 n will be a prime which also satisﬁes F
= O (poly( n )) .
As stated previously, the aggregate function f will just be the sum function.
Our construction for the neural network used in the k th iteration,  ( k ) , will take in the sum of the neighbors labels according to Equation (1.1) and output a one-hot encoded vector in { 0 , 1 } F .
Implementation of Neural Network.
Our construction for  ( k ) is the following: First recall the notation from (simpliﬁed) Equation 1.1:
[Equation]
For every node v , the input to  ( k ) is the sum of feature vectors of neighbors from the prior iteration,
[Equation]
, which is returned by the summing function f .
Given any such input x 2 Z F ,  ( k ) computes the inner product of x with a vector a 2 Z F where each entry of a is a uniformly random integer in [ F ] := { 0 , . .
. , F  1 } . 2.
( k ) then computes h x, a i mod F . 3.
Finally, we represent the value of z = h x, a i as a one-hot encoded vector in Z F where the z -th entry is equal to 1 and all other entries are 0 ’s.
Altogether,  ( k ) can be summarized as: h ( k ) v = One-Hot
[Equation]
Note that we set the initial labels h (0) u to be the same starting vector for all vertices (any onehot vector).
This matches the WL test which also initializes all nodes with the same initial label.
Furthermore, the weights of  ( k ) are independent: the random vector a is sampled independently for each iteration.
Correctness of Construction.
The following lemma proves that the above construction satisﬁes the requirement of Deﬁnition 1.1.
Its proof is given in Appendix B.
Lemma 2.1.
Let { h ( k  1) w : w 2 N ( v ) } and { h ( k  1) w : w 2 N ( u ) } denote the multiset of neighbor- hood labels for vertices v and u respectively.
If the multisets are distinct then the labels computed for v and u in the k th iteration are the same with probability at most O (1 /F ) .
If the multisets are the same then the labels are the same, i.e., the k th iteration is successful according to Deﬁnition 1.1.
Complexity of the GNN.
We now evaluate the size complexity of implementing our construction via a neural network  .
Note that Step 1 of the construction can be done with 1 layer as it simply involves taking an inner product.
The main challenge is to implement the modulo function.
We give the following construction in Section B.1 of the appendix.
Theorem 2.2.
Suppose F = poly( n ) .
There exists an explicit construction of a network which computes modulo F in the domain { 0 , . . .
, nF } using a ReLU network with O (log n ) hidden units and O (log n ) depth.
More generally, given an integer parameter t > 0 , the function can be computed with O (( Fn ) O
(1 /t ) log n ) hidden units and O ( t ) depth.
Directly appealing to the theorem above, we can implement modulo F required in Step 2 of the construction using a neural network with O (log n ) units, depth O (log n ) .
In addition, we need only
O (log n ) bits to represent the weights.
Finally, Step 3 of our construction requires outputting a one hot encoding.
We can do this by inputting z (the output of Step 2 of the construction) into F indicator functions, each of which detect if z is equal to a particular integer in [ F ] .
Each indicator function can be implemented via O (1) ReLU nodes as follows.
Let
[Equation]
which can be easily implemented as a ReLU network.
(Note min( a, b ) =
a + b  max( a, b ) and max( a, b ) = max( a  b, 0)
+ b = ReLU ( a  b ) + b .)
It can be checked that g (0) = 1 and g ( x ) =
for all other integers x 6 = 0 and that g can be implemented with O (1) ReLU function compositions.
Thus, Step 3 of the construction requires O (1) hidden layers and O ( F ) total hidden units.
Altogether we have proven the following result.
Theorem 2.3.
There exists a construction of a neural network  which performs a successful iteration according to Deﬁnition 1.1 with failure probability p =
O (1 / | F | ) .
has depth O (log n ) , O ( F ) hidden units, and requires O (log( nF )) bits of precision.
Furthermore, all labels in all iterations are vectors in { 0 , 1 } F .
More generally, given an integer parameter t > 0 , the function can be computed with n O (1 /t ) polylog( n ) hidden units and O ( t ) depth.
Remark 2.4.
In the standard WL test, the number of iterations is chosen to be O ( n ) .
Thus the right setting of F in Theorem 2.3 is F
= O (poly( n )) which gives us depth O (log n ) , O (poly( n )) hidden units, and requires O (log n ) bits of precision in addition to labels in dimension O (poly( n )) .
Second Construction: Polylogarithmic-size GNN via Pseudo-randomness
We now present a more efﬁcient construction of a GNN which simulates the WL test with an exponential improvement in the number of hidden units and label size.
To motivate the improvement, we consider Step 3 of the prior construction which outputs a one-hot encoding.
The one-hot encoding was useful as it allowed us to index into a uniformly random vector a (which we then sum over mod
F in order to hash the neighborhood’s labels).
However, this limited us to use feature vectors of a large dimension and required many hidden units to create one-hot vectors.
Instead of working with one-hot encodings as an intermediary, we will directly compute the entries of the random vector a as needed.
This has two advantages: we can signiﬁcantly reduce the dimension of the feature vectors as well as reduce the total size of the neural networks used.
We accomplish this via using pseudo-random vectors whose entries can be generated as needed with a small ReLU neural network (see Corollary 3.3).
This allows us to use node labels in dimension O (log n ) as opposed to O ( n ) .
The random vectors we employ have their entries generated from an " -biased sample space.
These are random vectors which are approximately uniform and they have been well-studied in the complexitytheory literature.
We recall some deﬁnitions below.
Deﬁnition 3.1 (Bias) .
Let X be a probability distribution over { 0 , 1 } m .
The bias of X with respect to a set of indices I ✓ { 1 , . . .
, m } is deﬁned as bias I ( X ) =
[Equation] =
⇤ where each sum is taken modulo 2 and the empty sum is deﬁned to be 0 .
Deﬁnition 3.2 ( " -biased Sample Space) .
A probability distribution over { 0 , 1 } m is called an " -biased sample space if bias I ( X )  " holds for all non-empty subsets I ✓ { 1 , . . .
, m } .
Note that the uniform distribution has bias 0 .
We now state our construction for the neural network used in the k th iteration,  ( k ) .
We recall that h ( k ) u denotes the label of a vertex u in the k th iteration.
Implementation of Neural Network.
Our construction for  ( k ) is the following:
The input vectors of  ( k ) , which are of the form h ( k  1) v for the k th iteration, are each assumed to be a feature vector in Z C log n for a sufﬁciently large constant C .
The output of  ( k ) will also be a feature vector in Z C log n .
Our aggregation function f will again be the summation function.
Let F be a prime of size poly( n ) which is at least 2 n . 3.
For each node v ,  ( k ) computes z v = h b,
[Equation]
i mod F where every entry of b is uniformly random in { 0 , . .
. , F  1 } .
Note P
[Equation]
( k  1) w is the output of the aggregation f . 4.
Let a t 2 { 0 , 1 } F for t = 1 , . .
, C log n be vectors which are independently drawn from an " -biased sample space for a sufﬁciently small constant " .
The output h ( k ) v will be a C log n dimensional binary vector where the t -th coordinate is equal to the z v -th coordinate of the vector a t .
In other words, h ( k ) v = ( a t ( z v ))
C log n t = where a t ( z v ) denotes the z v -th coordinate of a t .
Correctness of Construction.
We now prove the correctness of our construction.
We will refer to z v computed in Step 3 of the construction as the index of v for the k th iteration.
To prove the correctness of the above construction, it sufﬁces to prove the lemma below which shows our construction satisﬁes Deﬁnition 1.1.
Lemma 3.1.
Let { h ( k  1) w : w 2 N ( v ) } and { h ( k  1) w : w 2 N ( u ) } denote the multiset of neighbor- hood labels for vertices v and u respectively.
If the multisets are distinct then the labels computed for v and u in the k th iteration are distinct with probability 1  1 / poly( n ) .
If the multisets are the same then the labels are the same, i.e., the k th iteration is successful according to Deﬁnition 1.1.
We ﬁrst need the following auxiliary lemma about " -biased sample spaces, proven in Section C.
Lemma 3.2.
Let D be a probability distribution over { 0 , 1 } m that is an " -biased sample space.
Then, for any x, y 2 Z m such that x 6 =
y, P a ⇠ D [ h a, x i =
h a, y i ]  12 + " 2 .
Note that this lemma is necessary, as we will be computing dot products of a with integer vectors (over integers), not with binary vectors modulo 2 .
We are now ready to prove Lemma 3.1.
Proof of Lemma 3.1.
Let x 0 =
[Equation]
denote the input for v and analogously, deﬁne y
[Equation]
to be the input for u .
We ﬁrst show that if { h ( k  1) w : w 2 N ( v ) } is not equal to (as multisets)[-{ h ( k  1) w : w 2 N ( u ) } then x 0 
-]
{ h ( k  1) w : w 2 N ( u ) } then x 0 6 = y 0 with sufﬁciently large probability.
We further consider the case that k  2 since for k = 1 (the ﬁrst iteration), the statement follows since all node labels are initialized to be the same.
Let z
v be the indices computed in Step 3 of iteration k
(which are used to construct the node labels h
( k  1) v in iteration k  1 ).
Note that there is a one to one mapping between z 0 v and h ( k  1) v .
Thus we can assume { z 0 w : w 2 N ( v ) } 6 = { z [Equation] w 2 N ( u ) } without loss of generality.
Let ˜ a 1 be the ﬁrst " -biased vector used in the previous neural network  ( k  1) .
Note that the ﬁrst entry of x 0 is equal to
P w 2 N ( v ) ˜ a 1 ( z w ) , i.e., it is the dot product of ˜ a 1 with a suitable vector x which is the sum of one-hot encoding of the neighborhood of v .
The same statement is true for y 0 : the ﬁrst entry of y 0 is equal to the dot product of ˜ a 1 with a vector y which represents the one-hot encoding of the neighborhood of u .
This is because we computed the index z 0 v of every node in the previous iteration, as deﬁned in Step 3 of the construction, and passed along the coordinate of ˜ a 1 which corresponds to this computed index.
By assumption, we know that x 6 = y .
Therefore by
Lemma 3.2, we have that P ˜ a 1 [( x 0 )
= ( y 0 ) 1 ] =
P ˜ a 1 ( h ˜ a 1 , x i =
h ˜ a 1 , y i )  2 / 3 for a suitable " .
By independence of vectors ˜ a t , it follows that P [ x 0 = y 0 ]  (2 / 3) C log n  1 / poly( n ) for a suitable constant C > 0 .
We now condition on x 0 6 = y 0 .
Without loss of generality, suppose that their ﬁrst coordinates, x 0and y 0 1 , differ.
We know x 0 1
= y 0 1 mod F since x 0 1 6 = y 0 1 , they are both non-negative and bounded by n , and | x 0 1  y 0 1 |  O ( n ) whereas F is a prime at least 2 n .
It follows that the probability of the event h ( x 0  y 0 ) , b i =
mod F is at most 1 /F .
To see this, condition on all the entries of b except b 1 .
Then ( x 0  y 0 ) 1 · b 1 must be equal to a speciﬁc value modulo F for h ( x 0  y 0 ) , b i =
mod F to hold, as desired.
We now condition on this event which equivalently means we condition on z v 6 =
z u (see Step 3 of the construction).
Now our task is to show that h ( k ) v 6 = h ( k ) u with sufﬁciently high probability.
The ﬁrst coordinate of h ( k ) v is equal to h a 1 , e z v i where a 1 is the ﬁrst " -biased vector considered in Step 4 of the construction and e z v is the basis vector in dimension C log n which has a 1 entry only in the z v coordinate and 0 otherwise.
Therefore by Lemma 3.2, we have that , P a 1 [( h ( k ) v ) 1 =
( h ( k ) u ) 1 ]
= P a
[ h a 1 , e z v i
=
h a 1 , e z u i ]  2 / 3 for a suitable choice of " since e z u and e z v are distinct.
By independence of vectors a t , it follows that P [ h ( k ) v = h ( k ) u ]  (2 / 3) C log n  1 / poly( n ) for a suitable constant C > 0 .
This exactly means that the node labels of u and v in the next iteration are different with probability at least 1  1 / poly( n ) , as desired.
Lastly, it is clear that if the multiset of neighborhood labels of u and v are the same, then x 0 = y
and it’s always the case that h ( k ) v = h ( k ) u .
Complexity of the GNN.
We now analyze the overall complexity of representing  ( k ) as a ReLU neural network.
First we state guarantees on generating " -based vectors using a ReLU network.
The following corollary is proven in Appendix C.
Corollary 3.3.
Let s =
O (log F + log(1 / " )) .
For every " and F , there exists an explicit ReLU network C : { 0 , 1 } s [ [ F ] ! { 0 , 1 } F which takes as input s uniform random bits and an index i 2 F and outputs the i th coordinate of an " -biased vector in { 0 , 1 } F .
C uses O (log F ) bits of precision and has poly( s ) hidden units.
More generally, given an integer parameter t > 0 , the function can be computed with n O (1 /t ) polylog( n ) hidden units and O ( t ) depth.
We can now analyze the complexity of our construction.
The complexity can be computed by analyzing each step of the construction separately as follows:
For every node v , the sum of feature vectors of neighbors from the prior iteration, P w 2 N ( v ) h
( k  1) w , is returned by the aggregate function f . 2.
The inner product with the random vector b in Step 3 of the construction can be computed using one layer of the network.
Then computing modulo F can be constructed via Theorem 2.2. 3.
Given the inner product value z v which is the output of Step 3 of the construction, we compute all of the O (log n ) coordinates of h ( k ) v in parallel.
We recall that each coordinate of h ( k ) v is indexing onto O (log n ) " -biased random vectors and we use the same index for all vectors, namely the z v -th index.
This can be done as follows.
We ﬁrst have O (log n ) edges fanning-out from the node which computes z v .
For all t = 1 , . . .
, O (log n ) , the other endpoint of the t -th fan-out edge computes the value a t ( z v ) where a t is the t -th " -biased vector as stated in Steps 4 and 5 of the construction.
This can be done by appealing to the construction guaranteed by Corollary 3.3.
The result of this computation is exactly h ( k ) v .
Altogether, we have proven the following theorem.
Theorem 3.4.
There exists a construction of  which performs a successful WL iteration according to
Deﬁnition 1.1 with p  1 / poly( n ) .
has depth O (log n ) , O (poly(log n )) hidden units, and requires
O (log n ) bits of precision.
All labels in all iterations are binary vectors in { 0 , 1 } O (log n ) .
More generally, given an integer parameter t > 0 , the function can be computed with n O (1 /t ) polylog( n ) hidden units and O ( t ) depth.
Lower Bounds
We complement our construction with lower bounds on the label size and number of ReLU units required to simulate the WL test.
We outline these two lower bounds below and defer the full details to Appendix D.
Message Size.
Recall that in our construction, the message (label) size was O (log n ) bits.
Via communication complexity, we give a corresponding lower bound.
In particular, we construct a graph on which any (randomized) communication protocol which simulates WL as in Deﬁnition 1. must send at least ⌦ (log n ) bits along one edge of the graph.
As message-passing GNNs are a speciﬁc class of communication protocols, this immediately implies that the message sizes must have ⌦ (log n ) bits, so our construction is optimal in that respect.
The hard instance is formed by a graph which is a collection of disjoint star subgraphs of sizes ranging from 2 to ⇥ ( p n ) .
In order to perform a valid WL coloring, each node must essentially learn the size of its subgraph, requiring ⌦ (log( p n ))
=
⌦ (log n ) bits of communication.
In addition, this must be done in only 2 iterations as the depth of each subgraph is 2 , so some node must send ⌦ (log n ) bits to its neighbors in a single round.
See Appendix D.1 for the full details and proof.
Number of ReLU Units.
In order to show a lower bound on the number of units needed to implement a successful WL iteration, we rely on prior work lower bounding the number of linear regions induced by a ReLU network (for instance [MPCB14]).
In particular, these works show that
ReLU networks induce a partition of the input space into N convex regions (where N is a function of the size of the network) such that the network acts as a linear function restricted to any given region.
Using these results, we describe a ﬁxed graph and a distribution over inputs to the neural network
S ( k  1) u for all u 2 V (sums of the labels from the previous round) which includes O ( F ) potential special pairs of nodes (where F is deﬁned such that inputs S ( k  1) u 2 [ F ] t for some t ).
For each such pair u, v , their neighborhoods N ( u ) , N ( v ) have different multisets of inputs, but both multisets of inputs sum to the same value.
We show that if the number of linear regions is small, N = o ( F ) , then it is relatively likely that u, v will be in the same linear region and thus their sums will collide:
S ( k ) u = S ( k ) v even while their neighborhoods had distinct inputs in the ( k  1) st round.
This immediately gives a ⌦ (log F ) lower bound on the number of ReLU units (and thus number of parameters) with more reﬁned depth/width tradeoffs given in Appendix D.2.1.
Note that F is the size of each coordinate in the sum of labels.
Even if the labels are binary, F can be as large as n , depending on the max degree in the graph, which implies a ⌦ (log n ) lower bound on the number of
ReLU units.
See Appendix D.2 for full details and proof.  Experiments
To demonstrate the expressivity of our construction, i.e., that our small-sized GNN reliably simulates the WL test, we perform experiments on both synthetic and real world data sets.
Common to all of our experiments is that we start with some graph G =
( V, E ) (either real world or generated with respect to some probability distribution).
We then simulate a perfect run of the WL test on G where any two nodes which receive different multisets of labels in iteration k  1 get distinct labels in iteration k with probability 1 as well as a run of our construction from 4 Section 3.
At any point in time, the node labels induce partitions of V where two nodes are in the same class if they have the same labels.
Denote the partitions after k -iterations using the perfect simulation and our construction respectively by P k and P 0 k .
Letting k 0 be minimal such that P k 0  1 =
P k 0 (at which point the WL labels have converged), we consider the implementation using our GNN successful if
P k = P 0 k for all k  k 0 , i.e., if the the simulation using our implementation induced the same partitions as a perfect runs.
For all of our experiments it turned out that k 0  5 (see [BK22] for a discussion of this fast convergence).
[Figure]
Sparse Erd˝os-Rényi Graphs We generated Erd˝os-Rényi random graphs G ( n, p ) with p = 20 /n for a varying number of vertices n .
For each value of n , we generated ﬁve such graphs and for each of these ﬁve graphs, we ran 10 independent trials of our GNN implementation with message sizes t = 1 , 2 , . . . .
Averaging over the ﬁve graphs, we report the minimal t such that at least 70% of the 10 iterations successfully simulated the WL test.
See Figure 2a.
The average message size needed to achieve this is approximately 3 log n where the logarithmic dependence on n is as predicted theoretically and signiﬁcantly improves on the linear message size required for prior constructions.
Scale Free Graphs We generated samples of the scale free graphs from [BBCR03] with a varying number of vertices n using the implementation from [HSS08].
Our experiment design was the same as for Erd˝os-Rényi random graphs.
See Figure 2b.
Cora Graph We ﬁnally ran experiments on the real world graph Cora 5 which is the citation network of n = 2708 scientiﬁc publications.
We simulated our GNN with varying message lengths, for each message length reporting the fraction of successful runs of 30 independent trials.
See Figure 2c for a plot of the results.
We see that with message length 35 , all of the 30 trials successfully simulated the
WL test.
Acknowledgements Anders Aamand is funded by DFF-International Postdoc Grant 0164-00022B from the Independent Research Fund Denmark.
This research was also supported by the NSF
TRIPODS program (award DMS-2022448), NSF award CCF-2006664, Simons Investigator Award,
MIT-IBM Watson collaboration, GIST- MIT Research Collaboration grant, NSF Graduate Research
Fellowship under Grant No. 1745302, and MathWorks Engineering Fellowship.
Checklist
For all authors... (a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’s contributions and scope?
[Yes] (b) Did you describe the limitations of your work?
[Yes] (c) Did you discuss any potential negative societal impacts of your work?
[N/A]
This paper is focused on fundamental research on graph neural networks with no direct social impacts.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
[Yes] 2.
If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?
[Yes] (b) Did you include complete proofs of all theoretical results?
[Yes] 3.
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)?
[N/A] (b) Did you specify all the training details (e
g., data splits, hyperparameters, how they were chosen)?
[N/A] (c) Did you report error bars (e
g., with respect to the random seed after running experiments multiple times)?
[N/A] (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)?
[N/A]
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?
[N/A] (b) Did you mention the license of the assets?
[N/A] (c) Did you include any new assets either in the supplemental material or as a URL?
[N/A] (d) Did you discuss whether
and how consent was obtained from people whose data you’re using/curating?
[N/A] (e) Did you discuss whether the data you are using/curating contains personally identiﬁable information or offensive content?
[N/A] 5.
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]