Mitigating Hallucinations in Large Language Models via Self-Refinement-Enhanced Knowledge Retrieval

TL;DR

The paper tackles hallucinations in medical QA by introducing Re-KGR, a post-generation framework that identifies high-risk tokens via token-level uncertainty, extracts and refines factual triples, retrieves corroborating KG evidence, and rectifies content to boost truthfulness. By combining entropy and layer-divergence signals with selective triple extraction, it reduces the number of necessary KG lookups while maintaining or improving factual accuracy. Evaluations on MedQuAD with LLaMA-7b and DoLa show notable improvements in truthfulness (up to 0.610 on GPT-4 scoring) and substantial reductions in retrieval time (63–75%), demonstrating practical efficiency gains. The approach emphasizes minimal retraining, domain-specific KG integration, and potential generalization to other high-stakes domains where reliable external knowledge is available.

Abstract

Large language models (LLMs) have demonstrated remarkable capabilities across various domains, although their susceptibility to hallucination poses significant challenges for their deployment in critical areas such as healthcare. To address this issue, retrieving relevant facts from knowledge graphs (KGs) is considered a promising method. Existing KG-augmented approaches tend to be resource-intensive, requiring multiple rounds of retrieval and verification for each factoid, which impedes their application in real-world scenarios. In this study, we propose Self-Refinement-Enhanced Knowledge Graph Retrieval (Re-KGR) to augment the factuality of LLMs' responses with less retrieval efforts in the medical field. Our approach leverages the attribution of next-token predictive probability distributions across different tokens, and various model layers to primarily identify tokens with a high potential for hallucination, reducing verification rounds by refining knowledge triples associated with these tokens. Moreover, we rectify inaccurate content using retrieved knowledge in the post-processing stage, which improves the truthfulness of generated responses. Experimental results on a medical dataset demonstrate that our approach can enhance the factual capability of LLMs across various foundational models as evidenced by the highest scores on truthfulness.

Figures (3)

• Figure 1: Illustration of the proposed method, which consists of four components: (1) Entity Detection; (2) Triple Extraction; (3) Knowledge Retrieval; (4) Knowledge Verification & Rectification. Given an input question, the LLM generates an intermediate response along with relevant token logits, both of which are leveraged to identify critical word entities and subsequently extract a refined collection of factual statements, facilitating the KG retrieval and verification process. The lower-left subplot shows the mechanism of Entity Detection, which operates with the entropy of the next-token predictive probability distribution. This component dynamically identifies thresholds for anomaly detection on generated tokens for each sentence, contributing to the extraction of critical entities. The lower-right subplot visualizes a subgraph in the KG and presents an example query used for knowledge retrieval.
• Figure 2: Example of Triple Extraction Process. LLMs are prompted to extract factually related statements in the form of knowledge triples, with identified entities being utilized to select critical triples.
• Figure 3: Illustration of entity detection results across various criteria and corresponding outcomes of self-refinement triples. The upper figure displays the number of extracted entities, refined triples, and original triples for Re-KGR method using the LLaMA-7b model, while the lower one presents results for the method with the DoLa model. Grey dashed lines indicate the number of extracted entities and refined triples based on the entropy criterion $c_e$, respectively.