Contextualization Distillation from Large Language Model for Knowledge Graph Completion (2024)

Abstract: While textual information significantly enhances the performance of pre-trained language models (PLMs) in knowledge graph completion (KGC), the static and noisy nature of existing corpora collected from Wikipedia articles or synsets definitions often limits the potential of PLM-based KGC models. To surmount these challenges, we introduce the Contextualization Distillation strategy, a versatile plug-in-and-play approach compatible with both discriminative and generative KGC frameworks. Our method begins by instructing large language models (LLMs) to transform compact, structural triplets into context-rich segments. Subsequently, we introduce two tailored auxiliary tasks—reconstruction and contextualization—allowing smaller KGC models to assimilate insights from these enriched triplets. Comprehensive evaluations across diverse datasets and KGC techniques highlight the efficacy and adaptability of our approach, revealing consistent performance enhancements irrespective of underlying pipelines or architectures. Moreover, our analysis makes our method more explainable and provides insight into how to generate high-quality corpora for KGC, as well as the selection of suitable distillation tasks.

## introduction knowledge graph completion (kgc) is a fundamental task in natural language processing (nlp), aiming at unveiling hidden insights within diverse knowledge graphs to explore novel knowledge patterns. traditional kgc methods @xcite typically predict the missing part of the triplets by learning the representation of each entity and relation based on their structural information. however, such embeddingbased methods tend to overlook the rich textual in-methods h@1 h@3 h@8/10 chatgpt-1-shot 15.6 17.6 19.6 palm2-1-shot 15.7 20.8 25.4 kg-s2s @xcite 28.5 38.8 49. 3 j. g. ballard shanghai place_of_birth problem: succinct text: j.g. ballard was a novelist problem: static text: shanghai is a 2010 american mystery/thriller neonoir film directed by mikael håfström, ... problem: noisy text: in 1984, j.g. ballard won broad, ... british boy during the japanese occup ation of shanghai while it has been well-discovered that textual information can be beneficial for plm-based kgc models @xcite @xcite @xcite , prior attempts to augment kgc models with textual data from wikipedia article @xcite or synsets definitions @xcite have encountered certain limitations: (i) entity descriptions, often succinct and static, may inhibit the formation of a comprehensive understanding of entities within kgc models. (ii) the incorporation of triplet descriptions, albeit potentially enriching, can introduce substantial noise, particularly when derived through automatic entity alignment @xcite . figure 1 demonstrates an example to illustrate the aforementioned limitations. the description for the head "j. g. ballard" is limited and for the tail "shanghai", it mistakenly uses the definition of the movie also named "shanghai". also, while the two entities show up in the triplet description, it falls short in conveying the semantic essence of the relation "place_of_birth". in light of these limitations, our attention shifts to large language models (llms) @xcite @xcite , renowned for their capability in generating articulate and high-quality data @xcite @xcite . our exploration commences with a scrupulous evaluation of llms, such as chatgpt and palm2, in kgc, benchmarking them across several esteemed kgc datasets @xcite @xcite . utilizing 1-shot in-context learning (icl), we deduce missing heads or tails in triplets and report evaluation metrics. it reveals a significant performance discrepancy of two llms in comparison to kg-s2s @xcite despite its reliance on a smaller foundational model, t5-base @xcite . this insight propels us toward the conclusion that direct utilization of llms for kgc tasks, while intuitive, is outperformed by the fine-tuning of more diminutive, specialized kgc models. this observation aligns with findings from to optimally harness llms for kgc, we draw inspiration from recent works @xcite and introduce a novel approach, contextualization distillation. contextualization distillation first extracts descriptive contexts from llms with well-designed prompts, thereby securing dynamic, high-quality context for each entity and triplet. subsequent to this, two auxiliary tasks are proposed to train smaller kgc models with these informative, descriptive contexts. the plug-in-and-play characteristic of our contextualization distillation enables us to apply and evaluate it on various kgc datasets and baseline models. through extensive experiments, we affirm that contextualization distillation consistently enhances the performance of smaller kgc models, irrespective of architectural and pipeline disparities. additionally, we provide an exhaustive analysis of each step of contextualization distillation, encouraging further insights and elucidations. the contributions of this work can be summarized into three main aspects: ## contextualization distillation in this section, we first illustrate how we curate prompts to extract the descriptive context of each triplet from the llm. subsequently, we design a multi-task framework, together with two auxiliary tasks-reconstruction and contextualization-to train smaller kgc models with these high-quality context corpus. the overview pipeline of our method is illustrated in figure 2 . ## experiment in this section, we apply our contextualization distillation across a range of plm-based kgc baselines. we compare our enhanced model with our approach against the vanilla models using several kgc datasets. additionally, we do further analysis of each component in our contextualized distillation and make our method more explainable by conducting case studies. ## conclusion in this work, we propose contextualization distillation, addressing the limitation of the existing kgc textual data by prompting llms to generate descriptive context. to ensure the versatility of our approach across various plm-based kgc models, we have designed a multi-task learning framework. within this framework, we incorporate two auxiliary tasks, reconstruction and contextualization, which aid in training smaller kgc models in the informative descriptive context. we conduct experiments on several mainstream kgc benchmarks and the results show that our contextualization distillation consistently enhances the baseline model's performance. furthermore, we conduct in-depth analyses to make the effect of our method more explainable, providing guidance on how to effectively leverage llms to improve kgc as well. in the future, we plan to adapt our method to other knowledge-driven tasks, such as entity linking and knowledge graph question answering.