SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Aseem Baranwal
University of Waterloo, Waterloo, Canada
aseem.paranwal@uwaterloo.ca
Kimon Fountoulakis
kimon.fountoulakis@uwaterloo.ca
Aukosh Jagannath
Department of Statistics and Actuarial Science, Department of Applied Mathematics,
a.jagannath@uwaterloo.ca
We study the node classification problem on feature-decorated graphs in the sparse setting, i.e., when the expected degree of a node is  $O(1)$  in the number of nodes, in the fixed-dimensional asymptotic regime, i.e., the dimension of the feature data is fixed while the number of nodes is large. Such graphs are typically known to be locally tree-like. We introduce a notion of Bayes optimality for node classification tasks, called asymptotic local Bayes optimality, and compute the optimal classifier according to this criterion for a fairly general statistical data model with arbitrary distributions of the node features and edge connectivity. The optimal classifier is implementable using a message-passing graph neural network architecture. We then compute the generalization error of this classifier and compare its performance against existing learning methods theoretically on a well-studied statistical model with naturally identifiable signal-to-noise ratios (SNRs) in the data. We find that the optimal message-passing architecture interpolates between a standard MLP in the regime of low graph signal and a typical convolution in the regime of high graph signal. Furthermore, we prove a corresponding non-asymptotic result.
Graph Neural Networks (GNNs) have rapidly emerged as a powerful tool for learning on graph-structured data, where along with features of the entities, there also exists a relational structure among them. They have found numerous applications to a wide range of domains such as social networks [Backstrom and Leskovec, 2011], recommendation systems [Ying et al., 2018, Hao et al., 2020], chip design [Mirhoseini et al., 2021], bioinformatics [Scarselli et al., 2009, Zhang et al., 2021], computer vision [Monti et al., 2017], quantum chemistry [Gilmer et al., 2017], statistical physics [Battaglia et al., 2016, Bapat et al., 2020], and financial forensics [Zhang et al., 2017, Weber et al., 2019]. Most of the success with these applications has been possible due to the advent of the
message-passing paradigm in GNNs, however, designing optimal GNN architectures for such a wide variety of applications still remains a challenging task.
In this work, we are interested in the node classification problem on very sparse feature-decorated graphs that are locally tree-like. We focus on the regime where the dimension of the node features is fixed and the number of nodes is large. Our motivation for considering this regime is that many major benchmark datasets for node classification appear to scale in this fashion. For example, in the popular Open Graph Benchmark collection [Hu et al., 2020], the medium and large-scale node-property prediction datasets have roughly  $10^{6}$  nodes (ogbn-products, ogbn-mag) to about  $10^{8}$  nodes (ogbn-papers100M), each with roughly  $10^{2}$  features. Graphs with such properties exist naturally in social, informational and biological networks; for motivational examples, see Stelzl et al. [2005], Adcock et al. [2013]. We present a precise definition of optimality for node classification tasks on locally tree-like graphs in this scaling regime and compute the optimal classifier according to this definition, for a multi-class statistical data model where node features can have arbitrary continuous or discrete distributions. Subsequently, we show that a message-passing GNN architecture is able to realize the optimal classifier. Furthermore, we provide a theoretical analysis, comparing the generalization error of the optimal classifier with other architectures like GCN and simple MLPs. Our results support a recent work [Veličković, 2022] in the context of classification on sparse graphs. In particular, we show that when node features are accompanied by sparse graphical side information, message-passing graph neural networks are able to realize the optimal classification scheme, and as such, there does not exist a better architecture beyond the message-passing paradigm.
Related Work. There has been a tremendous amount of work on GNN architecture design, where the most popular designs are based on a convolutional architecture, with each layer of the neural network performing a weighted convolution (averaging) operation with immediate neighbours, e.g., graph convolutional networks (GCN) [Kipf and Welling, 2017, Chen et al., 2020] or graph attention networks (GAT) [Veličković et al., 2018]. These architectures are known to have several limitations regarding their expressive power (see, for e.g., Li et al. [2018], Oono and Suzuki [2020], Balcilar et al. [2021], Xu et al. [2021], Keriven [2022]).
An interesting line of research consists of both theoretical and empirical works that attempt to address these limitations by developing an understanding of GNN architectures within the scope of message-passing [Rong et al., 2020, Liu et al., 2022, Maskey et al., 2022], as well as beyond it [Maron et al., 2019, Murphy et al., 2019, Chen et al., 2019]. For example, Xu et al. [2018] propose an architecture with a technique called skip-connections, that flexibly leverages different ranges of neighbourhoods for each node to enable structure-awareness in node representations, Chen et al. [2020] propose a modification of the vanilla GCN with an initial residual that effectively relieves the problem of oversmoothing [Oono and Suzuki, 2020], and Keriven et al. [2021] study the universality of structural GNNs in the large random graph limit. However, this area of research still lacks a clear understanding of optimality in the context of graph learning problems, making it hard to design architectures for which a well-defined notion of optimality can be theoretically justified.
Several works have studied traditional message-passing GNN architectures like GCN and GAT using the binary contextual stochastic block model, see for example, Baranwal et al. [2021], Chien et al. [2022], Fountoulakis et al. [2022a,b], Javaloy et al. [2022], Baranwal et al. [2023]. These analyses rely heavily on two assumptions: first, the graph is not too sparse, i.e., for a graph with  $n$  nodes, the expected degree of a node is  $\Omega_{n}(\log^{2}n / n)$ , and second, the node features are modelled as a Gaussian mixture. The work by Wei et al. [2022] is of particular interest to us, where the authors take a Bayesian inference perspective to investigate the functions of non-linearity in GNNs for binary node classification. They characterize the max-a-posteriori estimation of a node label given the features of itself and its immediate neighbours. A similar perspective to that of Wei et al. [2022] is discussed in Gosch et al. [2023], where the latter authors derive insights into the robustness-accuracy trade-off in GNNs for node classification. In contrast to these inspiring works, we study the highly sparse regime where the expected degree of a node is  $O_{n}(1)$  and consider nodes beyond the immediate neighbours, at any fixed distance. (In fact, our non-asymptotic results allow distances of order  $c\log n$  for small enough  $c > 0$ , see Section 3.5 below.) Furthermore, our main result holds for a general multi-class statistical model with arbitrary continuous or discrete feature distributions and arbitrary edge-connectivity probabilities between all pairs of classes.
Our Contributions. In this paper, we use a multi-class statistical model with arbitrary node features and edge-connectivity profiles among all pairs of classes to study the node classification problem in the regime where the graph component of the data is very sparse, i.e., the expected degree is  $O(1)$ . The data model is described in Section 3.2. We state the following main results and findings:
1. We introduce a family of graph neural network architectures that are asymptotically (in the number of nodes  $n \to \infty$ ) Bayes optimal in a local sense for a general multi-class data model with arbitrary feature distributions. The optimality is stated precisely in Theorem 1.
2. We analyze the architecture in the simpler two-class setting with Gaussian features, explicitly characterizing the generalization error in terms of the natural signal-to-noise ratio (SNR) in the data, and perform a comparative study against other learning methods analyzed using the same statistical model (Theorems 2 and 3). We find two key insights:
- When the graph SNR is very low, the architecture reduces to a simple MLP that does not consider the graph, while if it is very high, our architecture reduces to a typical convolutional network that averages information from all nodes in the local neighbourhood. In the regime between the low and high SNRs, the architecture interpolates and performs better than both a simple MLP and a typical GCN.
- If the information in the graph is larger than a threshold, then a simple convolution is able to perform better than all methods that do not utilize the graph. Not surprisingly, this threshold aligns with the Kesten-Stigum weak-recovery threshold for community detection in sparse networks [Massoulie, 2014, Mossel et al., 2018].
3. In the non-asymptotic setting with a fixed number of nodes, we show that even for a logarithmic depth, the neighbourhoods of an overwhelming fraction of nodes are tree-like with high probability. Subsequently, we show that the optimal classifier in the non-asymptotic setting obtains an error that is close to that incurred by the optimal classifier in the asymptotic setting. This is formalized in Theorem 4.
Let us end this section by reiterating that we work in the fixed-dimensional regime. It is natural to wonder as to the performance of this architecture as compared to optimal algorithms in the high-dimensional setting. We present a numerical comparison of our work to the AMP-BP algorithm from Deshpande et al. [2018] in low and high-dimensional settings in Appendix B.
2 Architecture
This section explains the design of our GNN architecture for node classification. We perform two modifications to existing message-passing architectures. First, we decouple the layers in the neural network from the neighbourhood radius in the message-passing framework. This style of decoupled architecture has previously been studied, see for example, Nikolentzos et al. [2020], Feng et al. [2022], Baranwal et al. [2023]. Second, we introduce a learnable parameter that models edge connectivity between each pair of classes and helps construct the messages to propagate. In the following, for any matrix  $M$ , we denote row  $i$  of  $M$  by  $M_{i,:}$  and column  $i$  of  $M$  by  $M_{i,:i}$ .
Before stating our architecture, we need the following additional notation and pre-processing. Let  $\ell \geq 0$  and  $L > 0$  be fixed integers. Let  $C \geq 2$  be the number of classes. For given data  $(\mathbf{A}, \mathbf{X})$  where  $\mathbf{A} \in \mathbb{R}^{n \times n}$  is the adjacency matrix of an unweighted undirected graph, and  $\mathbf{X} \in \mathbb{R}^{n \times d}$  is the node feature matrix, we perform a pre-computation on the graph to construct a tensor  $\tilde{\mathbf{A}}$  as follows:
$$
\tilde {\mathbf {A}} ^ {(k)} = f (\mathbf {A} ^ {k}) \wedge \left(\neg f \left(\sum_ {m = 0} ^ {k - 1} \mathbf {A} ^ {m}\right)\right) \text {f o r} k \in \{1, \dots , \ell \},
$$
where  $f(M)$  for a matrix  $M$  returns the entry-wise flattened matrix with  $f(M)_{ij} = \mathbb{1}(M_{ij} > 0)$ , and  $(\wedge, \neg)$  denote the entry-wise bit-wise operators ('and', 'negation') respectively. Note here that  $\tilde{\mathbf{A}}^{(k)}$  is an  $n \times n$  binary matrix with  $\tilde{\mathbf{A}}_{uv}^{(k)} = 1$  if and only if  $v$  is present in the distance  $k$  neighbourhood of  $u$  but not within the distance  $(k - 1)$  neighbourhood. The idea behind this pre-processing step is the following: for each node  $u$  and each  $k \in [\ell]$ , we want to divide the radius  $\ell$  neighbourhood of  $u$  into  $\ell$  groups of nodes, where each group  $k \in [\ell]$  consists of nodes that are within discovered the neighbourhood at each distance from a given node.  $\tilde{\mathbf{A}}_{u,:}^{(k)}$  models a non-backtracking walk of length  $k$  that considers new nodes in the distance-  $k$  neighbourhood that were not discovered.
We can now define the graph neural network architecture as follows.
Architecture 1. Given input data  $(\mathbf{A},\mathbf{X})$  where  $\mathbf{A}\in \{0,1\}^{n\times n}$  is the adjacency matrix and  $\mathbf{X}\in \mathbb{R}^{n\times d}$  is the node feature matrix, define:
$$
\mathbf {H} ^ {(0)} = \mathbf {X}, \quad \mathbf {H} ^ {(l)} = \sigma_ {l} \left(\mathbf {H} ^ {(l - 1)} \mathbf {W} ^ {(l)} + \mathbf {1} _ {n} \mathbf {b} ^ {(l)}\right) \text {f o r} l \in [ L ],
$$
$$
\mathbf {Q} = \mathrm {s i g m o i d} (\mathbf {Z}), \qquad \mathbf {M} _ {u, i} ^ {(k)} = \max _ {j \in [ C ]} \left\{\mathbf {H} _ {u, j} ^ {(L)} + \log (\mathbf {Q} _ {i, j} ^ {k}) \right\} \mathrm {f o r} k \in [ \ell ], u \in [ n ], i \in [ C ].
$$
Then the predicted label is given by  $\hat{\mathbf{y}} = \{\hat{y}_u\}_{u\in [n]}$ , where
$$
\hat {y} _ {u} = \operatorname * {a r g m a x} _ {i \in [ C ]} \left(\mathbf {H} _ {u, c} ^ {(L)} + \sum_ {k = 1} ^ {\ell} \tilde {\mathbf {A}} _ {u,:} ^ {(k)} \mathbf {M} _ {:,: i} ^ {(k)}\right).
$$
Let us pause here to comment on the interpretation of the terms arising in this architecture. Here,  $\mathbf{H}^{(L)}$  is viewed as the output of a simple  $L$ -layer MLP with  $\{\sigma_l\}_{l\in [L]}$  being a set of non-linear functions. We have  $(\mathbf{W}^{(l)},\mathbf{b}^{(l)})_{l\in [L]}$  as the learnable parameters of this MLP, with suitable dimensions so that  $\mathbf{H}^{(L)}\in \mathbb{R}^{n\times C}$ . In addition, we introduce the learnable parameter  $\mathbf{Z}\in \mathbb{R}^{C\times C}$  which is used to model edge connectivity among all pairs of classes. The quantity  $\tilde{\mathbf{A}}_{u,:}^{(k)}\mathbf{M}_{:,i}^{(k)} = \sum_{v\in [n]}\tilde{\mathbf{A}}_{u,v}^{(k)}\mathbf{M}_{v,i}^{(k)}$  is viewed as the sum of messages  $\mathbf{M}_{v,i}^{(k)}$  passed by all distance  $k$  neighbours of node  $u$ .
Although Architecture 1 follows the style of convolutional architectures like GCN and GAT for collecting messages within a local neighbourhood, the novelty lies in the construction of the messages  $\mathbf{M}$ . Intuitively,  $\mathbf{Q} = \mathrm{sigmoid}(\mathbf{Z})$  learns the probabilities of edge connectivity between all pairs of classes so that  $\mathbf{Q}^k$  models the probability of observing a distance  $k$  path between a pair of nodes in two classes. To predict the label of node  $u$ , the messages from other nodes  $v$  are constructed based on their features  $\mathbf{X}_v$ , their distance  $k$  from node  $u$ , and the path probabilities  $\mathbf{Q}^k$ . We show in Theorem 1 that this architecture is in a sense (made precise in Definition 3.2) universally optimal among all node-classification schemes for sparse graphs. Our result thus aligns with the observations in Velicković [2022], showing that optimal neural network architectures for node classification on sparse graphs are implementable using the message-passing paradigm.
3 Theoretical Analysis and Discussion
In this section, we present a theoretical analysis of the message-passing GNN given in Architecture 1. We begin by defining a natural notion of optimality in our setting and show that among local learning methods on graphs, Architecture 1 is optimal according to this definition on a very general statistical model. We then compute the generalization error and compare the architecture to other well-studied methods like a simple MLP and a GCN.
3.1 Asymptotic Local Bayes Optimality
For classification tasks, it is natural to use a notion of generalization error in a "per sample" or online sense. Without graphical side information, the natural choice is the Bayes risk. With graphical information, however, there is an important obstruction: the number of samples is equal to the size of the corresponding graph. As such, a naive extension of the Bayes risk does not have this property.
A natural approach would be to consider the Bayes risk for estimators that take in the node, the data set, and the graph, i.e.,  $\hat{y}_v = \hat{y} (v,(X,G))$ . In this case, however, the risk necessarily implicitly depends on the sample size,  $n$ , through  $G$ . One might try to remove this dependence by taking the infinite sample size limit, but for a class of estimators this general, it is not clear that such a limit is well defined. To circumvent this issue, we restrict attention to node classifiers that are only allowed "local" information around the node. The large graph limit of the generalization error is then naturally interpreted via local weak convergence. (For the convenience of the reader, we briefly recall the notion of local weak convergence of sparse graphs in Appendix A.1. See also Ramanan [2021, Chapter 1] or Bordenave [2016, Section 3] for more detailed expositions.) In this limit, one can then interpret the generalization error as a per-sample error for the randomly rooted graph  $(G,u)$  where  $u$  is a uniform random vertex in  $V(G)$ . (Here and in the following a rooted graph is a pair of a graph  $G$
and a distinguished vertex,  $u$ , called the root.) With these observations, we are led to a natural notion of Bayes optimality, namely asymptotic local Bayes optimality which we define presently.
Before turning to this definition, we must first recall the notion of  $\ell$ -local classifiers. For a node  $v$  in a graph  $G$ , let  $\eta_k(v) = \{u \in V(G) : d(u,v) \leq k\}$  denote the ball of radius  $k$  for the canonical graph distance metric.
Definition 3.1 ( $\ell$ -local classifier). Let  $G = (\mathbf{A}, \mathbf{X})$  be a feature-decorated graph of  $n$  vertices with  $d$ -dimensional features  $\mathbf{X}_u$  for each vertex  $u$ . For a fixed radius  $\ell > 0$ , an  $\ell$ -local node-classifier is a function  $f$  that takes as input a root vertex  $u \in [n]$ , the adjacency matrix  $\mathbf{A}$  and the features of all nodes within the  $\ell$ -neighbourhood of  $u$ , i.e.,  $\{\mathbf{X}_v\}_{v \in \eta_\ell(u)}$ , and outputs a classification label for  $u$ .
Suppose now that we have a sequence of (random) feature decorated graphs  $(X_{n},G_{n})$  with  $|V(G)| = n$ . Let  $u_{n}$  denote a uniform at random vertex in  $G_{n}$ . Suppose finally that the rooted feature-decorated graphs  $(X_{n},G_{n},u_{n})$  locally weakly converge to  $(X,G,u)$ . We can then define the notion of asymptotically  $\ell$ -locally Bayes optimal classifiers for this sequence of problems.
Definition 3.2. We say that a classifier  $h_{\ell}^{*} \in \mathcal{C}_{\ell}$  is asymptotically  $\ell$ -locally Bayes optimal classifier of the root for the sequence  $\{(X_{n}, G_{n}, u_{n})\}$  if it minimizes the probability of misclassification of the root of the local weak limit,  $(X, G, u)$ , over the class  $\mathcal{C}_{\ell}$ , i.e.,
$$
h_{\ell}^{*} = \operatorname *{argmin}_{h\in \mathcal{C}_{\ell}}\mathbf{Pr}\left[h(u,\{\mathbf{X}_{v}\}_{v\in \eta_{\ell}(u,G)})\neq y_{u}\right].
$$
Before turning to our data model, we note here that the reader may ask whether or not the asymptotically  $\ell$ -locally Bayes optimal classifier is in any sense the limit of optimal  $\ell$ -local classifier of the random root,  $u_{n}$ . We show this in an appropriate sense in Theorem 4.
3.2 Data Model
Let us now turn to the data model that we use for our analysis. We work with the general multi-class contextual stochastic block model (CSBM) where each node belongs to one of  $C$  different classes labelled  $1, \ldots, C$ , and the node features have arbitrary continuous or discrete distributions. This model with  $C = 2$ , along with a specialization to Gaussian features has been extensively studied in several works on (semi)-supervised node classification and unsupervised community detection, see, for example, Deshpande et al. [2018], Lu and Sen [2020], Baranwal et al. [2021], Wei et al. [2022], Fountoulakis et al. [2022b], Baranwal et al. [2023]. Informally, a CSBM consists of a coupling of a stochastic block model (SBM) [Holland et al., 1983] with a mixture model where the components of the mixture have arbitrary distributions and are associated with the blocks of the SBM.
More formally, let  $n, d$  be positive integers such that  $n$  denotes the number of nodes and  $d$  denotes the dimension of the node features. Define  $y_{1},\ldots ,y_{n}\in \{1,\dots ,C\}$  as the latent variables (class labels) to be inferred. We will assume that the latent variables have a uniform prior, i.e.,  $y_{u}\sim \mathrm{Unif}(\{[C]\})$  for all  $u$ . For the relational part of the data, we have an undirected unweighted graph of  $n$  nodes,  $G = (V,E)$  with adjacency matrix  $\mathbf{A} = (a_{uv})_{u,v\in [n]}\sim \mathrm{SBM}(n,\mathbf{Q})$ , where  $\mathbf{Q} = \{q_{ij}\} \in [0,1]^{C\times C}$  is the edge-probability matrix, meaning that
$$
\Pr \left(a _ {u v} = 1 \mid y _ {u} = i, y _ {v} = j\right) = q _ {i j}.
$$
The node attributes,  $\mathbf{X} \in \mathbb{R}^{n \times d}$  are sampled from a mixture of  $C$  arbitrary continuous or discrete distributions,  $\mathbb{P} = \{\mathbb{P}_i\}_{i \in [C]}$ , where corresponding to the  $y_u$ , we have  $\mathbf{X}_u \sim \mathbb{P}_{y_u}$  for all  $u \in [n]$ .
We will view  $n$  as large and study the setting where  $d$  is fixed (does not grow with  $n$ ). We note here that in previous related works [Baranwal et al., 2021, Wei et al., 2022, Baranwal et al., 2023], crucial assumptions have been made about the distribution of the node features and the sparsity of the graph, i.e.,  $q_{ij} = \Omega_n(\log^2 n / n)$ . In contrast, we work in the extremely sparse setting where  $q_{ij} = b_{ij} / n$  for constants  $b_{ij} > 1$ , so we write  $\mathbf{Q} = \mathbf{B} / n$  where  $\mathbf{B} = \{b_{ij}\}_{i,j\in [C]}$ . Furthermore, the only assumption we need about the distributions  $\mathbb{P}_i$  is that  $\mathbb{P}_i$  are absolutely continuous with respect to some base measure, in which case their densities exist, denoted by  $\rho_i$ . For ease of reading, we encourage the
<sup>1</sup>It is also desirable for the empirical misclassification error to converge to the generalization error. If one works with the stronger notion of local convergence in probability, then this will hold as well. This later mode will hold in our examples but we leave this to future work.
reader to consider the case where  $\mathbb{P}_i$  are continuous or discrete, therefore, the base measure is simply the Lebesgue measure on  $\mathbb{R}$  or the counting measure on  $\mathbb{Z}$  respectively.
For a feature-decorated graph  $G = (\mathbf{A}, \mathbf{X}) = (\{a_{uv}\}_{u, v \in [n]}, \{\mathbf{X}_u\}_{u \in n})$  sampled from the model described above, we say that  $G \sim \mathrm{CSBM}(n, d, \mathbb{P}, \mathbf{Q})$  or  $G \sim \mathrm{CSBM}(n, d, \mathbb{P}, \mathbf{B} / n)$ .
3.3 Optimal Classifier
We are now ready to state our first main result that characterizes the asymptotically  $\ell$ -locally Bayes optimal classifier on the CSBM data described in Section 3.2.
Theorem 1 (Bayes optimal message-passing). For any  $\ell \geq 1$ , the asymptotically  $\ell$ -locally Bayes optimal classifier of the root for the sequence  $(G_{n},u_{n})\sim \mathrm{CSBM}(n,d,\mathbb{P},\mathbf{Q})$  is
$$
h _ {\ell} ^ {*} (u, \left\{\mathbf {X} _ {v} \right\} _ {v \in \eta_ {\ell} (u)}) = \operatorname * {a r g m a x} _ {i \in [ C ]} \left\{\log \rho_ {i} (\mathbf {X} _ {u}) + \sum_ {v \in \eta_ {\ell} (u) \backslash \{u \}} \mathcal {M} _ {i d (u, v)} (\mathbf {X} _ {v}) \right\},
$$
where  $\{\rho_i\}_{i\in [C]}$  are the densities associated with the distributions  $\mathbb{P}_i\in \mathbb{P}$ , and
$$
\mathcal {M} _ {i k} (\mathbf {x}) = \max  _ {j \in [ C ]} \left\{\log \rho_ {j} (\mathbf {x}) + \log \mathbf {Q} _ {i j} ^ {k} \right\}.
$$