Subgraph-Aware Training of Language Models for Knowledge Graph Completion Using Structure-Aware Contrastive Learning

TL;DR

This work tackles the limitation of text-only PLM fine-tuning for knowledge graph completion by injecting the knowledge graph's structural inductive bias into training. It introduces SATKGC, which combines subgraph-aware sampling (BRWR-based) and a subgraph-as-mini-batch (SaaM) with a proximity-aware contrastive objective and a frequency-aware mini-batch loss to handle long-tail distributions. The approach yields state-of-the-art results on WN18RR, FB15k-237, and Wikidata5M across transductive and inductive settings, demonstrating robust gains by leveraging KG topology. The proposed SaaM framework is encoder-agnostic and generalizable, offering a practical path to integrate structural information into PLM-based KGC systems.

Abstract

Fine-tuning pre-trained language models (PLMs) has recently shown a potential to improve knowledge graph completion (KGC). However, most PLM-based methods focus solely on encoding textual information, neglecting the long-tailed nature of knowledge graphs and their various topological structures, e.g., subgraphs, shortest paths, and degrees. We claim that this is a major obstacle to achieving higher accuracy of PLMs for KGC. To this end, we propose a Subgraph-Aware Training framework for KGC (SATKGC) with two ideas: (i) subgraph-aware mini-batching to encourage hard negative sampling and to mitigate an imbalance in the frequency of entity occurrences during training, and (ii) new contrastive learning to focus more on harder in-batch negative triples and harder positive triples in terms of the structural properties of the knowledge graph. To the best of our knowledge, this is the first study to comprehensively incorporate the structural inductive bias of the knowledge graph into fine-tuning PLMs. Extensive experiments on three KGC benchmarks demonstrate the superiority of SATKGC. Our code is available.

Figures (10)

• Figure 1: False positive (FP) ratio against the distance (i.e., length of the shortest path) between head and tail of a FP triple in KG across different text-based methods.
• Figure 2: False positive (FP) ratio against the degree of tail for a FP triple across different text-based methods.
• Figure 3: Overview of the proposed training framework, which consists of: (i) Random-walk Based Subgraph Sampling (before training); (ii) Subgraph as a Mini-batch; (iii) Proximity-aware Contrastive Learning; (iv) Frequency-aware Mini-batch Training.
• Figure 4: Example of BRWR-based subgraph sampling; (a) probability of selecting start entity $s$ between $h$ and $t$ of a center triple, where $t$ with a lower degree is more likely to be $s$ than $h$; (b) probability of selecting a neighbor of current entity $u$. A random walker is more likely to move to $v_1$ than to $v_2$ with its degree larger than $v_1$.
• Figure 5: Frequency distributions of entities for original KG, 100 mini-batches randomly sampled from $\mathcal{T}$, and those randomly sampled by SaaM on WN18RR and FB15k-237. The entities are sorted in the ascending order of their degrees.