Exponentially Improving the Complexity of Simulating the Weisfeiler-Lehman Test with Graph Neural Networks

Introduction

Graph Neural Networks (GNNs) have become a popular tool for machine learning on graph-structured data, with applications in social network prediction [HYL17], traffic 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 define the network structure: in each step $k$, a node aggregates messages from each of its neighbors and combines them using a function $\phi^{(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 refinement. This combinatorial procedure is a necessary but not sufficient 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.

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 $\phi^{(k)}$ in their construction depend on the structure of the underlying graph, which suffices for node classification, but is not sufficient for the context of graph classification.

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, definitely require close attention.”

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 efficient construction of a GNN architecture that is capable of simulating the WL test. For graphs with $n$ nodes, it can simulate $\mbox{poly}(n)$ steps of the WL test, such that the neural network implementing $\phi^{(k)}$ in each round has $\operatorname{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 $\phi^{(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)}\operatorname{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/\mbox{poly}(n)$. Thus, our construction can be viewed as creating a distribution over neural networks computing the function $\phi^{(k)}$.Note that selecting $\phi^{(k)}$ at random is quite different from random node initialization, e.g., as investigated in [ACGL21]. In particular, in our model all nodes use the same function $\phi^{(k)}$ (with the same parameters), without breaking the permutation invariance property of GNNs, as in the standard GNN model. In particular, this implies that, for each graph, there exists a single neural network implementing $\phi^{(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 first 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,\dots,F-1\}$, then the network implementing $\phi^{(k)}$ must use $\Omega(\log F)$ ReLU units.

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 $\phi$ is used in all iterations (i.e., $\phi=\phi^{(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 classifiers, 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 flurry 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 identifiers [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 Preliminaries

For the rest of the paper, we use $\mathcal{N}(v)$ to denote the neighborhood of $v$ in a graph $G(V,E)$ including $v$ itself, and we use $\{\cdot\}$ to denote multisets rather than sets.

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_{v}^{(k)}$ denote the embedding vector of node $v\in V$ in the $k$th iteration. The vectors $h_{v}^{(0)}$ represent the initial node embeddings. In every iteration $k\geq 1$, each node $v$ sends its current embedding $h_{v}^{(k-1)}$ to all its neighbors, and then computes its new embedding $h_{v}^{k}$ by the equation

where $\phi^{(k)}$ is implemented by a neural network with ReLU activations (note that $\phi^{(k)}$ may differ across different iterations $k$). The function $f$ is called the ‘aggregate’ function, and $\phi^{(k)}$ is called the ‘combine’ function. The embeddings $h_{v}^{(K)}$ at the final iteration $K$ can be used for node classification. For graph classification, they can be aggregated into a graph embedding $h_{G}$ with a ‘readout’ function,

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].

A WL iteration gets existing labels $\{h_{v}:v\in V\}$ for all nodes, and outputs new labels $\{h_{v}^{\prime}:v\in V\}$ given by

for an aggregate function $f$ and neural network $\phi$ with random weights. We say the iteration is successful if for all $v,u\in V$, the following holds:

If $\{h_{w}:w\in\mathcal{N}(v)\}=\{h_{w}:w\in\mathcal{N}(u)\}$ then $h_{v}^{\prime}=h_{u}^{\prime}$ with probability 1, and

If $\{h_{w}:w\in\mathcal{N}(v)\}\neq\{h_{w}:w\in\mathcal{N}(u)\}$ then $h_{v}^{\prime}\neq h_{u}^{\prime}$ with probability $1-p$,

where the probability is over the choices of the random weights of $\phi$.

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 final construction.

Recall that the $k$th WL iteration, for a node $v$, aggregates the labels of its neighbors from the previous iteration, $\mathcal{H}_{v}^{(k-1)}:=\{h_{w}^{(k-1)}:w\in\mathcal{N}(v)\}$, and hashes them into a new label $h_{v}^{(k)}$ for $v$. Our GNNs aggregate by summing, i.e., they sum $\mathcal{H}_{v}^{(k-1)}$ into $\mathcal{S}_{v}^{(k)}:=\sum_{w\in\mathcal{N}(v)}h_{w}^{(k-1)}$, and then hash the sum into $h_{v}^{(k)}$ using a ReLU neural network of our choice.

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_{v}^{(k-1)}$ 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 first failure mode, i.e., ensure that with high probability $\mathcal{S}_{u}^{(k)}\neq\mathcal{S}_{v}^{(k)}$ if $\mathcal{H}_{u}^{(k-1)}\neq\mathcal{H}_{v}^{(k-1)}$.

To this end, suppose for a moment that we had access to a truly random vector $a\in\{0,1\}^{F}$. Then each node $v$, instead of sending the one-hot encoding of its label $h_{v}^{(k-1)}$, we could instead send the dot product $\langle a,h_{v}^{(k-1)}\rangle$, which is the single bit $a_{h_{v}^{(k-1)}}$. Each node $u$ thus receives the bits $\left\{a_{h_{w}^{(k-1)}}:w\in\mathcal{N}(w)\right\}$ from its neighbors and aggregates them into the sum $\sum_{w\in\mathcal{N}(u)}\langle a,h_{w}^{(k-1)}\rangle$, which, by linearity, is equal to $\langle a,\mathcal{S}_{u}^{(k)}\rangle$ (using the notation from Construction 1). It is easy to observe that if $\mathcal{S}_{u}^{(k)}\neq\mathcal{S}_{v}^{(k)}$ then $\langle a,\mathcal{S}_{u}^{(k)}\rangle\neq\langle a,\mathcal{S}_{v}^{(k)}\rangle$ with probability at least $0.5$. Repeating this process $\log F$ independent times decreases the collision probability to the requisite $1/F$. Thus, we can define a new labeling scheme $\bar{h}_{v}^{(k)}\in\{0,1\}^{\log F}$ that concatenates $\log F$ dot products with independent random vectors $a^{1}\ldots a^{\log F}\in\{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 efficient pseudorandom analog. Note that the above approach goes through even if the probability that $\langle a,\mathcal{S}_{u}^{(k)}\rangle\neq\langle a,\mathcal{S}_{v}^{(k)}\rangle$ is slightly larger than $0.5$, say $0.5+\epsilon$ for a small constant $\epsilon$. It is well-known in complexity theory that there exist pseudo-random generators, called $\epsilon$-biased spaces, that generate vectors $a$ satisfying this property given only $O(\log F+\log(1/\epsilon))$ truly random bits.Technically, they guarantee this property only when $\mathcal{S}_{u}^{(k)}$ are binary vectors, but Lemma 3.2 shows how to extend this property to general integer vectors as well. Crucially, each bit of $a$ can be computed using a threshold circuit with size polynomial in $O(\log F+\log(1/\epsilon))$ 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 ($\epsilon$-biased) analog of $a$ yields our final construction.

First Construction: Polynomial-size GNN

Our first construction towards Definition 1.1 is exponentially larger compared to our final optimized construction of Section 3. Nevertheless, it is instructive and motivates our optimized construction.

Let $h_{u}^{(k)}$ denote the label of a vertex $u$ in the $k$th iteration. For our first construction, we will always maintain the invariant that $h_{u}^{(k)}$ will be a one-hot encoded vector in $\{0,1\}^{F}$ for all $u\in V$ and all iterations $k$. $F>2n$ will be a prime which also satisfies $F=O(\operatorname{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, $\phi^{(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}$.

Our construction for $\phi^{(k)}$ is the following: First recall the notation from (simplified) Equation 1.1:

$\phi^{(k)}$ then computes $\langle x,a\rangle\bmod F$.

Altogether, $\phi^{(k)}$ can be summarized as: $h_{v}^{(k)}=\text{One-Hot}\left(\left\langle\sum_{w\in\mathcal{N}(v)}h_{w}^{(k-1)},a\right\rangle\bmod F\right)$.

Note that we set the initial labels $h_{u}^{(0)}$ to be the same starting vector for all vertices (any one-hot vector). This matches the WL test which also initializes all nodes with the same initial label. Furthermore, the weights of $\phi^{(k)}$ are independent: the random vector $a$ is sampled independently for each iteration.

The following lemma proves that the above construction satisfies the requirement of Definition 1.1. Its proof is given in Appendix B.

Let $\{h_{w}^{(k-1)}:w\in\mathcal{N}(v)\}$ and $\{h_{w}^{(k-1)}:w\in\mathcal{N}(u)\}$ denote the multiset of neighborhood 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 Definition 1.1.

We now evaluate the size complexity of implementing our construction via a neural network $\phi$. 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.

Suppose $F=\operatorname{poly}(n)$. There exists an explicit construction of a network which computes modulo $F$ in the domain $\{0,\ldots,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

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=\text{ReLU}(a-b)+b$.) It can be checked that $g(0)=1$ and $g(x)=0$ for all other integers $x\neq 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.

There exists a construction of a neural network $\phi$ which performs a successful iteration according to Definition 1.1 with failure probability $p=O(1/|F|)$. $\phi$ 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)}\operatorname{polylog}(n)$ hidden units and $O(t)$ depth.

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(\operatorname{poly}(n))$ which gives us depth $O(\log n)$, $O(\operatorname{poly}(n))$ hidden units, and requires $O(\log n)$ bits of precision in addition to labels in dimension $O(\operatorname{poly}(n))$.

Second Construction: Polylogarithmic-size GNN via Pseudo-randomness

We now present a more efficient 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 significantly 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 $\varepsilon$-biased sample space. These are random vectors which are approximately uniform and they have been well-studied in the complexity-theory literature. We recall some definitions below.

A probability distribution over $\{0,1\}^{m}$ is called an $\varepsilon$-biased sample space if $\textup{bias}_{I}(X)\leq\varepsilon$ holds for all non-empty subsets $I\subseteq\{1,\ldots,m\}$.

Note that the uniform distribution has bias . We now state our construction for the neural network used in the $k$th iteration, $\phi^{(k)}$. We recall that $h_{u}^{(k)}$ denotes the label of a vertex $u$ in the $k$th iteration.

Our construction for $\phi^{(k)}$ is the following:

Let $F$ be a prime of size $\operatorname{poly}(n)$ which is at least $2n$.

For each node $v$, $\phi^{(k)}$ computes $z_{v}=\langle b,\sum_{w\in\mathcal{N}(v)}h_{w}^{(k-1)}\rangle\bmod F$ where every entry of $b$ is uniformly random in $\{0,\ldots,F-1\}$. Note $\sum_{w\in\mathcal{N}(v)}h_{w}^{(k-1)}$ is the output of the aggregation $f$.

Let $a^{t}\in\{0,1\}^{F}$ for $t=1,\ldots,C\log n$ be vectors which are independently drawn from an $\varepsilon$-biased sample space for a sufficiently small constant $\varepsilon$.

The output $h_{v}^{(k)}$ 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_{v}^{(k)}=(a^{t}(z_{v}))_{t=1}^{C\log n}$ where $a^{t}(z_{v})$ denotes the $z_{v}$-th coordinate of $a^{t}$.

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 suffices to prove the lemma below which shows our construction satisfies Definition 1.1.

Let $\{h_{w}^{(k-1)}:w\in\mathcal{N}(v)\}$ and $\{h_{w}^{(k-1)}:w\in\mathcal{N}(u)\}$ denote the multiset of neighborhood 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/\operatorname{poly}(n)$. If the multisets are the same then the labels are the same, i.e., the $k$th iteration is successful according to Definition 1.1.

We first need the following auxiliary lemma about $\varepsilon$-biased sample spaces, proven in Section C.

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$.

Let $x^{\prime}=\sum_{w\in\mathcal{N}(v)}h_{w}^{(k-1)}$ denote the input for $v$ and analogously, define $y^{\prime}=\sum_{w\in\mathcal{N}(u)}h_{w}^{(k-1)}$ to be the input for $u$. We first show that if $\{h_{w}^{(k-1)}:w\in\mathcal{N}(v)\}$ is not equal to (as multisets) $\{h_{w}^{(k-1)}:w\in\mathcal{N}(u)\}$ then $x^{\prime}\neq y^{\prime}$ with sufficiently large probability. We further consider the case that $k\geq 2$ since for $k=1$ (the first iteration), the statement follows since all node labels are initialized to be the same. Let $z_{v}^{\prime}$ be the indices computed in Step $3$ of iteration $k-1$ (which are used to construct the node labels $h_{v}^{(k-1)}$ in iteration $k-1$). Note that there is a one to one mapping between $z_{v}^{\prime}$ and $h_{v}^{(k-1)}$. Thus we can assume $\{z_{w}^{\prime}:w\in\mathcal{N}(v)\}\neq\{z_{w}^{\prime}:w\in\mathcal{N}(u)\}$ without loss of generality.

We now condition on $x^{\prime}\neq y^{\prime}$. Without loss of generality, suppose that their first coordinates, $x^{\prime}_{1}$ and $y^{\prime}_{1}$, differ. We know $x^{\prime}_{1}\neq y^{\prime}_{1}\bmod F$ since $x^{\prime}_{1}\neq y^{\prime}_{1}$, they are both non-negative and bounded by $n$, and $|x^{\prime}_{1}-y^{\prime}_{1}|\leq O(n)$ whereas $F$ is a prime at least $2n$. It follows that the probability of the event $\langle(x^{\prime}-y^{\prime}),b\rangle=0\bmod F$ is at most $1/F$. To see this, condition on all the entries of $b$ except $b_{1}$. Then $(x^{\prime}-y^{\prime})_{1}\cdot b_{1}$ must be equal to a specific value modulo $F$ for $\langle(x^{\prime}-y^{\prime}),b\rangle=0\bmod F$ to hold, as desired. We now condition on this event which equivalently means we condition on $z_{v}\neq z_{u}$ (see Step $3$ of the construction).

We now analyze the overall complexity of representing $\phi^{(k)}$ as a ReLU neural network. First we state guarantees on generating $\varepsilon$-based vectors using a ReLU network. The following corollary is proven in Appendix C.

Let $s=O(\log F+\log(1/\varepsilon))$. For every $\varepsilon$ and $F$, there exists an explicit ReLU network $C:\{0,1\}^{s}\cup[F]\rightarrow\{0,1\}^{F}$ which takes as input $s$ uniform random bits and an index $i\in F$ and outputs the $i$th coordinate of an $\varepsilon$-biased vector in $\{0,1\}^{F}$. $C$ uses $O(\log F)$ bits of precision and has $\operatorname{poly}(s)$ hidden units. More generally, given an integer parameter $t>0$, the function can be computed with $n^{O(1/t)}\operatorname{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, $\sum_{w\in\mathcal{N}(v)}h_{w}^{(k-1)}$, is returned by the aggregate function $f$.

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.

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_{v}^{(k)}$ in parallel. We recall that each coordinate of $h_{v}^{(k)}$ is indexing onto $O(\log n)$ $\varepsilon$-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 first have $O(\log n)$ edges fanning-out from the node which computes $z_{v}$. For all $t=1,\ldots,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 $\varepsilon$-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_{v}^{(k)}$.

Altogether, we have proven the following theorem.

There exists a construction of $\phi$ which performs a successful WL iteration according to Definition 1.1 with $p\leq 1/\operatorname{poly}(n)$. $\phi$ has depth $O(\log n)$, $O(\operatorname{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)}\operatorname{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.

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 Definition 1.1 must send at least $\Omega(\log n)$ bits along one edge of the graph. As message-passing GNNs are a specific class of communication protocols, this immediately implies that the message sizes must have $\Omega(\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 $\Theta(\sqrt{n})$. In order to perform a valid WL coloring, each node must essentially learn the size of its subgraph, requiring $\Omega(\log(\sqrt{n}))=\Omega(\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 $\Omega(\log n)$ bits to its neighbors in a single round. See Appendix D.1 for the full details and proof.

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 fixed graph and a distribution over inputs to the neural network $S_{u}^{(k-1)}$ for all $u\in V$ (sums of the labels from the previous round) which includes $O(F)$ potential special pairs of nodes (where $F$ is defined such that inputs $S_{u}^{(k-1)}\in[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_{u}^{(k)}=S_{v}^{(k)}$ even while their neighborhoods had distinct inputs in the $(k-1)$st round.

This immediately gives a $\Omega(\log F)$ lower bound on the number of ReLU units (and thus number of parameters) with more refined depth/width tradeoffs given in Section 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 $\Omega(\log n)$ lower bound on the number of ReLU units. See Section 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 fromSince our goal is to test whether our protocol correctly simulates WL test with small messages, we are not implementing the actual GNNs but instead we are simulating their computation. Further, for simplicity, we replaced the $\varepsilon$-biased sample space with a random string, which guarantees $\varepsilon=0$. 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 $\mathcal{P}_{k}$ and $\mathcal{P}_{k}^{\prime}$. Letting $k_{0}$ be minimal such that $\mathcal{P}_{k_{0}-1}=\mathcal{P}_{k_{0}}$ (at which point the WL labels have converged), we consider the implementation using our GNN successful if $\mathcal{P}_{k}=\mathcal{P}_{k}^{\prime}$ for all $k\leq 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}\leq 5$ (see [BK22] for a discussion of this fast convergence).

We generated Erdős-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 five such graphs and for each of these five graphs, we ran 10 independent trials of our GNN implementation with message sizes $t=1,2,\dots$. Averaging over the five graphs, we report the minimal $t$ such that at least $70\%$ of the 10 iterations successfully simulated the WL test. See Figure 2(a). The average message size needed to achieve this is approximately $3\log n$ where the logarithmic dependence on $n$ is as predicted theoretically and significantly improves on the linear message size required for prior constructions.

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ős-Rényi random graphs. See Figure 2(b).

We finally ran experiments on the real world graph Corahttps://graphsandnetworks.com/the-cora-dataset/ which is the citation network of $n=2708$ scientific 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 2(c) for a plot of the results. We see that with message length $35$, all of the $30$ trials successfully simulated the WL test.

Anders Aamand is supported 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 AI Lab, GIST- MIT Research Collaboration grant, NSF Graduate Research Fellowship under Grant No. 1745302, and MathWorks Engineering Fellowship.

References

Appendix A Omitted Proofs of Section 1

Consider some iteration $i\leq T$. Suppose we have the following guarantee on the node label inputs for the $i$th iteration (note the inputs are the output labels of the previous iteration):

Appendix B Omitted Proofs of Section 2

We now reduce the modulo $F$ case to the construction of the triangular wave function.

Appendix C Omitted Proofs for Section 3

We first need to define the circuit class TC0.

For inputs $x_{1},\ldots,x_{m}\in\{0,1\}$ the output of a threshold gate, TH, is

TC0 is the class of boolean functions computed by constant-depth $\operatorname{poly}(m)$-size circuits with threshold gates.

It is known that $\varepsilon$-biased vectors can be generated using an efficient circuit in TC0.

Let $s=O(\log F+\log(1/\varepsilon))$. For every $\varepsilon$ and $F$, there exists an explicit TC0 circuit $C:\{0,1\}^{s}\cup\{0,1\}^{\lceil log_{2}F\rceil}\rightarrow\{0,1\}$ which takes as input $s$ uniform random bits and an index $i\in[F]$ and outputs the $i$th coordinate of an $\varepsilon$-biased vector in $\{0,1\}^{F}$. $C$ uses $\operatorname{poly}(s)$ threshold gates.

Note that the guarantees of Theorem C.1 are not directly applicable since we need to use a ReLU network instead of threshold gates. Nevertheless, since the circuit $C$ guaranteed by Theorem C.1 has integer inputs in all gates, we can easily approximate each threshold gates using an appropriately scaled ReLU. This is a straightforward and known reduction but we briefly outline a procedure in Lemma C.2.

Consider the threshold gate TH: $\{0,1\}^{m}\rightarrow\{0,1\}$ which computes the threshold $\sum_{i=1}^{m}a_{i}x_{i}\geq\theta$. Assume that $a_{i},\theta$ are all integers bounded by $\operatorname{poly}(m)$. TH can be computed by a ReLU network using $O(\log m)$ bits of precision and a constant number of parameters.

It is for all integers $x\leq 0$ and $1$ for all integers $\geq 1$, i.e., it computes the threshold $``x\geq 0"$. By shifting and scaling $g$, we can now compute the threshold $``x\geq\theta"$ for any integer $\theta$. Finally, the sum $\sum_{i=1}^{m}a_{i}x_{i}$ can be computed using one additional layer. Since all parameters are integers, we only require $O(\log m)$ bits of precision to store the shifting and scaling factors. ∎

Lastly, we remark that as per the definition of a threshold gate in Definition C.1, Theorem C.1 requires the index $i\in[F]$ to be inputted as a binary string with its bits given on individual nodes. However, this presents a slight inconsistency with the statement of Theorem C.1 and its corollary, Corollary 3.3 which is used in the construction of Section 3. Specifically, Step $3$ of the construction of Section 3 outputs the actual integer $i\in[F]$ which we use as the index for our $\varepsilon$-biased vector, which does not match the format required by Theorem C.1. This inconsistency is straightforward to fix without having any impact whatsoever in the asymptotic size complexity of the neural network. We simply take the integer $i$ outputted by Step $3$ of the construction and compute the $j$th bit of $i$ for all $1\leq j\leq O(\log F)$ in parallel. The $j$th bit is exactly equal to if and only if $(i\bmod 2^{j+1})<2^{j}$ and $1$ otherwise. Note that $2^{j+1}=O(F)$ for all $j$ and we can easily compute each $\bmod 2^{j+1}$ by appealing to Theorem 2.2. This only requires $O(1)$ extra depth and an additional $O(\operatorname{poly}(\log n))$ hidden units and $O(\log n)$ bits of precision. The more general trade-off of Theorem 3.4 also readily holds.

Appendix D Lower bounds

In this appendix we provide lower bounds on the complexity of graph neural networks that are able to simulate the WL test. We present both a communication complexity lower bound and a lower bound on the number of ReLU units of the GNN. More concretely, in Section D.1, we prove that in order to maintain the invariant that with at least some constant probability, nodes with isomorphic neighborhoods get the same label while nodes with non-isomorphic neighborhoods get different labels, some message sent between nodes must be of length at least $\Omega(\log n)$. This bound matches the upper bound of Theorem 3.4. Second, in Section D.2, we consider a more specific although still fairly general lower bound model which captures the implementation of the WL test using neural networks. We suppose that the messages sent between nodes are $t$-dimensional vectors with integral entries. We moreover suppose that each node combines its received messages by summing them to get a vector in $[F]^{t}$ (here, $[F]=\{0,1,\dots,F-1\}$) and applying a collectively agreed upon neural network $\phi$ with at most $H$ ReLU units to this sum. We show that if the combination of summing neighborhoods and applying the neural network maps distinct multisets to distinct elements with at least some constant probability, then $H=\Omega(\log F)$. Moreover, parametrizing in terms of the depth and width of the neural network, we obtain a more fine-grained lower bound, demonstrating that for shallow neural networks, we need even more ReLU units. In Remark D.5, we point out that our lower bound holds even if the aggregation function $f$ is itself a neural network with a bounded number of ReLU units. As a node in an $n$-node graph could have up to $n-1$ neighbours, we need at least $F=\Omega(n)$ in order to store the sum of the messages from the neighbors of the nodes. With this assumption, the lower bound thus becomes $\Omega(\log n)$ which matches our upper bound up to $\operatorname{polylog}(n)$ factors. It remains an interesting open problem to bridge the gap between the upper and lower bound.

For both our lower bounds we assume that the nodes have access to an infinite public string of random bits. In Section D.2, this is the string which the nodes use to collectively agree on some neural network network $f$ with respect to some distribution on such networks with at most $H$ ReLU units.

We consider a forest graph $G$ composed of pieces $G_{1},G_{1}^{\prime},G_{2},G_{2}^{\prime}\dots,G_{m},G_{m}^{\prime}$, for $m=\Theta(\sqrt{n})$. Each piece $G_{k}$ consists of a “top” node $u_{k}$, which is only connected to a “middle” node $v_{k}$, which in turn is connected to $k$ “bottom” nodes $w_{k,1},\dots,w_{k,k}$, and $G_{k}^{\prime}$ is simply a duplicate of $G_{k}$ (with vertices $u_{k}^{\prime},v_{k}^{\prime}$, and $w_{k,j}^{\prime}$ for $1\leq j\leq k$). See Figure 3 for a depiction of $G$. We note that after two rounds, each $u_{k}$ (and $u_{k}^{\prime}$) should know the respective value of $k$, because the local graph of depth $2$ around $u_{k}$ is distinct for each $k\in[m]$.

Suppose there exists a public random string $r$ that every node of $G$ has access to, and each node additionally has some independent private randomness. Suppose there is a communication protocol where by the end, with probability at least $3/4$, the following hold.

For every $k\in[m]$, the top nodes $u_{k}$ and $u_{k}^{\prime}$ output the same value.

Then, there must be some $k$ such that the edge $(u_{k},v_{k})$ or the edge $(u_{k}^{\prime},v_{k}^{\prime})$ has at least $\Omega(\log n)$ total bits of communication. Hence, if there are only $O(1)$ rounds of communication, one of those rounds must have sent $\Omega(\log n)$ bits of communication across the edge.

First, we note that we may assume the communication is one-way from $v_{k}$ to $u_{k}$. This is because the node $v_{k}$ can simulate all communication from $u_{k}$, as $u_{k}$ has no information about neighbors apart from $v_{k}$. So, we just need to show the one-way communication complexity is $\Omega(\log n)$. Next, we will assume there is no public randomness - we will remove this assumption at the end. So, each $u_{k}$ (resp., $u_{k}^{\prime}$) receives at most $b$ bits of information from $v_{k}$ (resp., $v_{k}^{\prime}$). If $v_{k}$ sends a randomized message of length $b$ to $u_{k}$ and $u_{k}$ uses this message to produce some output $o_{k}$, with probability at least $\frac{3}{4}$ the outputs $o_{k}$ must all be distinct. In addition, for each $k$, the outputs of the duplicate copies of $u_{k}$ must be the same with probability at least $\frac{3}{4}$. Our goal is to show that $b=\Omega(\log n).$

To finish, we revisit the fact that we assumed there was no public randomness. Let us reintroduce the random string $r$ that every node of $G$ is given. We assume that with probability at least $3/4$, the top nodes $u_{k},u_{k}^{\prime}$ have the same output for all $k\in[m]$ and that the nodes $u_{1},\dots,u_{m}$ output pairwise distinct values. But as this event happens with probability at least $3/4$ over a random string $r$, there must exist a choice of $r$ for which it happens with probability at least $3/4$ conditioned on $r$. But then we are back to the case where there is no public randomness, as desired. ∎

D.2 Lower Bound: ReLU Units

We would like our GNN to satisfy that for any $n$-node graph $G$, and arbitrary inputs $x_{1},\dots,x_{n}$ to the nodes, it holds with probability at least $9/10$ over the randomness of $D$ that $z_{i}\neq z_{j}$ for all $i,j$ such that the multisets $\{x_{k}\mid k\in\mathcal{N}(i)\}$ and $\{x_{k}\mid k\in\mathcal{N}(j)\}$ are different. The following theorem provides a lower bound on the number of ReLU units $H$ needed for this property to hold.

Suppose that the neural networks in $\Phi_{H}$ have at most $H\leq\lg F-4$ ReLU units. Then there exists a graph $G$ on $n$ nodes and inputs $x_{1},\dots,x_{n}\in[F]^{t}$ such that if $\mathcal{N}(1),\dots,\mathcal{N}(n)$ are the neighborhoods of the nodes of $G$, then with probability at least $1-\left(\frac{3}{4}\right)^{n/6}$, there exists $i,j\in[n]$ such that $z_{i}=z_{j}$ even though the multisets $\{x_{k}\mid k\in\mathcal{N}(i)\}$ and $\{x_{k}\mid k\in\mathcal{N}(j)\}$ are different. Thus, to simulate the WL test with neural networks from $\Phi_{H}$, we need $H>\log F-3$ ReLU units.

Before proving the theorem, we first explain how to interpret it as a lower bound for the computational complexity of implementing a WL iteration as in Definition 1.1 as a neural network. As an initial observation, note that in any iteration $k$, if for two nodes $u$ and $v$, the sums $S_{u}^{(k-1)}$ and $S_{v}^{(k-1)}$ are distinct (recall that $\mathcal{S}^{(k-1)}_{u}=\sum_{u^{\prime}\in\mathcal{N}(u)}h^{(k-1)}_{u}$), then for the WL iteration to be successful, we must also have that $h^{(k)}_{u}\neq h^{(k)}_{v}$. This is because, $S_{u}^{(k-1)}\neq S_{v}^{(k-1)}$ implies that the multisets $\{h_{u^{\prime}}^{(k-1)}\mid u^{\prime}\in\mathcal{N}(u)\}$ and $\{h_{v^{\prime}}^{(k-1)}\mid v^{\prime}\in\mathcal{N}(v)\}$ are also distinct. But then Definition 1.1 yields that we need $h_{u}^{(k)}\neq h_{v}^{(k)}$ (at least with some probability $1-p$). Now, consider two nodes $i=u$ and $j=v$ such that in some iteration of the WL test, the multisets of sums $\{\mathcal{S}_{u^{\prime}}^{(k-1)}\mid u^{\prime}\in\mathcal{N}(u)\}$ and $\{\mathcal{S}_{v^{\prime}}^{(k-1)}\mid v^{\prime}\in\mathcal{N}(v)\}$ are different. This corresponds to the multisets $\{x_{k}\mid k\in\mathcal{N}(i)\}$ and $\{x_{k}\mid k\in\mathcal{N}(j)\}$ being different. We would like to argue that for the WL iteration to be successful according to Definition 1.1, for each such pair of nodes $u,v$, we must have that also $\mathcal{S}_{u}^{(k)}\neq\mathcal{S}_{v}^{(k)}$ with some good probability (the sums of labels in the next iteration differ). Theorem D.4 tells us that the probability of this happening is very low if we use too few ReLU units. Now why do we require that $\mathcal{S}_{u}^{(k)}\neq\mathcal{S}_{v}^{(k)}$ for such a pair of nodes $u,v$?

Since the multisets $\{\mathcal{S}_{u^{\prime}}^{(k-1)}\mid u^{\prime}\in\mathcal{N}(u)\}$ and $\{\mathcal{S}_{v^{\prime}}^{(k-1)}\mid v^{\prime}\in\mathcal{N}(v)\}$ are different, by the initial observation, the multisets $\{h_{u^{\prime}}^{(k)}\mid u^{\prime}\in\mathcal{N}(u)\}$ and $\{h_{v^{\prime}}^{(k)}\mid v^{\prime}\in\mathcal{N}(v)\}$ must also be distinct for the WL test to be successful. But since the multisets of labels $\{h_{u^{\prime}}^{(k)}\mid u^{\prime}\in\mathcal{N}(u)\}$ and $\{h_{u^{\prime}}^{(k)}\mid v^{\prime}\in\mathcal{N}(v)\}$ are distinct it follows by another application of Definition 1.1, that we must also have that $h_{u}^{(k+1)}\neq h_{v}^{(k+1)}$. However, the only way this can happen is if $\mathcal{S}_{u}^{(k)}\neq\mathcal{S}_{v}^{(k)}$ as otherwise these two sums will be mapped to the same label by $\phi^{(k+1)}$.

We remark that this lower bound applies to an isolated WL iteration rather than a full sequence of iterations. In particular, the inputs $x_{1},\dots,x_{n}\in[F]^{t}$ (corresponding to the sums $\{\mathcal{S}_{u}^{(k-1)}\mid u\in G\}$) are adversarially chosen while in reality these inputs are not arbitrary but are the result of a prior WL iteration. Our construction in Section 3 indeed works against such adversarially chosen sums in the sense that different multisets of sums are mapped (via applying $\phi^{(k)}$ and summing the outputs for each multisets) to different sums with high probability, and as such our lower bound is exactly a lower bound for this harder problem. However, in general the sums $S_{v}^{(k-1)}$ are not adversarially chosen, and it would very be interesting to find a lower bound that does not require this assumption but works all the way from a graph and its initial labels.

in spite of the multisets, $\{x_{6j},x_{6j+1},x_{6j+2}\}$ and $\{x_{6j+3},x_{6j+4},x_{6j+5}\}$ being different. Now, $\phi$ is a neural network, so this identity does not need to hold. The idea is however, that if $\phi$ has only few ReLU units, then we obtain a good upper bound on the number of linear regions by Theorem D.3 and since $a$ and $b$ are random, the set $F_{a,b}$ is likely to be fully contained in one of these regions. And since $\phi$ restricted to this region in linear, (D.1) holds in this case.

In other words, if for a fixed $i\leq n_{0}$, $(a,b)$ is chosen such that $L_{a,b}$ is fully contained in one of the convex regions, then (D.1) is satisfied. It follows that for any given $i<n_{0}$,

Since the event $(z_{6i+1}=z_{6i+4})_{i\leq n_{0}}$ are independent (as we choose independent $(a,b)$ for each pair of paths $(P_{i},P_{i}^{\prime})$), it follows that

If in particular, $H\leq\lg F-4$, we obtain that

As the multisets $\{x_{k}\mid k\in\mathcal{N}(6i+1)\}$ and $\{x_{k}\mid k\in\mathcal{N}(6i+4)\}$ are different, this completes the proof. ∎

In analogue with our construction in Section 2 and Section 3, we assumed in the above proof that the aggregate function $f$ is the summation function. As such, $f$ is just another neural network but without a single ReLU unit. We can therefore think of the combined computation performed by $f$ and $\phi$ (illustrated in Figure 5) as the result of applying a single neural network. It follows from this observation (and the proof of Theorem D.4) that in the more general setting where $f$ is a function in $\Phi_{H_{1}}$ and where the neural network $\phi\in\Phi_{H_{2}}$, then we must have that $H_{1}+H_{2}>\log F-4$ in order to successfully simulate the WL test.

The proof of Theorem D.4 used that the number of convex linear regions of any neural network with at most $H$ ReLU’s is at most $2^{H}$. However, in many cases one can obtain better upper bounds on the number of such regions, and this directly translates to a better lower bound than the one given in Theorem D.4. Indeed, if the family of neural networks $\Phi$ satisfies that the domain of any $\phi\in\Phi$ can be partitioned into at most $K$ convex regions such that $\phi$ restricted to each of these regions is linear, then the lower bound in (D.2) instead becomes

In particular, when using independence of the events $(z_{6i+1}=z_{6i+4})_{i<n_{0}}$, we just need $K\leq\frac{F}{24}$, say, to get that an error occur with probability at least $1-2^{-\Omega(n)}$. Plugging in the bound $K\leq 2^{H}$ of Theorem D.3 gave the desired bound of Theorem D.4 which led to the $\Omega(\log F)$ lower bound. If we instead use the more fine grained theorem below, we obtain better bounds for shallow neural networks with low input dimension as stated in Corollary D.7.

Any ReLU neural network with input dimension $t$, width $w$, and depth $d$ has at most $O(w^{d\cdot t})$ linear regions.

Let $\Phi_{t,d,w}$ consist of all ReLU neural networks with input dimension $t$, depth $d$, and width $w$. Suppose that we are in the setting of Theorem D.4, except that the neural network is picked from $\Phi_{t,d,w}$. Then the conclusion of the theorem holds as long as $w^{d\cdot t}\leq cF$ for a small enough constant $c$. In particular, to simulate the WL test with neural networks from $\Phi_{t,d,w}$, we need $dw=\Omega(dF^{\frac{1}{dt}})$ ReLU units.

As an example, for shallow neural networks with low input dimension, say with $d,t=O(1)$, this lower bound becomes $dw=F^{\Omega(1)}$, i.e. polynomial rather than logarithmic in the size of the underlying field.

For certain architectures of the neural networks one can obtain even stronger bounds on the number of linear regions (see e.g., Proposition 3 in [Mon17] and Theorem 1 in [STR18]). These bounds are parametrized in the number of ReLU units in each of the $d$ layers of the neural networks. One can therefore obtain even more fine grained lower bounds on the number of ReLU units if one makes more assumptions on the family of neural networks but the bounds are more opaque and we refrain from stating them here.

D.2.2 Description complexity

In this subsection, we consider more general function classes $\Phi$ that do not necessarily have to consist of neural networks. We prove that for any aggregation function $f$, if for any two distinct multisets of labels each of size at most $n$, there exists a function $\phi\in\Phi$ such that the multisets are mapped to different labels by $\phi\circ f$, then $|\Phi|=\Omega\left(\frac{n}{\log n}\right)$. It follows that the description complexity of $\Phi$ must be $\Omega(\log n)$. Our bound is combinatorial, and does not employ the linear structure of $F^{t}$. Hence we may just put $t=1$ and think of $F$ as a set rather than a vector space.

Let $F$ and $n$ be natural numbers and let $N=\sum_{i=0}^{n}\binom{i+F-1}{i}$ be the number of multisets of $[F]$ of size at most $n$. Suppose that $f$ is any aggregation function mapping multisets of $[F]$ to some range $\mathcal{R}$. Let $\Phi$ be any set of functions from $\mathcal{R}$ to $F$ and assume that $|\Phi|<\frac{\log N}{\log F}$. Then there exists distinct multisets $A$ and $B$ each with at most $n$ elements from $[F]$ such that $\phi(f(A))=\phi(f(B))$ for all $\phi\in\Phi$.