Efficient Link Prediction via GNN Layers Induced by Negative Sampling

Let $\mathcal{G}=(\mathcal{V},\mathcal{E},X)$ be a graph with node set $\mathcal{V}$, corresponding $d_{x}$-dimensional node features $X\in\mathbb{R}^{n\times d_{x}}$, and edge set $\mathcal{E}$, where $|\mathcal{V}|=n$. We use $A$ to denote the adjacency matrix and $D$ for the degree matrix. The associated Laplacian matrix is defined by $L\triangleq D-A$. Furthermore, $Y\in\mathbb{R}^{n\times d}$ refers to node embeddings of size $d$ we seek to learn via a node-wise link prediction procedure. Specifically the node embedding for node $i$ is $y_{i}$, which is equivalent to the $i$-th row of $Y$.

We begin by introducing a commonly-used loss for link prediction, which is defined over the training set $\mathcal{E}_{train}\subset\mathcal{E}$. For both node-wise and edge-wise methods, the shared goal is to obtain an edge probability score $p(v_{i},v_{j})=\sigma(e_{ij})$ for all edges $(v_{i},v_{j})\in\mathcal{E}_{train}$ (as well as negatively-sampled counterparts to be determined shortly), where $\sigma$ is a sigmoid function and $e_{ij}$ is a discriminative representation for edge $(v_{i},v_{j})$. Proceeding further, for every true positive edge $(v_{i},v_{j})$ in the training set, $N\geq 1$ negative edges $(v_{i^{a}},v_{j^{a}})_{a=1,...,N}$ are randomly sampled from the graph for supervision purposes. We are then positioned to express the overall link prediction loss as $~{}~{}~{}\mathcal{L}_{link}\triangleq~{}=$

where each edge probability is computed with the corresponding edge representation. The lingering difference between node- and edge-wise methods then lies in how each edge representation $e_{ij}$ is actually computed.

For node-wise methods, $e_{ij}=h(y_{i},y_{j})$, where $y_{i}$ and $y_{j}$ are node-wise embeddings and $h$ is a decoder function ranging in complexity from a parameter-free inner-product to a multi-layer MLP. While decoder structure varies [42, 43, 44, 45, 16], of particular note for its practical effectiveness is the HadamardMLP approach, which amounts to simply computing the hadamard product between $y_{i}$ and $y_{i}$ and then passing the result through an MLP with trainable parameters $\theta$. Fast, sublinear inference times are possible with HadamardMLP using an algorithm from [16]. In contrast, the constituent node embeddings themselves are typically computed with some form of trainable GNN encoder model $g$ of the form $y_{i}=g(x_{i},\mathcal{G}_{i})$ and $y_{j}=g(x_{j},\mathcal{G}_{j})$, where $\mathcal{G}_{i}$ and $\mathcal{G}_{j}$ are the subgraphs containing nodes $v_{i}$ and $v_{j}$, respectively.

Turning to edge-wise methods, the edge representation $e_{ij}$ relies on the subgraph $\mathcal{G}_{ij}$ defined by both $v_{i}$ and $v_{j}$. In this case $e_{ij}=h_{e}(v_{i},v_{j},\mathcal{G}_{ij})$, where $h_{e}$ is an edge encoder GNN whose predictions can generally not be decomposed into a function of individual node embeddings as before. Note also that while the embeddings from node-wise subgraphs for all nodes in the graph can be produced by a single GNN forward pass, a unique/separate edge-wise subgraph and corresponding forward pass are needed to make predictions for each candidate edge. This explains why edge-wise models endure far slower inference speeds in practice.

Prior work on optimization-based node embeddings [46, 26, 27, 19, 28, 29] largely draw on energy functions related to [47], which facilitates the balancing of local consistency relative to labels or a base predictor, with global constraints from graph structure. However, these desiderata alone are inadequate for the link prediction task, where we would also like to drive individual nodes towards regions of the embedding space where they are maximally discriminative with respect to their contributions to positive and negative edges. To this end we take additional inspiration from triplet ranking loss functions [48] that are explicitly designed for learning representations that can capture relative similarities or differences between items.

With these considerations in mind, we initially posit the energy function    $\ell_{node}~{}~{}\triangleq||Y-f(X;W)||^{2}_{F}+$

where $f(X;W)$ (assumed to apply row-wise to each individual node feature $x_{i}$) represents a base model that processes the input features using trainable weights $W$, $dist(y_{i},y_{j})$ is a distance metric, while $\lambda$ and $\lambda_{K}$ are hyperparameters that control the impact of positive and negative edges. Moreover, $\mathcal{V}_{(i,j)}^{K}$ is the set of negative destination nodes sampled for edge $(v_{i},v_{j})$ and $|\mathcal{V}_{(i,j)}^{K}|=K$. Overall, the first term pushes the embeddings towards the processed input features, while the second and third terms apply penalties to positive and negative edges in a way that is loosely related to the aforementioned triplet ranking loss (more on this below).

If we choose $dist(i,j)=||y_{i}-y_{j}||^{2}$ and define edges of the negative graph $\mathcal{G}^{-}$ as $\mathcal{E}^{-}\triangleq\{(i,j^{\prime})|~{}j^{\prime}\in\mathcal{V}_{(i,j)}^{K}\text{ for }(i,j)\in\mathcal{E}\}$, we can rewrite (2) as $~{}~{}~{}\ell_{node}=$

where $L^{-}$ is the Laplacian matrix of $\mathcal{G}^{-}$. To find the minimum, we compute the gradients

where $Y^{(0)}=f(X;W)$. The corresponding gradient descent updates then become $Y^{(t+1)}=$

where the step size is $\frac{\alpha}{2}$. We note however that (3) need not generally be convex or even lower bounded. Moreover, the gradients may be poorly conditioned for fast convergence depending on the Laplacian matrices involved. Hence we consider several refinements next to stabilize the learning process.

Lower-Bounding the Negative Graph. Since the regularization of negative edges brings the possibility of an ill-posed loss surface (a non-convex loss surface that may be unbounded from below) , we introduce a convenient graph-aware lower-bound analogous to the max operator used by the triplet loss. Specifically, we update (3) to the form $~{}~{}~{}\ell_{node}~{}=||Y-f(X;W)||^{2}_{F}+$

noting that we use $\text{Softplus}(x)=\log(1+e^{x})$ instead of $max(\cdot,0)$ to make the energy differentiable. Unlike triplet loss that includes positive term in the $max(\cdot,0)$ function, we only restrict the lower bound for the negative term, because we still want the positive part to impact our model when the negative part hits the bound.

Normalization. We use the normalized Laplacian matrices of original graph $\mathcal{G}$ and negative graph $\mathcal{G}^{-}$ instead to make the gradients smaller. Also, for the gradients of the first term in (3), we apply a re-scaling by summation of degrees of both graphs. The modified gradients are $\frac{\partial\ell_{node}}{\partial Y}=$

where $D$ and $D^{-}$ are diagonal degree matrices of $\mathcal{G}$ and $\mathcal{G}^{-}$. The normalized Laplacians are $\tilde{L}=D^{-\frac{1}{2}}LD^{-\frac{1}{2}},\tilde{L}^{-}={D^{-}}^{-\frac{1}{2}}L^{-}{D^{-}}^{-\frac{1}{2}}$ leading to the corresponding energy

Learning to Combine Negative Graphs. We now consider a more flexible implementation of negative graphs. More concretely, we sample $K$ negative graphs ${\{\mathcal{G}^{-}_{(k)}\}}_{k=1,...,K}$, in which every negative graph consists of one negative edge per positive edge ($\mathcal{G}^{-}_{(k)}\triangleq\{(i,j^{\prime})|~{}j^{\prime}\in\mathcal{V}_{(i,j)}^{1}\text{ for }(i,j)\in\mathcal{E}\}$) (note that the superscript on $\mathcal{V}$ implies a single negative sample per positive edge). And we set learnable weights $\lambda_{K}^{k}$ for the structure term of each negative graph, which converts the energy function to $~{}\ell_{node}~{}=$

Also in practice we normalize this energy function to

where $D^{-}_{K}=\sum_{k=1}^{K}D^{-}_{k}$, $D^{-}_{k}$ is the degree matrix of $L^{-}_{k}$ and $\tilde{L^{-}_{k}}={D^{-}_{K}}^{-\frac{1}{2}}L^{-}_{k}{D^{-}_{K}}^{-\frac{1}{2}}.$ The lower bound is also added as before. Overall, the motivation here is to inject trainable flexibility into the negative sample graph, which is useful for increasing model expressiveness.

Combining the modifications we discussed in last section (and assuming a single, fixed $\lambda_{K}$ here for simplicity; the more general, learnable case with multiple $\lambda_{K}^{k}$ from Section IV-B naturally follows), we obtain the final energy function

with the associated gradients$~{}\frac{\partial\ell_{node}}{\partial Y}~{}=$

where $Q=\gamma|\mathcal{E}|-\text{tr}[Y^{\top}\tilde{L}^{-}Y]$ and $\sigma(x)$ is the sigmoid function. The final updates for our model then become

where the diagonal scaling matrices $(C_{1},C_{2})$ and scalar coefficients $(c_{3},c_{4})$ are given by

with $Q^{(t)}=\gamma|\mathcal{E}|-\text{tr}[({Y^{(t)}})^{\top}\tilde{L}^{-}Y^{(t)}]$, $I$ as an $n\times n$ identity matrix, and $\frac{\alpha}{2}$ as the step size.

From the above expressions, we observe that the first and second terms of (IV-C) can be viewed as rescaled skip connections from the previous layer and input layer/base model, respectively. As we will later show, these scale factors are designed to facilitate guaranteed descent of the objective from (IV-C). Meanwhile, the third term of (IV-C) represents a typical GNN graph propagation layer while the fourth term is the analogous negative sampling propagation unique to our model. In the context of Chinese philosophy, the latter can be viewed as the Yin to the Yang which is the third term, and with a trade-off parameter that can be learned when training the higher-level objective from (1), the Yin/Yang balance can in a loose sense be estimated from the data; hence the name YinYanGNN for our proposed approach.

We illustrate key aspects of YinYanGNN in Figure 1 and the explicit computational steps are shown in Algorithm 1. And we emphasize that although the derivations of YinYanGNN may be somewhat complex, the algorithm that results is quite simple and efficient.

YinYanGNN Training. Upon inspection of Algorithm 1, we observe that lines 3 to 9 depict the training process for each epoch. Here we first sample the negative graph on line 3 which is used (unique to our model) on lines 4 to 6 to produce the YinYanGNN encoder node embeddings $Y^{(T)}$; this completes one model forward pass. Next, on line 7 we again sample negative edges as needed by the supervised training loss. Given these negative samples, along with predictive probabilities obtained via node embeddings from $Y^{(T)}$ passed through the HadamardMLP decoder, we compute the link prediction training loss defined by (1) on line 8. Finally, we update model parameters $\{W,\theta\}$ (and optionally $\lambda_{K}$) by standard backpropagation on line 9; this executes one model backward pass (in our experiments we use the Adam optimizer).

YinYanGNN Inference. For inference, we basically follow one forward pass of the training process from above to obtain encoder node embeddings $Y^{(T)}$. These are then applied to the HadamardMLP decoder to produce link prediction scores or edge probability estimates analogous to line 8, but without the loss computation.

YinYanGNN has a time complexity given by $O(T|\mathcal{E}|Kd+nPd^{2})$ for one forward pass, where as before $n$ is the number of nodes, $T$ is the number of propagation layers/iterations, $P$ is the number of MLP layers in the base model $f(X;W)$, and $d$ is the hidden embedding size. Notably, this complexity is of the same order as a vanilla GCN model, one of the most common GNN architectures [7], which has a forward-pass complexity of $O(T(|\mathcal{E}|d+nd^{2}))$.

We now drill down further into the details of overall inference speed. We denote the set of source nodes for test-time link prediction as $\mathcal{V}_{src}$, and for each node we examine all the other nodes in the graph, which means we have to calculate roughly $|\mathcal{V}_{src}|n$ edges. We compare YinYanGNN’s time with SEAL [9] and a recently proposed fast baseline BUDDY [12] in Table II. Like other node-wise methods, we split the inference time of our model into two parts: computing embeddings and decoding (for SEAL and BUDDY they are implicitly combined). The node embedding computation can be done only once so it does not depend on $\mathcal{V}_{src}$. Our decoding process uses HadamardMLP to compute scores for each destination node (which can also be viewed as being for each edge) and get the top nodes (edges). From this it is straightforward to see that the decoding time dominates the overall computation of embeddings. So for the combined inference time, SEAL is the slowest because of the factor $|\mathcal{E}|n$ while BUDDY and our model are both linear in the graph node number $n$ independently of $|\mathcal{E}|$. However, BUDDY has a larger factor including the compution of subgraph structure $b^{2}h$, which will be much slower as our experiments will show. Moreover, unlike SEAL or BUDDY, we can apply Flashlight [16] to node-wise methods like YinYanGNN, an accelerated decoding method based on maximum inner product search (MIPS) that allows HadamardMLP to have sublinear decoding complexity in $n$. See Section VI-C for related experiments and discussion.

Convergence criteria for the energies from (3) and (IV-C) are as follows (see Appendix X-B for proofs):

Previous work [49, 50] has demonstrated how randomly dropping edges of two isomorphic graphs can produce non-isomorphic, distinguishable subgraphs, and in so doing, enhance GNN expressiveness. Our incorporation of negative samples within the YinYanGNN encoder/forward pass can loosely be viewed in similar terms, but with a novel twist: Instead of randomly dropping edges in the original graph, we randomly add typed negative edges that are likewise capable of breaking otherwise indistinguishable symmetries.

Figure 2 serves to illustrate how the inclusion of negative samples within the forward pass of our model can potentially increase the expressiveness beyond traditional node-wise embedding approaches. As observed in the figure, $v_{2}$ and $v_{3}$ are isophormic nodes in the original graph (solid lines). However, when negative samples/edges are included the isomorphism no longer holds, meaning that link $(v_{1},v_{2})$ and link $(v_{1},v_{3})$ can be distinguished by a node-wise embedding method even without unique discriminating input features. Moreover, when combined with the flexibility of learning to balance multiple negative sampling graphs as in (10) through the trainable weights $\{\lambda_{K}^{k}\}$, the expressiveness of YinYanGNN becomes strictly greater than or equal to a vanilla nodewise embedding method (with equivalent capacity) that has no explicit access to the potential symmetry-breaking influence of negative samples.

Critically though, these negative samples are not arbitrarily inserted into our modeling framework. Rather, they emerge by taking gradient steps (12) over a principled regularization factor (within (10)) designed to push the embeddings of nodes sharing a negative edge apart during the forward pass. In a Section VII-E ablation we compare this unique YinYanGNN aspect with the alternative strategy of using random node features for breaking isomorphisms.

We conclude here by remarking that the goal of YinYanGNN, and the supporting analysis of this section, is not to explicitly match the expressiveness of edge-wise link prediction methods; even in principle this is not feasible given the additional information edge-wise methods have access to. Instead, the guiding principle of YinYanGNN is more aptly distilled as follows: conditioned on the limitations of a purely node-wise link prediction architecture, to what extent can we improve model expressiveness (and ultimately accuracy) without compromising the scalability which makes node-wise methods attractive in the first place.

Accuracy results are displayed in Table III, where we observe that YinYanGNN achieves the best performance on 6 out of 7 datasets (while remaining competitive across all 7) even when compared against more time-consuming or inference-inefficient edge-wise methods. And we outperform node-wise methods by a large margin (including several others shown in Appendix X-A), demonstrating that YinYanGNN can achieve outstanding predictive accuracy without sacrificing efficiency. Similarly, as included in the Appendix X-A, YinYanGNN also outperforms a variety of non-GNN link prediction baselines.

Inference Time. We next present comparisons in terms of inference speed, which is often the key factor determining whether or not a model can be deployed in real-world scenarios. For example, in an online system, providing real-time recommendations may require quickly evaluating a large number of candidate links. Figure 3 reports the results, again relative to both edge- and node-wise baseline models. Noting the log-scale time axis, from these results we observe that YinYanGNN is significantly faster than all the edge-wise models, and nearly identical to the fellow node-wise approaches as expected. And for the fastest edge-wise model, Neo-GNN, YinYanGNN is still simultaneously more efficient (Figure 3) and also much more accurate (Table III). Additional related details and inference-time results are deferred to Appendix X-A We also remark that, as a node-wise model, the efficiency of YinYanGNN can be further improved to sublinear complexity using Flashlight [16]; however, at the time of submission public Flashlight code was not available, so we differ consideration to future work.

Decoding Time. Besides full inference times, in practice, we can save the embeddings output from the encoder and only compare the decoding step, which is a key differentiator. To this end, we compare the decoding speed of YinYanGNN with SEAL, a highly-influential edge-wise model, and BUDDY, which represents a strong baseline with a good balance between accuracy and speed for edge-wise GNNs. We conduct the experiments according to [16]. Table IV displays the results, where our model achieves the fastest performance compared with these two baselines. We also note that although the time complexity of BUDDY is linear in $n$ (which is the same as our model without Flashlight), the different scaling coefficients caused by the required subgraph information calculation still make BUDDY much slower than YinYanGNN.

Finally, although the code for Flashlight is not yet publicly available, we can mimic the acceleration it will produce as follows. As Flashlight is a repeated version of MIPS, which can be applied to dot product prediction, we also run MIPS on our model with dot product as the decoder to see how fast MIPS can accelerate. We run the experiments on CPU and achieve a 20x speedup (compare last two columns of Table IV with and without MIPS).

Training Time. We note that as is often the case for link prediction, our focus is on efficient inference, since the number of candidate node pairs grows quadratically in the graph size. That being said, we nonetheless provide a representative comparison here of training times. Specifically, we measure the training times on OGBL-PPA using identical training sets and platforms. BUDDY requires approximately 3 hours, whereas our YinYanGNN model accomplishes the same training task in approximately 12.5 minutes, underscoring a substantial discrepancy in training efficiency.

General Comments. After the edge-wise SEAL algorithm was originally published, not many new scalable node-wise alternatives have been proposed, likely because it has been difficult to achieve competitive performance. For example, there are actually very few node-wise entries atop the OGB link prediction leaderboard; almost all are less efficient edge-wise methods or related.

In fact, over the course of preparing this submission, there was no model of any kind with OGB leaderboard performance above ours on all datasets, and only a single published node-wise model on the leaderboard above ours on any dataset: this is model CFLP w/ JKNet [60], and only when applied to DDI. But when we examine the CFLP paper, we find that their performance is much lower than YinYanGNN on other non-OGB datasets as shown in Table V.

Neural Common Neighbors (NCN). Another recent approach with strong leaderboard performance worth mentioning here is NCN [61], which relies on a node-wise encoder as with YinYanGNN. However, the key feature of NCN is a decoder based on collecting common neighbors of nodes spanning each candidate link, which leads to a significantly higher inference cost. Even so, the NCN decoder itself can be easily be integrated with any node-wise encoder, including YinYanGNN, to potentially increase accuracy, albeit with a much higher inference time. Results on OGB-Citation2 (the largest OGB link prediction task) are listed in Table VI. These results demonstrate that YinYanGNN with an NCN decoder can actually outperform both the original NCN MRR from [61] as well as all of the edge-wise models in Table III. The main downside is that inference time increases by $\sim 40\times$ (see Table VI) over YinYanGNN in its original form. This further highlights the challenge of achieving the highest accuracy while preserving the full efficiency of purely node-wise models.

PLNLP and GIDN. Beyond the above, we conclude by commenting on two additional leaderboard methods. First, there is an additional, unpublished node-wise model on the leaderboard called PLNLP [42] that does lead to strong accuracy on both Collab and DDI (two relatively small datasets), but the main contribution is a new training loss that is orthogonal to our method. In fact any node-wise model, including ours, could be trained with their loss to improve performance. The only downside is that the PLNLP approach seems to be dataset dependent, with lesser performance on other benchmarks, e.g., PLNLP only achieves 32.38 Hits@100 on PPA and 84.92 MRR on Citation2. Additionally, the GIDN model [62] follows the same loss as PLNLP and again achieves good results restricted to the smaller-scale Collab and DDI benchmarks (no other results are shown). However, the companion GIDN paper [62] is unpublished and extremely brief, with just a couple short paragraphs explaining the method and no assessment of inference time complexity. Hence we are unsure of how this model compares against node-wise models when applied to large-scale datasets as is our focus.

As the integration of both positive and negative samples within a unified node-wise embedding framework is a critical component of our model, in Table VII we report results both with and without the inclusion of the negative sampling penalty in our lower-level embedding model from (IV-C). Clearly the former displays notably superior performance as expected.

One benefit of learning $\{\lambda_{K}^{k}\}$ is that it allows us to better match the performance of larger $K$ with fixed $\lambda_{K}$ using a smaller $K$ but with learnable $\{\lambda_{K}^{k}\}$. The latter is more practical given the greater efficiency of a smaller $K$. Figures 5 and 5 demonstrate this phenomena, whereby a learnable $\{\lambda_{K}^{k}\}$ achieves better accuracy for smaller values of $K$ on Cora and Pubmed datasets (for larger $K$ performance is similar).

For bigger datasets where using larger $K$ can be prohibitively expensive, this capability can be exploited to achieve higher accuracy. For example, on PPA and Citation2 a larger $K$ may well improve accuracy, but this is not computationally cheap without sampling or other approximations. In contrast, we can improve performance with small $K$ just by learning $\{\lambda_{K}^{k}\}$ as shown in Table VIII. (In contrast, for dense graphs a larger $K$ is less likely to be helpful, in part because the extra edges will make the negative graph much more dense and more similar to complete graph, which results in the loss of structural information. Not surprisingly then, we found that learning $\{\lambda_{K}^{k}\}$ did not improve performance on the dense graph DDI dataset where $K=1$ was optimal.)

As shown in Table IX, the YinYanGNN architecture is tolerant to the use of different negative samplers, from sampling negative edges uniformly in the graph (namely global uniform sampler) to sampling negative edges by fixing the source node and sampling negative destination nodes (namely source uniform sampler).

As we discussed in Section IV, heuristically embedding negative samples within an existing GNN architecture is another possible strategy. But in this way the negative sampling terms are not integrated by a principled energy function. We follow the analysis of negative sampling in energy function, and modify the GCN architecture to incorporate negative sampling in the forward model as follows

where $W^{(t)}$ are the layer-wise weights and $K$ is the negative edge number for each positive edge. Also, $\lambda_{K}$ is the hyperparameter to control the impact of negative edges as before. As shown in Table X, negative sampling is also useful in some datasets for GCNs, but the performance is still not as good compared to YinYanGNN, where the negative samples are anchored to the proposed energy function.

In Section V-C, we show that random edges can in principle break the stated isomorphism. However, another potential way to accomplish this effect is to use random features that disambiguate each node. However, this can introduce undesirable auxilary effects, which become apparent when we analyze the resulting energy function when applying random features instead of negative sampling.

When we add random features to (3) we obtain the modified form

For simplicity of illustration, we consider an $f$ that obeys the distributive property, i.e., We then compute the gradient

From these expressions we observe that random features merely introduce random blurriness to the gradient updates; they do not emerge from an otherwise useful regularization factor. In contrast, with YinYanGNN, the negative samples create gradients that push the node embeddings of nodes that do not share a true edge apart, a useful form of structural regularization that is independent of any mechanism to break isomorphisms per se.

We support these points with additional experiments in Table XI, where random features are introduced and compared against our original YinYanGNN model. From these results we observe that random features can actually degrade performance and introduce instabilities as evidenced by the reduced accuracy and larger standard deviations.

And finally, in Table XII, we examine how YinYanGNN performance is impacted as we vary the relative proportion assigned to negative graphs, which amounts to the relative weighting of $\lambda$ to $\lambda_{K}$. Specifically, we choose $\lambda=1$ and vary $\lambda_{K}$ from $0.1$ to $10$ (in each case $\lambda_{K}$ is fixed, not learnable) and train YinYanGNN on small datasets Cora, Citeseer, and Pubmed as well as large dataset Citation2. Performance is generally best with $\lambda_{K}\in[0.1,1]$ across all datasets, with the smaller datasets favoring $\lambda_{K}=1$ and the largest $\lambda_{K}=0.1$.

In this section we show further details of inference time results, and compare YinYanGNN with more diverse classical baselines, expressive node-wise GNNs, and distillation methods.