Efficient Link Prediction via GNN Layers Induced by Negative Sampling
Yuxin Wang, Xiannian Hu, Quan Gan, Xuanjing Huang, Xipeng Qiu, David Wipf
TL;DR
This work introduces YinYanGNN, a novel node-wise GNN for link prediction that brings negative sampling into the forward pass by optimizing an energy function balancing positive and negative edges. The model computes end-to-end trainable node embeddings via energy minimization, preserving node-wise efficiency while achieving accuracy competitive with edge-wise methods. Key contributions include an energy-based formulation with normalization and multiple negative graphs, convergence guarantees, and extensive experiments demonstrating strong performance and favorable inference times on standard benchmarks. The approach offers a practical middle ground between expressive edge-wise models and fast node-wise predictors, with potential for further acceleration via specialized decoding tech like Flashlight.
Abstract
Graph neural networks (GNNs) for link prediction can loosely be divided into two broad categories. First, \emph{node-wise} architectures pre-compute individual embeddings for each node that are later combined by a simple decoder to make predictions. While extremely efficient at inference time, model expressiveness is limited such that isomorphic nodes contributing to candidate edges may not be distinguishable, compromising accuracy. In contrast, \emph{edge-wise} methods rely on the formation of edge-specific subgraph embeddings to enrich the representation of pair-wise relationships, disambiguating isomorphic nodes to improve accuracy, but with increased model complexity. To better navigate this trade-off, we propose a novel GNN architecture whereby the \emph{forward pass} explicitly depends on \emph{both} positive (as is typical) and negative (unique to our approach) edges to inform more flexible, yet still cheap node-wise embeddings. This is achieved by recasting the embeddings themselves as minimizers of a forward-pass-specific energy function that favors separation of positive and negative samples. Notably, this energy is distinct from the actual training loss shared by most existing link prediction models, where contrastive pairs only influence the \textit{backward pass}. As demonstrated by extensive empirical evaluations, the resulting architecture retains the inference speed of node-wise models, while producing competitive accuracy with edge-wise alternatives. We released our code at https://github.com/yxzwang/SubmissionverOfYinYanGNN.