SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Noam Razin, Tom Verbin, Nadav Cohen
{noamrazin,tomverbin}@mail.tau.ac.il,cohennadav@cs.tau.ac.il
Graph neural networks (GNNs) are widely used for modeling complex interactions between entities represented as vertices of a graph. Despite recent efforts to theoretically analyze the expressive power of GNNs, a formal characterization of their ability to model interactions is lacking. The current paper aims to address this gap. Formalizing strength of interactions through an established measure known as separation rank, we quantify the ability of certain GNNs to model interaction between a given subset of vertices and its complement, i.e. between the sides of a given partition of input vertices. Our results reveal that the ability to model interaction is primarily determined by the partition's walk index — a graph-theoretical characteristic defined by the number of walks originating from the boundary of the partition. Experiments with common GNN architectures corroborate this finding. As a practical application of our theory, we design an edge sparsification algorithm named Walk Index Sparsification (WIS), which preserves the ability of a GNN to model interactions when input edges are removed. WIS is simple, computationally efficient, and in our experiments has markedly outperformed alternative methods in terms of induced prediction accuracy.<sup>1</sup> More broadly, it showcases the potential of improving GNNs by theoretically analyzing the interactions they can model.
Graph neural networks (GNNs) are a family of deep learning architectures, designed to model complex interactions between entities represented as vertices of a graph. In recent years, GNNs have been successfully applied across a wide range of domains, including social networks, biochemistry, and recommender systems (see, e.g., [36, 59, 45, 49, 96, 104, 101, 18]). Consequently, significant interest in developing a mathematical theory behind GNNs has arisen.
One of the fundamental questions a theory of GNNs should address is expressivity, which concerns the class of functions a given architecture can realize. Existing studies of expressivity largely fall into three categories. First, and most prominent, are characterizations of ability to distinguish nonisomorphic graphs [103, 74, 72, 70, 6, 15, 10, 17, 43, 42, 80], as measured by equivalence to classical Weisfeiler-Leman graph isomorphism tests [99]. Second, are proofs for universal approximation of continuous permutation invariant or equivariant functions, possibly up to limitations in distinguishing some classes of graphs [73, 55, 25, 69, 3, 42]. Last, are works examining specific properties of GNNs such as frequency response [77, 5] or computability of certain graph attributes, e.g. moments, shortest paths, and substructure multiplicity [35, 9, 26, 39, 69, 23, 17, 105].
A major drawback of many existing approaches — in particular proofs of equivalence to Weisfeiler-Leman tests and those of universality — is that they operate in asymptotic regimes of unbounded
<sup>1</sup>An implementation of WIS is available at https://github.com/noamrazin/gnn_interactions.
Figure 1: Illustration of our main theoretical contribution: quantifying the ability of GNNs to model interactions between vertices of an input graph. Consider a partition of vertices  $(\mathcal{I},\mathcal{I}^c)$ , illustrated on the left, and a depth  $L$  GNN with product aggregation (Section 3). For graph prediction, as illustrated on the right, the strength of interaction the GNN can model between  $\mathcal{I}$  and  $\mathcal{I}^c$ , measured via separation rank (Section 2.2), is primarily determined by the partition's  $(L - 1)$ -walk index — the number of length  $L - 1$  walks emanating from  $\mathcal{C}_{\mathcal{I}}$  which is the set of vertices with an edge crossing the partition. The same holds for vertex prediction, except that there walk index is defined while only considering walks ending at the target vertex.
network width or depth. Moreover, to the best of our knowledge, none of the existing approaches formally characterize the strength of interactions GNNs can model between vertices, and how that depends on the structure of the input graph and the architecture of the neural network.
The current paper addresses the foregoing gaps. Namely, it theoretically quantifies the ability of fixed-size GNNs to model interactions between vertices, delineating the impact of the input graph structure and the neural network architecture (width and depth). Strength of modeled interactions is formalized via separation rank [12] — a commonly used measure for the interaction a function models between a subset of input variables and its complement (the rest of the input variables). Given a function and a partition of its input variables, the higher the separation rank, the more interaction the function models between the sides of the partition. Separation rank is prevalent in quantum mechanics, where it can be viewed as a measure of entanglement [62]. It was previously used for analyzing variants of convolutional, recurrent, and self-attention neural networks, yielding both theoretical insights and practical tools [30, 33, 61, 62, 64, 100, 65, 85]. We employ it for studying GNNs.
Key to our theory is a widely studied correspondence between neural networks with polynomial nonlinearity and tensor networks $^2$  [32, 29, 30, 34, 90, 61, 62, 7, 56, 57, 63, 64, 83, 100, 84, 85, 65]. We extend this correspondence, and use it to analyze message-passing GNNs with product aggregation. We treat both graph prediction, where a single output is produced for an entire input graph, and vertex prediction, in which the network produces an output for every vertex. For graph prediction, we prove that the separation rank of a depth  $L$  GNN with respect to a partition of vertices is primarily determined by the partition's  $(L - 1)$ -walk index — a graph-theoretical characteristic defined to be the number of length  $L - 1$  walks originating from vertices with an edge crossing the partition. The same holds for vertex prediction, except that there walk index is defined while only considering walks ending at the target vertex. Our result, illustrated in Figure 1, implies that for a given input graph, the ability of GNNs to model interaction between a subset of vertices  $\mathcal{I}$  and its complement  $\mathcal{I}^c$ , predominantly depends on the number of walks originating from the boundary between  $\mathcal{I}$  and  $\mathcal{I}^c$ . We corroborate this proposition through experiments with standard GNN architectures, such as Graph Convolutional Network (GCN) [59] and Graph Isomorphism Network (GIN) [103].
Our theory formalizes conventional wisdom by which GNNs can model stronger interaction between regions of the input graph that are more interconnected. More importantly, we show that it facilitates an edge sparsification algorithm that preserves the expressive power of GNNs (in terms of ability to model interactions). Edge sparsification concerns removal of edges from a graph for reducing computational and/or memory costs, while attempting to maintain selected properties of the graph (cf. [11, 93, 48, 20, 86, 98, 67, 24]). In the context of GNNs, our interest lies in maintaining prediction accuracy as the number of edges removed from the input graph increases. We propose an algorithm for removing edges, guided by our separation rank characterization. The algorithm, named Walk Index Sparsification (WIS), is demonstrated to yield high predictive performance for GNNs (e.g. GCN and GIN) over standard benchmarks of various scales, even when removing a significant portion of
$^{2}$ Tensor networks form a graphical language for expressing contractions of tensors — multi-dimensional arrays. They are widely used for constructing compact representations of quantum states in areas of physics (see, e.g., [97, 79]).
edges. WIS is simple, computationally efficient, and in our experiments has markedly outperformed alternative methods in terms of induced prediction accuracies across edge sparsity levels. More broadly, WIS showcases the potential of improving GNNs by theoretically analyzing the interactions they can model, and we believe its further empirical investigation is a promising direction for future research.
The remainder of the paper is organized as follows. Section 2 introduces notation and the concept of separation rank. Section 3 presents the theoretically analyzed GNN architecture. Section 4 theoretically quantifies (via separation rank) its ability to model interactions between vertices of an input graph. Section 5 proposes and evaluates WIS — an edge sparsification algorithm for arbitrary GNNs, born from our theory. Lastly, Section 6 concludes. Related work is discussed throughout, and for the reader's convenience, is recapitulated in Appendix B.
For  $N\in \mathbb{N}$ , let  $[N]\coloneqq \{1,\ldots ,N\}$ . We consider an undirected input graph  $\mathcal{G} = (\mathcal{V},\mathcal{E})$  with vertices  $\mathcal{V} = [|\mathcal{V}|]$  and edges  $\mathcal{E}\subseteq \{\{i,j\} :i,j\in \mathcal{V}\}$ . Vertices are equipped with features  $\mathbf{X}\coloneqq (\mathbf{x}^{(1)},\dots ,\mathbf{x}^{(|\mathcal{V}|)})\in \mathbb{R}^{D_x\times |\mathcal{V}|}$  one  $D_{x}$ -dimensional feature vector per vertex  $(D_{x}\in \mathbb{N})$ . For  $i\in \mathcal{V}$ , we use  $\mathcal{N}(i)\coloneqq \{j\in \mathcal{V}:\{i,j\} \in \mathcal{E}\}$  to denote its set of neighbors, and, as customary in the context of GNNs, assume the existence of all self-loops, i.e.  $i\in \mathcal{N}(i)$  for all  $i\in \mathcal{V}$  (cf. [59, 50]). Furthermore, for  $\mathcal{I}\subseteq \mathcal{V}$  we let  $\mathcal{N}(\mathcal{I})\coloneqq \cup_{i\in \mathcal{I}}\mathcal{N}(i)$  be the neighbors of vertices in  $\mathcal{I}$ , and  $\mathcal{I}^c\coloneqq \mathcal{V}\setminus \mathcal{I}$  be the complement of  $\mathcal{I}$ . We use  $\mathcal{C}_{\mathcal{I}}$  to denote the boundary of the partition  $(\mathcal{I},\mathcal{I}^c)$ , i.e. the set of vertices with an edge crossing the partition, defined by  $\mathcal{C}_{\mathcal{I}}\coloneqq \{i\in \mathcal{I}:\mathcal{N}(i)\cap \mathcal{I}^c\neq \emptyset \} \cup \{j\in \mathcal{I}^c:\mathcal{N}(j)\cap \mathcal{I}\neq \emptyset \}$ . Lastly, we denote the number of length  $l\in \mathbb{N}_{\geq 0}$  walks from any vertex in  $\mathcal{I}\subseteq \mathcal{V}$  to any vertex in  $\mathcal{I}\subseteq \mathcal{V}$  by  $\rho_l(\mathcal{I},\mathcal{J})$ . In particular,  $\rho_l(\mathcal{I},\mathcal{J}) = \sum_{i\in \mathcal{I},j\in \mathcal{J}}\rho_l(\{i\},\{j\})$
Note that we focus on undirected graphs for simplicity of presentation. As discussed in Section 4, our results are extended to directed graphs in Appendix D.
2.2 Separation Rank: A Measure of Modeled Interaction
A prominent measure quantifying the interaction a multivariate function models between a subset of input variables and its complement (i.e. all other variables) is known as separation rank. The separation rank was introduced in [12], and has since been employed for various applications [51, 47, 13]. It is also a common measure of entanglement, a profound concept in quantum physics quantifying interaction between particles [62]. In the context of deep learning, it enabled analyses of expressivity and generalization in certain convolutional, recurrent, and self-attention neural networks, resulting in theoretical insights and practical methods (guidelines for neural architecture design, pretraining schemes, and regularizers — see [30, 33, 61, 62, 64, 100, 65, 85]).
Given a multivariate function  $f: (\mathbb{R}^{D_x})^N \to \mathbb{R}$ , its separation rank with respect to a subset of input variables  $\mathcal{I} \subseteq [N]$  is the minimal number of summands required to express it, where each summand is a product of two functions — one that operates over variables indexed by  $\mathcal{I}$ , and another that operates over the remaining variables. Formally:
Definition 1. The separation rank of  $f: (\mathbb{R}^{D_x})^N \to \mathbb{R}$  with respect to  $\mathcal{I} \subseteq [N]$  is:
$$
\begin{array}{l} \operatorname {s e p} (f; \mathcal {I}) := \min  \left\{R \in \mathbb {N} _ {\geq 0}: \exists g ^ {(1)}, \dots , g ^ {(R)}: (\mathbb {R} ^ {D _ {x}}) ^ {| \mathcal {I} |} \to \mathbb {R}, \bar {g} ^ {(1)}, \dots , \bar {g} ^ {(R)}: (\mathbb {R} ^ {D _ {x}}) ^ {| \mathcal {I} ^ {c} |} \to \mathbb {R} \right. \\ \left. \text {s . t .} f (\mathbf {X}) = \sum_ {r = 1} ^ {R} g ^ {(r)} \left(\mathbf {X} _ {\mathcal {I}}\right) \cdot \bar {g} ^ {(r)} \left(\mathbf {X} _ {\mathcal {I} ^ {\mathrm {c}}}\right) \right\}, \tag {1} \\ \end{array}
$$
where  $\mathbf{X} := (\mathbf{x}^{(1)}, \ldots, \mathbf{x}^{(N)}), \mathbf{X}_{\mathcal{I}} := (\mathbf{x}^{(i)})_{i \in \mathcal{I}}$ , and  $\mathbf{X}_{\mathcal{I}^c} := (\mathbf{x}^{(j)})_{j \in \mathcal{I}^c}$ . By convention, if  $f$  is identically zero then  $\operatorname{sep}(f; \mathcal{I}) = 0$ , and if the set on the right hand side of Equation (1) is empty then  $\operatorname{sep}(f; \mathcal{I}) = \infty$ .
3Due to the existence of self-loops,  $\mathcal{C}_{\mathcal{I}}$  is exactly the shared neighbors of  $\mathcal{I}$  and  $\mathcal{I}^c$  , i.e.  $\mathcal{C}_{\mathcal{I}} = \mathcal{N}(\mathcal{I})\cap \mathcal{N}(\mathcal{I}^{c})$  4For  $l\in \mathbb{N}_{\geq 0}$  , a sequence of vertices  $i_0,\dots ,i_l\in \mathcal{V}$  is a length  $l$  walk if  $\{i_{l^{\prime} - 1},i_{l^{\prime}}\} \in \mathcal{E}$  for all  $l^{\prime}\in [l]$
Interpretation If  $\mathrm{sep}(f;\mathcal{I}) = 1$ , the function is separable, meaning it does not model any interaction between  $\mathbf{X}_{\mathcal{I}}$  and  $\mathbf{X}_{\mathcal{I}^c}$ , i.e. between the sides of the partition  $(\mathcal{I},\mathcal{I}^c)$ . Specifically, it can be represented as  $f(\mathbf{X}) = g(\mathbf{X}_{\mathcal{I}})\cdot \bar{g} (\mathbf{X}_{\mathcal{I}^c})$  for some functions  $g$  and  $\bar{g}$ . In a statistical setting, where  $f$  is a probability density function, this would mean that  $\mathbf{X}_{\mathcal{I}}$  and  $\mathbf{X}_{\mathcal{I}^c}$  are statistically independent. The higher  $\mathrm{sep}(f;\mathcal{I})$  is, the farther  $f$  is from separability, implying stronger modeling of interaction between  $\mathbf{X}_{\mathcal{I}}$  and  $\mathbf{X}_{\mathcal{I}^c}$ .
3 Graph Neural Networks
Modern GNNs predominantly follow the message-passing paradigm [45, 50], whereby each vertex is associated with a hidden embedding that is updated according to its neighbors. The initial embedding of  $i \in \mathcal{V}$  is taken to be its input features:  $\mathbf{h}^{(0,i)} := \mathbf{x}^{(i)} \in \mathbb{R}^{D_x}$ . Then, in a depth  $L$  message-passing GNN, a common update scheme for the hidden embedding of  $i \in \mathcal{V}$  at layer  $l \in [L]$  is:
$$
\mathbf {h} ^ {(l, i)} = \text {A G G R E G A T E} \left(\left\{\left\{\mathbf {W} ^ {(l)} \mathbf {h} ^ {(l - 1, j)}: j \in \mathcal {N} (i) \right\} \right\}\right), \tag {2}
$$
where  $\{\cdot\}$  denotes a multiset,  $\mathbf{W}^{(1)} \in \mathbb{R}^{D_h \times D_x}$ ,  $\mathbf{W}^{(2)} \in \mathbb{R}^{D_h \times D_h}$ , ...,  $\mathbf{W}^{(L)} \in \mathbb{R}^{D_h \times D_h}$  are learnable weight matrices, with  $D_h \in \mathbb{N}$  being the network's width (i.e. hidden dimension), and AGGREGATE is a function combining multiple input vectors into a single vector. A notable special case is GCN [59], in which AGGREGATE performs a weighted average followed by a non-linear activation function (e.g. ReLU). Other aggregation operators are also viable, e.g. element-wise sum, max, or product (cf. [49, 53]). We note that distinguishing self-loops from other edges, and more generally, treating multiple edge types, is possible through the use of different weight matrices for different edge types [49, 88]. For conciseness, we hereinafter focus on the case of a single edge type, and treat multiple edge types in Appendix D.
After  $L$  layers, the GNN generates hidden embeddings  $\mathbf{h}^{(L,1)},\ldots ,\mathbf{h}^{(L,|\mathcal{V}|)}\in \mathbb{R}^{D_h}$ . For graph prediction, where a single output is produced for the whole graph, the hidden embeddings are usually combined into a single vector through the AGGREGATE function. A final linear layer with weights  $\mathbf{W}^{(o)}\in \mathbb{R}^{1\times D_h}$  is then applied to the resulting vector.6 Overall, the function realized by a depth  $L$  graph prediction GNN receives an input graph  $\mathcal{G}$  with vertex features  $\mathbf{X}\coloneqq (\mathbf{x}^{(1)},\dots,\mathbf{x}^{(|\mathcal{V}|)})\in$ $\mathbb{R}^{D_x\times |\mathcal{V}|}$ , and returns:
$$
\text {(g r a p h p r e d i c t i o n)} \quad f ^ {(\theta , \mathcal {G})} (\mathbf {X}) := \mathbf {W} ^ {(o)} \operatorname {A G G R E G A T E} \left(\left\{\left\{\mathbf {h} ^ {(L, i)}: i \in \mathcal {V} \right\} \right\}\right), \tag {3}
$$
with  $\theta := (\mathbf{W}^{(1)}, \dots, \mathbf{W}^{(L)}, \mathbf{W}^{(o)})$  denoting the network's learnable weights. For vertex prediction tasks, where the network produces an output for every  $t \in \mathcal{V}$ , the final linear layer is applied to each  $\mathbf{h}^{(L,t)}$  separately. That is, for a target vertex  $t \in \mathcal{V}$ , the function realized by a depth  $L$  vertex prediction GNN is given by:
$$
(\text {v e r t e x p r e d i c t i o n}) \quad f ^ {(\theta , \mathcal {G}, t)} (\mathbf {X}) := \mathbf {W} ^ {(o)} \mathbf {h} ^ {(L, t)}. \tag {4}
$$
Our aim is to investigate the ability of GNNs to model interactions between vertices. Prior studies of interactions modeled by different deep learning architectures have focused on neural networks with polynomial non-linearity, building on their representation as tensor networks [32, 30, 34, 90, 61, 62, 7, 56, 63, 64, 83, 100, 84, 85, 65]. Although neural networks with polynomial non-linearity are less common in practice, they have demonstrated competitive performance [28, 31, 91, 94, 27, 37, 53], and hold promise due to their compatibility with quantum computation [46, 14] and fully homomorphic encryption [44]. More importantly, their analyses brought forth numerous insights that were demonstrated empirically and led to development of practical tools for widespread deep learning models (with non-linearities such as ReLU).
Following the above, in our theoretical analysis (Section 4) we consider GNNs with (element-wise) product aggregation, which are polynomial functions of their inputs. Namely, the AGGREGATE operator from Equations (2) and (3) is taken to be:
$$
\operatorname {A G G R E G A T E} (\mathcal {X}) := \odot_ {\mathbf {x} \in \mathcal {X}} \mathbf {x}, \tag {5}
$$
In GCN, AGGREGATE also has access to the degrees of vertices, which are used for computing the averaging weights. We omit the dependence on vertex degrees in our notation for conciseness.
<sup>6</sup>We treat the case of output dimension one merely for the sake of presentation. Extension of our theory (delivered in Section 4) to arbitrary output dimension is straightforward — the results hold as stated for each of the functions computing an output entry.
where  $\odot$  stands for the Hadamard product and  $\mathcal{X}$  is a multiset of vectors. The resulting architecture can be viewed as a variant of the GNN proposed in [53], where it was shown to achieve competitive performance in practice. Central to our proofs are tensor network representations of GNNs with product aggregation (formally established in Appendix E), analogous to those used for analyzing other types of neural networks. We empirically demonstrate our theoretical findings on popular GNNs (Section 4.2), such as GCN and GIN with ReLU non-linearity, and use them to derive a practical edge sparsification algorithm (Section 5).
We note that some of the aforementioned analyses of neural networks with polynomial non-linearity were extended to account for additional non-linearities, including ReLU, through constructs known as generalized tensor networks [29]. We thus believe our theory may be similarly extended, and regard this as an interesting direction for future work.
4 Theoretical Analysis: The Effect of Input Graph Structure and Neural Network Architecture on Modeled Interactions
In this section, we employ separation rank (Definition 1) to theoretically quantify how the input graph structure and network architecture (width and depth) affect the ability of a GNN with product aggregation to model interactions between input vertices. We overview the main results and their implications in Section 4.1, while deferring the formal analysis to Appendix A due to lack of space. Section 4.2 provides experiments demonstrating our theory's implications on common GNNs, such as GCN and GIN with ReLU non-linearity.
4.1 Overview and Implications
Consider a depth  $L$  GNN with width  $D_{h}$  and product aggregation (Section 3). Given a graph  $\mathcal{G}$ , any assignment to the weights of the network  $\theta$  induces a multivariate function —  $f^{(\theta, \mathcal{G})}$  for graph prediction (Equation (3)) and  $f^{(\theta, \mathcal{G}, t)}$  for prediction over a given vertex  $t \in \mathcal{V}$  (Equation (4))— whose variables correspond to feature vectors of input vertices. The separation rank of this function with respect to  $\mathcal{I} \subseteq \mathcal{V}$  thus measures the interaction modeled across the partition  $(\mathcal{I}, \mathcal{I}^c)$ , i.e. between the vertices in  $\mathcal{I}$  and those in  $\mathcal{I}^c$ . The higher the separation rank is, the stronger the modeled interaction.
Key to our analysis are the following notions of walk index, defined by the number of walks emanating from the boundary of the partition  $(\mathcal{I},\mathcal{I}^c)$ , i.e. from vertices with an edge crossing the partition induced by  $\mathcal{I}$  (see Figure 1 for an illustration).
Definition 2. Let  $\mathcal{I} \subseteq \mathcal{V}$ . Denote by  $\mathcal{C}_{\mathcal{I}}$  the set of vertices with an edge crossing the partition  $(\mathcal{I}, \mathcal{I}^c)$ , i.e.  $\mathcal{C}_{\mathcal{I}} := \{i \in \mathcal{I} : \mathcal{N}(i) \cap \mathcal{I}^c \neq \emptyset\} \cup \{j \in \mathcal{I}^c : \mathcal{N}(j) \cap \mathcal{I} \neq \emptyset\}$ , and recall that  $\rho_l(\mathcal{C}_{\mathcal{I}}, \mathcal{J})$  denotes the number of length  $l \in \mathbb{N}_{\geq 0}$  walks from any vertex in  $\mathcal{C}_{\mathcal{I}}$  to any vertex in  $\mathcal{J} \subseteq \mathcal{V}$ . For  $L \in \mathbb{N}$ :
- (graph prediction) we define the  $(L - 1)$ -walk index of  $\mathcal{I}$ , denoted  $\mathrm{WI}_{L - 1}(\mathcal{I})$ , to be the number of length  $L - 1$  walks originating from  $\mathcal{C}_{\mathcal{I}}$ , i.e.  $\mathrm{WI}_{L - 1}(\mathcal{I}) \coloneqq \rho_{L - 1}(\mathcal{C}_{\mathcal{I}}, \mathcal{V})$ ; and
- (vertex prediction) for  $t \in \mathcal{V}$  we define the  $(L - 1, t)$ -walk index of  $\mathcal{I}$ , denoted  $\mathrm{WI}_{L - 1, t}(\mathcal{I})$ , to be the number of length  $L - 1$  walks from  $\mathcal{C}_{\mathcal{I}}$  that end at  $t$ , i.e.  $\mathrm{WI}_{L - 1, t}(\mathcal{I}) \coloneqq \rho_{L - 1}(\mathcal{C}_{\mathcal{I}}, \{t\})$ .
As our main theoretical contribution, we prove:
Theorem 1 (informally stated). For all weight assignments  $\theta$  and  $t\in \mathcal{V}$
$$
\left(\text {g r a p h p r e d i c t i o n}\right) \quad \log \big (\operatorname {s e p} \big (f ^ {(\theta , \mathcal {G})}; \mathcal {I} \big) \big) = \mathcal {O} \big (\log (D _ {h}) \cdot \mathrm {W I} _ {L - 1} (\mathcal {I}) \big),
$$
$$
\left(\text {v e r t e x p r e d i c t i o n}\right) \quad \log \left(\operatorname {s e p} \left(f ^ {\left(\theta , \mathcal {G}, t\right)}; \mathcal {I}\right)\right) = \mathcal {O} \left(\log \left(D _ {h}\right) \cdot \mathrm {W I} _ {L - 1, t} (\mathcal {I})\right).
$$
Moreover, nearly matching lower bounds hold for almost all weight assignments.<sup>7</sup>