Statistical Guarantees for Link Prediction using Graph Neural Networks

Graph Neural Networks (GNNs) have emerged as a powerful tool for link prediction (Zhang & Chen, 2018; Zhang, 2022). A significant advantage of GNNs lies in their adaptability to different graph types. Traditional link prediction heuristics tend to presuppose network characteristics. For example, the common neighbors heuristic presumes that nodes that share many common neighbors are more likely to be connected, which is not necessarily true in biological networks (Kovács, 2019). In contrast, GNNs inherently learn predictive features through the training process, presenting a more flexible and adaptable method for link prediction.

This paper provides statistical guarantees for link prediction using GNNs, in graphs generated by the graphon model. A graphon is specified by a symmetric measurable kernel function $W:\Omega^{2}\rightarrow[0,1]$. A graph $G_{n}=(V_{n},E_{n})$ with the vertex set $V_{n}=\{1,2,\cdots,n\}$ is sampled from $W$ as follows: (i) each vertex $i\in V_{n}$ draws latent feature $(\omega_{i})\overset{i.i.d}{\sim}\mu$ for some probability distribution $\mu$ on $\Omega\subset\mathbb{R}^{q}$; (ii) the edges of $G_{n}$ are generated independently and with probability $W_{n,i,j}:=\rho_{n}\cdot W(\omega_{i},\omega_{j}),$ where $\rho_{n}\in(0,1]$ is a constant called the sparsifying factor¹¹1Throughout the paper, we will assume $(\omega_{i})\sim\text{Unif}[0,1]$. This is without loss of generality. For any graphon $\tilde{W}$ with features in some arbitrary $\Omega\subset\mathbb{R}^{q}$ sampled from $\mu$ on $\Omega$, there exists some graphon $W$ with latent features drawn from $\text{Unif}[0,1]$ so that the graphs generated from these two graphons are equivalent in law. See Remark 4 in (Davison & Austern, 2023) for more details..

The graphon model includes various widely researched graph types, such as Erdos-Renyi, inhomogeneous random graphs, stochastic block models, degree-corrected block models, random exponential graphs, and geometric random graphs as special cases; see (Lovász, 2012) for a more detailed discussion.

This paper’s key contribution is the analysis of a linear Graph Neural Network model (LG-GNN) which can provably estimate the edge probabilities in a graphon model. Specifically, we present a GNN-based algorithm that yields estimators, denoted as $\hat{p}_{i,j}$, that converge to true edge probabilities $\rho_{n}W(\omega_{i},\omega_{j})$. Crucially, these estimators have mean squared error converging to 0 at the rate $o_{n\to\infty}(\rho_{n}^{2})$. To our knowledge, this work is the first to rigorously characterize the ability of GNNs to estimate the underlying edge probabilities in general graphon models.

The estimators $\hat{p}_{i,j}$ are constructed in two main steps. We first employ LG-GNN (Algorithm 1), to embed the vertices of $G_{n}$. Concretely, for each vertex $i\in[n]$, LG-GNN computes a set $\Lambda_{i}=\{\lambda_{i}^{0},\lambda_{i}^{1},\dots,\lambda_{i}^{L}\},$ where $\lambda_{i}^{k}\in\mathbb{R}^{d_{n}}$ and $d_{n}$ is the embedding dimension. Then, $\Lambda_{i},\Lambda_{j}$ are used to construct estimators $\hat{q}_{i,j}^{(k)}$ for the moments of $W$. We refer the reader to Section 4 for the formal definition of the moments. Intuitively, the $k$th moment $W_{n,i,j}^{(k)}$ represents the probability that there is a path of length $k$ between vertices with latent features $\omega_{i}$ and $\omega_{j}$ in $G_{n}$.

The next step is to show that when the number of distinct nonzero eigenvalues of $W$, denoted $m_{W}$, is finite, then the edge probabilities $W_{n,i,j}$ can be written as a linear function of the moments $W_{n,i,j}^{(2:m_{W}+1)}$. This naturally motivates Algorithm 2, which learn the $W_{n,i,j}$’s from the moment estimators $\hat{q}_{i,j}^{(2:m_{W}+1)}$, using a constrained regression. The regression coefficients $\hat{\beta}^{n,m_{W}}$ are then used to produce estimators $\hat{p}_{i,j}$ for $W_{n,i,j}$.

The main result of the paper (stated in Theorem 4.4) presents the convergence rate of the mean square error of $\hat{\beta}^{n,m_{W}}$. It shows that if $L$, the number of message-passing layers in LG-GNN, is at least $m_{W}-1$, then the mean square error converges to 0. That implies that our estimators for the edge probabilities $W_{n,i,j}$ are consistent. For $L<m_{W}-1$, the theorem provides the rate at which the mean square error decreases when $L$ increases. The second main result, stated in Proposition 4.5, gives statistical guarantees on how well LG-GNN can detect in-community edges in a symmetric stochastic block model. A notable feature of this theorem is that the implied convergence rate is much faster than that of Theorem 4.4, which demonstrates mathematically that ranking high and low probability edges is easier than estimating the underlying probabilities of edges.

Finally, we would like to highlight two key aspects of the results of the paper. Firstly, our statistical guarantees for edge prediction are proven for scenarios when node features are absent and the initial node embeddings $\lambda_{i}^{0}$ are chosen at random. This underscores that effective link prediction can be achieved solely through the appropriate selection of GNN architecture, even in the absence of additional node data. The second point relates to graph sparsity: although graphons typically produce dense graphs, introducing the sparsity factor $\rho_{n}$ results in vertex degrees of $O(\rho_{n}\cdot n)$, facilitating the exploration of sparse graphs. Our findings are pertinent for $\log(n)/n\ll\rho_{n}\leq 1$. Note that a sparsity of $\log(n)/n$ is the necessary threshold for connectivity (Spencer, 2001), highlighting the generality of our results.

While the primary focus of this paper is theoretical, we complement our theoretical analysis with experimental evaluations on real-world datasets (specifically, the Cora dataset) and graphs derived from random graph models. Our empirical observations reveal that in scenarios where node features are absent, LG-GNN exhibits performance comparable to the traditional Graph Convolutional Network (GCN) on simple random graphs, and surpasses GCN in more complex graphs sampled from graphons. Additionally, LG-GNN presents two further benefits: LG-GNN does not involve any parameter tuning (e.g., through the minimization of a loss function), resulting in significantly faster operation, and it avoids the common oversmoothing issues associated with the use of numerous message-passing layers.

Link prediction on graphs have a wide range of applications in domains ranging from social network analysis to drug discovery (Hasan & Zaki, 2011; Abbas et al., 2021). A survey of techniques and applications can be found in (Kumar et al., 2020; Martínez et al., 2016; Djihad Arrar, 2023).

Much of the existing theory on GNNs is regarding their expressive power. For example, (Xu et al., 2018; Morris et al., 2021) show that GNNs with deterministic node initializations have expressive power bounded by that of the 1-dimensional Weisfeiler-Lehman (WL) graph isomorphism test. Generalizations such as $k$-GNN (Morris et al., 2021) have been proposed to boost the expressive power higher in the WL-hierarchy. The Structural Message Passing GNN (SGNN) (Vignac et al., 2020) was also proposed and was shown to be universal on graphs with bounded degrees, and converges to continuous ”c-SGNNs” (Keriven et al., 2021), which were also shown to be universal on many popular random graph models. Lastly, (Abboud et al., 2020) showed that GNNs that use random node initializations are universal, in that they can approximate any function defined on graphs with fixed order.

A recent wave of works focus on deriving statistical guarantees for graph representation algorithms. A common data-generating model for the graph is a graphon model (Lovász & Szegedy, 2006; Borgs et al., 2008, 2012). A large literature has been devoted to establishing guarantees for community detection on graphons such as the stochastic block model; see (Abbe, 2018) for an overview. For this task, spectral embedding methods have long been proposed (see (Deng et al., 2021; Ma et al., 2021) for some recent examples). Lately, statistical guarantees for modern random walk-based graph representation learning algorithms have also been obtained. Notably (Davison & Austern, 2021; Barot et al., 2021; Qiu et al., 2018; Zhang & Tang, 2021) characterize the asymptotic properties of the embedding vectors obtained by deepwalk, node2vec, and their successors and obtain statistical guarantees for downstream tasks such as edge prediction. Recently, some works also aim at obtaining learning guarantees for GNNs. Stability and transferability of certain untrained GNNs have been established in (Ruiz et al., 2021; Maskey et al., 2023; Ruiz et al., 2023; Keriven et al., 2020). For example (Keriven et al., 2020) shows that for relatively sparse graphons, the embedding produced by an untrained GNN will converge in $L^{2}$ to a limiting embedding that depends on the underlying graphon. They use this to study the stability of the obtained embeddings to small changes in the training distribution. Other works established generalization guarantees for GNNs. Those depend respectively on the number of parameters in the GNN (Maskey et al., 2022), or on the VC dimension and Radamecher complexity (Esser et al., 2021) of the GNN.

Differently from those two lines of work, our paper studies when link prediction is possible using GNNs, establishes statistical guarantees for link prediction, and studies how the architecture of the GNN influences its performance. More similar to our paper is (Kawamoto et al., 2018), which exploits heuristic mean-field approximations to predict when community detection is possible using an untrained GNN. Note, however, that contrary to us, their results are not rigorous and the accuracy of their approximation is instead numerically evaluated. (Lu, 2021) formally established guarantees for in-sample community detection for two community SBMs with a GNN trained via coordinate descent. However, our work establishes learning guarantees for general graphons beyond two-community SBMs, both in the in-sample and out-sample settings. Moreover, the link prediction task we consider, while related to community detection, is still significantly different. (Baranwal et al., 2021) studies node classification for contextual SBMs and shows that an oracle GNN can significantly boost the performance of linear classifiers. Another related work (Magner et al., 2020) studies the capacity of GNN to distinguish different graphons when the number of layers grows at least as $L=\Omega(\log(n))$. Interestingly they find that GNN struggles in differentiating graphons whose expected degree sequence is not sufficiently heterogeneous, which unfortunately occurs for many graphon models, including the symmetric SBM. It is interesting to note that in Proposition 5.1 we will show that this is also the regime where the classical GCN fails to provide reliable edge probability prediction. Finally, some learning guarantees have also been derived for other graph models. Notably (Alimohammadi et al., 2023) studied the convergence of GraphSAGE and related GNN architectures under local graph convergence.

In this section, we present our assumptions, some background regarding GNNs, and the link prediction goals that we focus on.

We introduce the Linear Graphon Graph Neural Network, or LG-GNN in Algorithm 1. The algorithm starts by assigning each node $i$ a random feature $Z_{i}\sim\frac{1}{\sqrt{d_{n}}}\mathcal{N}(0,I_{d_{n}}),$ where $d_{n}=\Omega(1/\rho_{n})$ is the embedding dimension. The first message passing layer computes $\lambda_{i}^{0}$ by summing $Z_{j}$ for all $j\in N(i)$, scaled by $\frac{1}{\sqrt{n}}.$ The subsequent layers normalize the $\lambda_{j}^{k}$’s by $1/n$ before adding them to $\lambda_{i}^{k}$. We show in Proposition D.3 and Lemma D.4 that this procedure essentially counts the number of paths between pairs of vertices. Specifically, $\operatorname{\mathbb{E}}[\langle\lambda_{i}^{k_{1}},\lambda_{j}^{k_{2}}\rangle|A,(\omega_{\ell})]$ is a linear combination of the ”empirical moments” of $W$ (Equation 7). The second stage of Algorithm 1 then recovers these empirical moments by decoupling the aforementioned linear equations. We refer the reader to Section D.2 for more details and intuition behind LG-GNN.

We note that the scaling of $1/\sqrt{n}$ in the first message passing layer is crucial in allowing the embedding vectors $(\lambda_{\ell}^{k})$ to learn information about the latent features $(\omega_{\ell})$ asymptotically. We show in Proposition 5.1 that without this construction, the classical GCN is unable to produce meaningful emebddings with random feature initializations.

As mentioned in Section 4, in the context of random node initializations, a naive choice of GNN architecture can cause learning to fail. In the following proposition, we demonstrate that for a large class of graphons, the Classical GCN architecture with random initializations results in embeddings that cannot be informative in out-of-sample graphon estimation. To make this formal, we assume that at training, only $n-m$ vertices are observable. We denote by $G_{n}|_{V_{n-m}}$ the induced subgraph with vertex set $V_{n-m}=\{1,\dots,n-m\}$. And we consider graphons that are such that

Note that many graphons satisfy this assumption, including symmetric SBMs.

We show that this leads to suboptimal risk for graphon estimation. For simplicity we show this for dense graphons, e.g when $\rho_{n}=1$.

Proposition 5.1 and Proposition 5.2 imply that in the out-of-sample setting, the embeddings produced by Equation 1 with random node feature initializations will lead to sub-optimal estimators for the edge probability $W(\omega_{i},\omega_{i})$. A key feature of the proof of Proposition 5.1 is that $\sum_{u\in N(v)}\frac{\lambda_{u}^{k-1}}{\sqrt{|N(u)\|N(v)|}}$ concentrates to 0 very quickly with random node initializations. This demonstrates the importance of the (subtle) construction of the first round of message passing $\lambda_{u}^{0}$ in Algorithm 1. We also note that Proposition 5.2 doesn’t necessarily imply that predicted probabilities $\hat{p}_{e_{i}}$ will be ineffective at ranking test edges, though we do see in the experiments that the performance is decreased for the out-of-sample case.

We remark that a key feature of LG-GNN is that it uses the embedding vector $\lambda_{i}^{k}$ produced at each layer. Indeed $\hat{p}_{i,j}$ depends on all of the terms $\{\langle\lambda_{i}^{0},\lambda_{j}^{0}\rangle,\langle\lambda_{i}^{0},\lambda_{j}^{1}\rangle,\dots,\langle\lambda_{i}^{0},\lambda_{j}^{L}\rangle\}$. This is in contrast to many classical ways of using GNNs for link prediction that depend only on $\langle\lambda_{i}^{L},\lambda_{j}^{L}\rangle.$ The following proposition shows that this construction is necessary to obtain consistent estimators.

To illustrate this, consider the following example for $L=0.$

While this result is for the specific case of Algorithm 1, which in particular contains no non-linearities, we anticipate that this general procedure of learning a function that maps a set of dot products $\{\langle\lambda_{i}^{k_{1}},\lambda_{j}^{k_{2}}\rangle\}_{k_{1},k_{2}}$ to a predicted probability, instead of just $\langle\lambda_{i}^{L},\lambda_{j}^{L}\rangle$, can lead to better performance for practioners on various types of GNN architectures.

We compare experimentally a GCN, LG-GNN, and PLSG-GNN. We perform experiments on the Cora dataset (McCallum et al., 2000) in the in-sample setting. We also show results for various random graph models. The results for random graphs below are in the out-sample setting, more results are in Appendix I. We report the AUC-ROC and Hits@k metric, and also a custom metric called the Probability Ratio@k, which is more suited to the random graph setting. We refer the reader to Appendix I for a more complete discussion.

LG-GNN and PLSG-GNN perform similarly to the classical GCN in settings with no node features and can outperform it on more complex graphons. One major advantage of LG-GNN/PLSG-GNN is that they do not require extensive tuning of hyperparameters (e.g., through minimizing a loss function) and hence run much faster and are easier to fit. For example, training the 4-layer GCN resulted in convergence issues, even on a wide set of learning rates.

This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.

The first author would like to thank Qian Huang for helpful discussions regarding the experiments. The authors would also like to thank the Simons Institute for the Theory of Computing, specifically the program on Graph Limits and Processes on Networks. Part of this work was done while Austern and Saberi were at the Simon’s Institute for the Theory of Computing.

This research is supported in part by the AFOSR under Grant No. FA9550-23-1-0251, and by an ONR award N00014-21-1-2664.