GRAIL: Guided Retrieval via Layer-Wise Discrepancy from Extrapolated Final Distributions
摘要
Large Language Models (LLMs) have shown strong performance in various NLP tasks, yet their reliance on internal knowledge makes them prone to hallucinations, especially in complex or knowledge-intensive scenarios. Retrieval-Augmented Generation (RAG) introduces external knowledge to address this issue. However, most existing RAG methods rely on static retrieval strategies that overlook the model’s internal reasoning state, often leading to redundant or noisy content that degrades generation quality. To address the critical challenge of determining when retrieval should be triggered during generation, we propose GRAIL-a Guided Retrieval method via Layer-wise Discrepancy from Extrapolated Final Distributions. GRAIL models the evolution of high-layer probability distributions to construct a stable predictive target: the Extrapolated-Enhanced Final Distribution (EEFD). By comparing intermediate-layer predictions with the EEFD, the model identifies tokens that exhibit significant distributional deviation at later stages of inference-suggesting semantic uncertainty or incomplete convergence. GRAIL dynamically triggers retrieval at such points to supplement external knowledge. Experimental results on multiple single-hop and multi-hop QA datasets demonstrate that GRAIL consistently delivers strong or competitive performance, confirming its effectiveness.