# Retrieval-Augmented Language Model for Extreme Multi-Label Knowledge Graph Link Prediction

arXiv: https://arxiv.org/abs/2405.12656 · HTML: https://mlsanity.com/paper/2405.12656/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Knowledge-graph and structured-knowledge RAG, one of 15 directions in Retrieval-augmented generation.
- Builds on 13 of its 17 references that are in the map.
- First posted on arXiv 21 May 2024.

## Placed in

- Retrieval-augmented generation / [Knowledge-graph and structured-knowledge RAG](https://mlsanity.com/rag/knowledge-graph-and-structured-knowledge-rag/)

## Builds on (most cited first, 13 of 13)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [SimCSE: Simple Contrastive Learning of Sentence Embeddings](https://mlsanity.com/paper/2104.08821/) (2021)
- [Atlas: Few-shot Learning with Retrieval Augmented Language Models](https://mlsanity.com/paper/2208.03299/) (2022)
- [On the Origin of Hallucinations in Conversational Models: Is it the Datasets or the Models?](https://mlsanity.com/paper/2204.07931/) (2022)
- [Knowledge-Augmented Language Model Prompting for Zero-Shot Knowledge Graph Question Answering](https://mlsanity.com/paper/2306.04136/) (2023)
- [Neural Path Hunter: Reducing Hallucination in Dialogue Systems via Path Grounding](https://mlsanity.com/paper/2104.08455/) (2021)
- [AttentionXML: Label Tree-based Attention-Aware Deep Model for High-Performance Extreme Multi-Label Text Classification](https://mlsanity.com/paper/1811.01727/) (2018)
- [LightXML: Transformer with Dynamic Negative Sampling for High-Performance Extreme Multi-label Text Classification](https://mlsanity.com/paper/2101.03305/) (2021)
- [RHO (ρ): Reducing Hallucination in Open-domain Dialogues with Knowledge Grounding](https://mlsanity.com/paper/2212.01588/) (2022)
- [RAVEN: In-Context Learning with Retrieval-Augmented Encoder-Decoder Language Models](https://mlsanity.com/paper/2308.07922/) (2023)
- [Multi-task Pre-training Language Model for Semantic Network Completion](https://mlsanity.com/paper/2201.04843/) (2022)
- [Born for Auto-Tagging: Faster and better with new objective functions](https://mlsanity.com/paper/2206.07264/) (2022)

## Abstract (arXiv)

Extrapolation in Large language models (LLMs) for open-ended inquiry encounters two pivotal issues: (1) hallucination and (2) expensive training costs. These issues present challenges for LLMs in specialized domains and personalized data, requiring truthful responses and low fine-tuning costs. Existing works attempt to tackle the problem by augmenting the input of a smaller language model with information from a knowledge graph (KG). However, they have two limitations: (1) failing to extract relevant information from a large one-hop neighborhood in KG and (2) applying the same augmentation strategy for KGs with different characteristics that may result in low performance. Moreover, open-ended inquiry typically yields multiple responses, further complicating extrapolation. We propose a new task, the extreme multi-label KG link prediction task, to enable a model to perform extrapolation with multiple responses using structured real-world knowledge. Our retriever identifies relevant one-hop neighbors by considering entity, relation, and textual data together. Our experiments demonstrate that (1) KGs with different characteristics require different augmenting strategies, and (2) augmenting the language model's input with textual data improves task performance significantly. By incorporating the retrieval-augmented framework with KG, our framework, with a small parameter size, is able to extrapolate based on a given KG. The code can be obtained on GitHub: https://github.com/exiled1143/Retrieval-Augmented-Language-Model-for-Multi-Label-Knowledge-Graph-Link-Prediction.git