Information Flow Reveals When to Trust Language Models
Abstract
In retrieval-augmented generation, language models can generate incorrect responses if they fail to utilize query-relevant content from the retrieved evidence. This shifts the focus of uncertainty quantification (UQ) toward assessing contextual grounding, i.e., whether predictions are supported by query-relevant tokens. Recent UQ methods unpack language models to characterize how inputs are processed. Nevertheless, these methods focus on a few layers and overlook the whole progressive propagation within the model, thereby failing to fully capture the grounding dynamics essential for reliable uncertainty estimation. We use information flow to build a layer-wise trace that reveals each context token’s contribution to the output, providing an interpretable basis for assessing reliability. From this analysis, we introduce two measures to calibrate prediction confidence. The first, \textit{simulatability}, posits that a prediction is more likely to be correct when context token contributions align closely with their true relevance. The second, \textit{concentration}, asserts that a response is more likely to be correct when it is derived from a narrow, focused subset of tokens. Experiments show that our method achieves an average AUROC of 0.709, exceeding the runner-up performance of 0.676, while maintaining moderate computational cost.
Lay Summary
Language models are increasingly used to answer questions, summarize information, and assist with decision-making. However, these systems can sometimes give convincing but incorrect answers, making it difficult for users to know when to trust them. Existing confidence scores mainly look at the model’s final answer, but they often ignore how the model actually used the provided information to produce that answer. In this work, we study how information moves through the layers of a language model while it generates a response. We introduce two new ways to track this process: one measures the order in which important pieces of information emerge, and the other measures how strongly each input token contributes to the final answer. By analyzing these information flows, our method can estimate whether the model’s response is truly grounded in the relevant context. Experiments on several question-answering benchmarks show that our approach provides more reliable confidence estimates than existing methods. This could help make future AI systems more transparent, trustworthy, and safer to use in real-world applications where reliability matters.