TPA: Next Token Probability Attribution for Detecting Hallucinations in RAG

AI summary

Overview

Research area: mechanistic interpretability applied to hallucination detection in Retrieval-Augmented Generation (RAG), within Natural Language Processing.

Technical level: Advanced. The paper works directly with transformer residual-stream internals (attention heads, FFN blocks, LayerNorm, unembedding matrices) and assumes familiarity with Pre-LN decoder architectures.

Scope: This paper introduces TPA (Next Token Probability Attribution), a framework that decomposes each generated token's probability into seven additive source contributions, aggregates them by Part-of-Speech tags, and uses the resulting features to classify hallucinated RAG responses across five LLMs.

What This Paper Is About

RAG systems ground language models in retrieved documents, but they can still produce responses that contradict the retrieved context. Prior mechanistic work, notably ReDeEP, explained such failures as a binary conflict in which internal parametric knowledge stored in FFNs overrides the retrieved context. TPA argues that this view is incomplete and instead attributes every token's probability to seven distinct sources so that detection rests on the model's full internal mechanics rather than on partial proxy signals. The goal is a supervised binary classifier that flags hallucinated responses using syntax-conditioned attribution patterns.

Key Contributions

1. TPA, a framework that mathematically attributes each token's probability to seven sources: Query, RAG Context, Past Token, Self Token, FFN, Final LayerNorm, and Initial Embedding. The attribution scores sum to the token's final probability, giving a complete decomposition of the generation process (Theorem 1).

2. A syntax-aware aggregation mechanism that quantifies how attribution sources drive distinct parts of speech. This lets the detector flag anomalies in specific linguistic categories, such as unusual LayerNorm contribution to numerals, while ignoring benign grammatical patterns.

3. A fine-grained attribution procedure that splits each attention block's probability contribution across individual heads via logit-based apportionment, then maps those head contributions back to input tokens using attention weights, justified by a first-order Taylor expansion (Proposition 1).

4. Extensive experiments across five LLMs (Llama2-7B/13B, Llama3-8B, Mistral-7B, and Qwen3-8B) reporting state-of-the-art performance on the RAGTruth and Dolly (AC) benchmarks, plus SHAP-based interpretability analysis.

Main Findings

• RAGTruth results (Llama2-7B, Llama2-13B): TPA achieves statistically significant Rank-1 results (p < 0.05) on both F1 and AUC. Llama2-7B: AUC 0.7873, Recall 0.8328, F1 0.7238. Llama2-13B: AUC 0.8681, Recall 0.7913, F1 0.7975. The strongest baseline on Llama2-7B, Novo, reached AUC 0.7608, Recall 0.8274, F1 0.7057.

• RAGTruth results (Llama3-8B): TPA ranks first in F1 (0.7843) and Recall (0.7860), but is statistically comparable to the strongest baselines. Novo reports the higher AUC on this backbone (0.8258 versus TPA's 0.8211).

• Mistral-7B: The largest improvement appears here, with TPA reaching F1 0.8702, AUC 0.9096, and Recall 0.9200, outperforming Novo (F1 0.8000, AUC 0.8419, Recall 0.8765) by 7% on F1. The authors interpret this as good transfer to newer architectures with sliding-window attention.

• Qwen3-8B: TPA outperforms Novo and TSV on F1 (0.6006), AUC (0.8236), and Recall (0.7130), versus Novo's 0.5468, 0.7919, and 0.6259, and TSV's 0.5493, 0.7710, and 0.6259.

• Dolly (AC) scaling trend: TPA trails baselines such as ReDeEP (F1 0.7833) on Llama2-7B with F1 0.7527, but secures significant Rank-1 performance on Llama2-13B (AUC 0.8159, Recall 0.9741, F1 0.8075) and the best F1 on Llama3-8B (0.7529).

• Ablation of attribution sources: Removing RAG or FFN consistently degrades accuracy (a 3.01% F1 drop on Llama-2-7B). Removing LayerNorm causes a sharp 5.83% F1 drop on Llama-3-8B. Removing Self yields marginal gains in some cases, such as on Mistral-7B, but the full feature set is retained for cross-architecture robustness.

• POS aggregation ablation: TPA-POS attains the best F1 in all seven backbone-dataset settings, with gains of 1.8–4.6% F1 over TPA-Mean on RAGTruth and up to 22.5% on Dolly (LLaMA2-13B). TPA-Mean uses a 7-dimensional feature and TPA-Stat a 14-dimensional one, both pooling over all tokens.

• Interpretability (SHAP): The classifier depends on features beyond the FFN-RAG conflict. LN_NUM plays a decisive role in Llama2-7B, and QUERY_NOUN appears as a top predictor in Mistral-7B. RAG_NOUN dominates in Llama2 and Mistral, whereas Llama3-8B relies on RAG_ADP (Adpositions). The LN_NUM signal reverses between Llama2-7B and Llama2-13B, where higher values correlate with factuality.

• Runtime: The attribution pipeline adds three post-hoc stages with total complexity O(LT|V|d + LTd² + LHT²). On an A100-40GB, feature extraction takes about 20 seconds per response in the sequential implementation, roughly 17 GPU-hours per LLM per dataset, while classifier inference is under one second per response.

Methodology in Plain English

The method starts from the fact that a Pre-LN transformer builds its final representation by repeatedly adding components to a residual stream: an initial token-plus-position embedding, then for each layer an attention update and an FFN update, and finally a LayerNorm adjustment. The authors introduce a probe function, similar in spirit to the logit lens, that reads out the probability of the target next token directly from any intermediate state. By measuring the probed probability before and after each component, they isolate how much each component changed the probability of the token that was actually generated, and Theorem 1 shows these changes sum exactly to the final output probability.

Attention blocks are the hard part because the output of an attention block is a sum of head outputs, but the probe passes through a Softmax that is not additive. The authors therefore compute each head's scalar contribution in logit space, then distribute the block's probability contribution across heads in proportion to the exponential of those logit contributions. A first-order Taylor expansion motivates this as an approximation, and the normalization keeps the head scores summing to the layer total. Once head contributions are known, attention weights are used to split each head's contribution across four disjoint input regions: the user query, the retrieved context, previously generated tokens, and the current token itself.

This yields a seven-dimensional attribution vector per token. Because a high FFN contribution is unremarkable for a function word but suspicious for a proper noun, the vectors are averaged within Part-of-Speech categories. SpaCy assigns tags, and sub-word tokens inherit the tag of their parent word through a tag-propagation rule. The final feature vector is 126-dimensional: seven sources crossed with 18 universal POS tags. An ensemble of five XGBoost classifiers with different random seeds performs the binary hallucination classification. Feature extraction needs only a single teacher-forced pass over the generated response, avoiding autoregressive resampling.

Why This Matters

Impact on research: The paper challenges the assumption that RAG hallucination is reducible to an FFN-versus-context conflict. It shows that LayerNorm and query-related contributions carry independent detection signal, and that these signals differ across architectures, which argues for learnable, syntax-aware detectors over static heuristics. It also connects probability attribution methods to a downstream reliability task with quantified gains.

Real-world applications (drawn from the paper's stated motivations):

• Clinical decision support, where hallucinated medication dosages can harm patients.
• Legal research assistants, where fabricated case citations have led to sanctioned court filings.
• Autonomous agents that act on retrieved evidence, where ungrounded intermediate responses propagate into downstream actions.
• General RAG question answering over retrieved documents, where responses inconsistent with the retrieved context need to be flagged before they reach users.

Industry relevance: Teams deploying white-box, open-weight LLMs in RAG pipelines can apply TPA as a post-hoc detector without retraining the generator, since it operates on a single teacher-forced pass plus a lightweight classifier. The reported cost of roughly 20 seconds per response for feature extraction and under one second for classification indicates a practical trade-off for high-stakes filtering. The requirement for white-box access, however, rules out closed-source API models.

Future Directions

• Extending TPA to phrase-level or span-level attribution to improve efficiency.
• Active mitigation through online monitoring and intervention when risky patterns appear, such as abnormal FFN or LayerNorm reliance.
• Studying TPA under noisy or adversarial retrieval, where expected source shifts could themselves serve as detection signals.
• Extending POS-based aggregation to non-English and specialized domains such as code and scientific text, using multilingual spaCy or AST node types.

Open questions raised but not resolved include how well the first-order approximation holds for low-confidence generations, whether the architecture-specific signals such as the LN_NUM reversal can be predicted in advance, and how the method degrades when the assumption of relevant and correct retrieved context is violated.

Target Audience

Researchers and graduate students working on hallucination detection, RAG reliability, and mechanistic interpretability of transformers will get the most from this paper, since it sits at the intersection of probability attribution and applied detection. Practitioners building RAG systems on open-weight models, such as Llama, Mistral, or Qwen backbones, will find the detection recipe and runtime figures directly relevant. Readers interested in SHAP-based model diagnostics and syntax-aware feature engineering will also benefit from the interpretability section. The paper is not suited to those without white-box access to model internals, as the authors themselves note in the limitations.

Authors’ abstract

Detecting hallucinations in Retrieval-Augmented Generation remains a challenge. Prior approaches attribute hallucinations to a binary conflict between internal knowledge stored in FFNs and the retrieved context. However, this perspective is incomplete, failing to account for the impact of other components of the LLM, such as the user query, previously generated tokens, the self token, and the final LayerNorm adjustment. To comprehensively capture the impact of these components on hallucination detection, we propose TPA which mathematically attributes each token's probability to seven distinct sources: Query, RAG Context, Past Token, Self Token, FFN, Final LayerNorm, and Initial Embedding. This attribution quantifies how each source contributes to the generation of the next token. Specifically, we aggregate these attribution scores by Part-of-Speech (POS) tags to quantify the contribution of each model component to the generation of specific linguistic categories within a response. By leveraging these patterns, such as detecting anomalies where Nouns rely heavily on LayerNorm, TPA effectively identifies hallucinated responses. Extensive experiments show that TPA achieves state-of-the-art performance.