A Scalable Pretraining Framework for Link Prediction with Efficient Adaptation

2025-08-06

AI Summary

Link prediction faces challenges including label sparsity, sensitivity to initialization, and poor generalization under distributional shift. To address these, we propose the first scalable pre-training framework specifically designed for link prediction. First, we systematically analyze the transferability of node-level and edge-level modules and design a late-fusion strategy to jointly model pairwise relational patterns. Second, we incorporate a Mixture-of-Experts (MoE) architecture to adaptively capture heterogeneous patterns across diverse pre-training sources, thereby mitigating negative transfer. Third, we introduce a parameter-efficient fine-tuning method that achieves rapid adaptation to new datasets with minimal parameter updates. Evaluated on 16 cross-domain benchmarks, our approach achieves state-of-the-art performance in low-resource settings, reduces adaptation overhead by over four orders of magnitude compared to full fine-tuning, while matching its accuracy.

Technology Category

• Machine Learning: Mixture of Experts (MoE)
• Data Mining & Knowledge Management: Linked Open Data, Knowledge Graphs & KB Completion
• Natural Language Processing: Language Grounding & Multi-modal NLP

Application Category

• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Web Mining and Content Analysis: Large pretrained models with web data
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

Abstract

Link Prediction (LP) is a critical task in graph machine learning. While Graph Neural Networks (GNNs) have significantly advanced LP performance recently, existing methods face key challenges including limited supervision from sparse connectivity, sensitivity to initialization, and poor generalization under distribution shifts. We explore pretraining as a solution to address these challenges. Unlike node classification, LP is inherently a pairwise task, which requires the integration of both node- and edge-level information. In this work, we present the first systematic study on the transferability of these distinct modules and propose a late fusion strategy to effectively combine their outputs for improved performance. To handle the diversity of pretraining data and avoid negative transfer, we introduce a Mixture-of-Experts (MoE) framework that captures distinct patterns in separate experts, facilitating seamless application of the pretrained model on diverse downstream datasets. For fast adaptation, we develop a parameter-efficient tuning strategy that allows the pretrained model to adapt to unseen datasets with minimal computational overhead. Experiments on 16 datasets across two domains demonstrate the effectiveness of our approach, achieving state-of-the-art performance on low-resource link prediction while obtaining competitive results compared to end-to-end trained methods, with over 10,000x lower computational overhead.

Problem

• Addressing limited supervision in link prediction
• Improving generalization under distribution shifts
• Reducing computational overhead for adaptation

Innovation

• Late fusion strategy for node-edge integration
• Mixture-of-Experts framework for diverse pretraining
• Parameter-efficient tuning for fast adaptation