Self-supervised Quantized Representation for Seamlessly Integrating Knowledge Graphs with Large Language Models

TL;DR

This paper tackles the challenge of pairing Knowledge Graphs with Large Language Models by learning self-supervised, discrete, language-aligned codes for each entity. It introduces SSQR, a two-stage framework that (1) encodes KG structure with a GCN and compresses it into quantized codes via a VQ-style codebook while aligning semantics through LLM embeddings, and (2) uses these codes as direct input features to tune LLMs for KG tasks. The authors demonstrate that SSQR achieves state-of-the-art or competitive performance on KG link prediction and triple classification with only 16 tokens per entity, significantly reducing prompting overhead. The work highlights the potential of discrete, language-native KG representations to enable seamless, scalable KG reasoning inside LLMs, while discussing limitations such as computational cost and generalization across diverse KGs.

Abstract

Due to the presence of the natural gap between Knowledge Graph (KG) structures and the natural language, the effective integration of holistic structural information of KGs with Large Language Models (LLMs) has emerged as a significant question. To this end, we propose a two-stage framework to learn and apply quantized codes for each entity, aiming for the seamless integration of KGs with LLMs. Firstly, a self-supervised quantized representation (SSQR) method is proposed to compress both KG structural and semantic knowledge into discrete codes (\ie, tokens) that align the format of language sentences. We further design KG instruction-following data by viewing these learned codes as features to directly input to LLMs, thereby achieving seamless integration. The experiment results demonstrate that SSQR outperforms existing unsupervised quantized methods, producing more distinguishable codes. Further, the fine-tuned LLaMA2 and LLaMA3.1 also have superior performance on KG link prediction and triple classification tasks, utilizing only 16 tokens per entity instead of thousands in conventional prompting methods.

Figures (13)

• Figure 1: Illustration of different strategies to integrate KGs with LLMs. (a) The direct method utilizes (sampled) graph structures and semantic text as inputs. (b) Our method for seamlessly integrating KGs with LLMs using learned quantized and discrete codes.
• Figure 2: The statistics of 2-hop sampled neighbors and needed tokens (by LLaMA2) for entities in FB15k-237.
• Figure 3: The overall architecture of our study. (a) is for SSQR learning. (b) is for instruction tuning for KG tasks, where the learned quantized representations serve as features. Icons and represent the status of the module during training, indicating if it is frozen or being updated, respectively.
• Figure 4: The cosine similarity of quantized representations on the WN18RR dataset (sampled 8 entities).
• Figure 5: The effects of codebook length ($M$) and sequence length ($N$) for each entity.