LinkGPT: Teaching Large Language Models To Predict Missing Links

TL;DR

LinkGPT tackles link prediction on text-attributed graphs by teaching LLMs to reason over both node-level and pairwise structural information. It introduces a three-component framework: neighborhood-aware node encoding paired with a dedicated pairwise encoder, a two-stage instruction-tuning regime, and a retrieval-rerank inference pipeline that uses BM25 and a Yes/No indexing to efficiently rank candidates. Across four real-world TAG datasets, LinkGPT achieves state-of-the-art LP performance and demonstrates strong zero-shot and few-shot generalization, while delivering about a $10\times$ reduction in inference time through the retrieval scheme. This work substantiates the practicality of end-to-end LLM-based LP on structured text data and offers a scalable approach for cross-domain graph reasoning in real-world applications.

Abstract

Large Language Models (LLMs) have shown promising results on various language and vision tasks. Recently, there has been growing interest in applying LLMs to graph-based tasks, particularly on Text-Attributed Graphs (TAGs). However, most studies have focused on node classification, while the use of LLMs for link prediction (LP) remains understudied. In this work, we propose a new task on LLMs, where the objective is to leverage LLMs to predict missing links between nodes in a graph. This task evaluates an LLM's ability to reason over structured data and infer new facts based on learned patterns. This new task poses two key challenges: (1) How to effectively integrate pairwise structural information into the LLMs, which is known to be crucial for LP performance, and (2) how to solve the computational bottleneck when teaching LLMs to perform LP. To address these challenges, we propose LinkGPT, the first end-to-end trained LLM for LP tasks. To effectively enhance the LLM's ability to understand the underlying structure, we design a two-stage instruction tuning approach where the first stage fine-tunes the pairwise encoder, projector, and node projector, and the second stage further fine-tunes the LLMs to predict links. To address the efficiency challenges at inference time, we introduce a retrieval-reranking scheme. Experiments show that LinkGPT can achieve state-of-the-art performance on real-world graphs as well as superior generalization in zero-shot and few-shot learning, surpassing existing benchmarks. At inference time, it can achieve $10\times$ speedup while maintaining high LP accuracy.

Figures (7)

• Figure 1: Overview of the LinkGPT framework. The framework consists of three main components: (1) node encoding and pairwise encoding, (2) two-stage instruction tuning, and (3) a retrieval-rerank scheme for inference. In this figure, $\Rightarrow$means that the corresponding module is tuned during stage 1 and is frozen during stage 2, and vice versa. Note that the two LLMs in this figure are the same model.
• Figure 2: Instruction templates for link prediction and neighbor prediction tasks. Only the answers (highlighted in green) are used for calculating the loss.
• Figure 3: Performance of few-shot in-context learning. (Left): When the model is fine-tuned and evaluated on MAE-Geology, in-context learning does not significantly improve the performance. (Right): For cross-domain generalization, in-context learning boosts the performance significantly.
• Figure 4: Performance of the retrieval-rerank scheme of LinkGPT. The line chart represents the final MRR, using the right axis, while the bar chart represents the total time required for each $(s, \mathcal{C})$ pair during inference, using the left axis. Here $s$ denotes the source node and $\mathcal{C}$ denotes the candidate target node set.
• Figure 5: Visualization of two node encoding methods.