Prompt injection defence via internal attribution
Research
Read article
Research
Closing the hallucination attribution gap
When AI hallucinates, ask: where did the answer come from?
Large language models have become highly fluent. This fluency has a problematic side effect: a model can produce an answer that reads as plausible while being unsupported, incorrect, or fabricated. In search, retrieval-augmented generation, compliance workflows, and tool-using agents, this is hard to detect, because a fabricated answer is as well-formed as a correct one.
We distinguish two kinds of hallucination. Factual hallucinations misstate the world. Faithfulness hallucinations fail to stay grounded in the information the model was given. We focus on the faithfulness case: the model is provided with the evidence it needs, e.g. a retrieved document, a table, or a user request, and still produces an output the context does not support. Grounding does not require copying. A model may paraphrase the user, normalise a date, or convert three hundred euros into €300 and remain faithful. The failure we target is a value that never originated in the input.
The core idea: hallucinations leave an attribution gap
A contextual hallucination cannot occur if the response can be traced back to the supporting evidence. If a generated value was determined by the input, it is grounded. Attribution makes this property verifiable. For each generated token, token-wise relevance indicates which parts of the prompt, e.g. the user request, a retrieved document, a table, or a tool result, contributed to it. A grounded value concentrates relevance on the corresponding evidence. A fabricated value leaves relevance on the context weak, absent, or misplaced.
We call this missing connection the attribution gap, and use it to flag ungrounded responses automatically.
A high-level primer on LRP
Layer-wise Relevance Propagation (LRP) is an explainability method for neural networks, and the foundation of layerwise's name. It traces a model's output backward through the network to estimate which input sources and internal components were most relevant for that output. LRP was introduced to explain the predictions of nonlinear classifiers by decomposing a decision into relevance scores over input features, such as the pixels of an image (Bach et al., 2015).
For language models the input features are tokens rather than pixels. Given a generated output, LRP propagates relevance backward layer by layer and produces a token-wise relevance map: which parts of the context the model used, and which it ignored. Our attention-aware extension, AttnLRP, attributes transformer outputs not only to input tokens but also to internal representations, while remaining efficient enough for large models (Achtibat et al., 2024).
Turning the trace into a score
We read the same attribution in two ways, a coarse measure and a finer one.
1. The fraction — coarse
The share of the answer's total relevance that lands on the trusted source. A high share indicates grounding, a low share indicates fabrication. A single number and a single threshold, with no training, computed cheaply.
2. Circuit-matching — fine-grained
Which attention heads performed the reading, compared against the signature characteristic of grounded answers. This reads the full head signature rather than a single aggregate, which keeps it reliable when the data distribution changes.
Why output-grading is insufficient
A fabricated value is, by construction, indistinguishable from a real one in the text. An invented order identifier has the right number of digits; a fabricated budget is a plausible round figure. A second model asked to grade the answer — the LLM-as-judge approach — therefore fails on the cases that matter, and increasing its size does not resolve this, because the discriminating information is not in the words but in where the answer came from.
| detecting ungrounded answers (AUROC) | reading the model(ours) | best LLM judge |
|---|---|---|
| Retrieval (RAG) | ≈0.99±0.01 | 0.92±0.04 |
| Tool-call arguments | 0.85±0.02 | 0.63–0.72±0.02 |
Why context usage matters in RAG and tool-using agents
Retrieval-Augmented Generation was developed in part to reduce hallucination by combining a model's parametric memory with retrieved documents from an external corpus (Lewis et al., 2020). It cannot guarantee grounded answers: the retriever may return irrelevant or incomplete context, the context may conflict, chunking may separate related evidence, or the generator may simply fail to use what was retrieved. A model can receive the right document and still produce the wrong answer.
The same problem appears one step earlier in tool-using agents, when the model constructs the tool call itself. A user may provide only part of the information — a destination and a travel date — and the model must decide which parameters to pass to a flight-search API. Here a hallucination takes the form of a parameter the user never provided, or one that should have been left unspecified.
User request:
I need a flight from Berlin to Paris tomorrow morning.
Generated tool call:
search_flights(
origin="Berlin",
destination="Paris",
date="tomorrow",
departure_time="morning",
budget_eur=1000
)Origin, destination, date, and departure time are grounded in the request. The value budget_eur=1000 is not; the user never mentioned a budget. With LRP we can check whether the budget value received relevance from the user prompt. Near-zero relevance on the user context is a strong indication of a tool-value hallucination.
Benchmark: tool-value hallucination in flight booking
To evaluate this, we built a flight-booking benchmark. The system receives the user context for a booking and the model must generate a tool call. We examine cases where the model fills tool-parameter values that the user did not provide. The focal value is the flight budget, present in only 33% of examples. When the budget is absent, a faithful model leaves the parameter unspecified; a hallucinating one invents a value.
Published benchmarks: NQ-Swap and ShortcutsBench
The same signal holds on standard datasets, in both regimes. For retrieval we use NQ-Swap, in which the model receives a document whose fact has been deliberately swapped, and we measure whether it copies the document or reverts to memorised knowledge. For tool use we use ShortcutsBench, real application tool-calls in which some arguments are absent from the request and must be asked for, not guessed. In both, reading the model separates grounded from ungrounded answers more sharply than the LLM judge.
| detecting ungrounded answers (AUROC) | fraction(ours) | circuit-matching(ours) | LLM judge |
|---|---|---|---|
| NQ-Swapretrieval (RAG) | 0.89±0.05 | 0.96±0.04 | 0.59±0.08 |
| ShortcutsBenchtool arguments | 0.82±0.02 | 0.85±0.02 | 0.72±0.02 |
Robustness under distribution shift
A detector is only useful after deployment if it generalises beyond its calibration data. We evaluate two shifts. For retrieval we train on one benchmark and test on a different, longer-context one. For tools we test on APIs not seen during calibration, the case of an integration added after deployment. The grounding trace is a property of how the answer was produced, not of the topic, the wording, or the tool being called, and it transfers in both cases.
| train → test (AUROC) | fraction(ours) | circuit-matching(ours) | LLM judge |
|---|---|---|---|
| NQ-Swap → long-contextretrieval (RAG) | 0.99±0.01 | 0.99±0.01 | 0.92±0.04 |
| ShortcutsBench → unseen toolstool use | 0.82±0.02 | 0.85±0.02 | 0.72±0.02 |
The grounding trace reflects the generation process rather than the surface form, so it remains informative on documents and tools the detector was never calibrated on.
From detection to control
The same readout that detects ungrounded generation can also govern it. Once we can see which evidence was used and which was ignored, the model can correct itself before the answer is returned.
Retrieval — flag and regenerate
12% → 1%
On NQ-Swap (n = 150), when the circuit-matching gate flags an answer we regenerate once with an instruction to defer to the document. This reduces parametric-override failures from 12% to about 1% (16 of 18 caught), with no degradation on already-correct answers. On a harder constructed long-context conflict set the same gate reduces failures from 29% to 3%.
Tool use — ask instead of fabricate
43% → 15%
On ShortcutsBench (n = 550), when a required argument is absent the correct action is to ask rather than invent. Left to its own judgment the model either fabricates every missing value or, when told to be cautious, queries the user about 80% of the time. Gating that decision on the head signature reduces harmful fabrications from 43% to 15% without the needless questions.
The readout that detects ungrounded generation therefore also controls it, across both retrieval and agents. Once grounding is inspectable, a system can:
Show the evidence
Highlight the evidence behind each generated value, and warn when an answer is weakly grounded.
Reject or regenerate
Re-run unsupported answers with stronger context constraints instead of returning them.
Audit and debug
Check whether tool outputs were used, and surface retrieved passages the model ignored.
Train for grounding
Use attribution as a training signal, so grounding is encouraged during fine-tuning, not only inspected afterwards.
Scope and outlook
A word on what we are solving. We are not trying to judge whether a statement is true about the world — that is factuality. We focus on faithfulness: whether an answer stays grounded in the evidence the model was actually given. On that question we do two things — we detect when grounding breaks, and we improve it — across both retrieval and tool use.
Grounding is just the first use case. The idea underneath it is simple and general: every answer a model gives is the end of a computation, and that computation can be read. Until now the signal most systems relied on was the output text, and the usual fix for bad text was to bolt on a second model to grade it. Reading the mechanism that produced the answer gives you different evidence — and it is cheap. It is already there in every model you run, it costs a fraction of another forward pass, and it tells you how the answer was formed, not just how it reads.
The same readout is already telling apart failures the output hides: prompt injections that hijack a trusted instruction, models that play down what they can do, reasoning that is there for show rather than load-bearing, and decisions that quietly turn on attributes they should ignore. Each one is really the same question — where did this come from? — rather than whether it looks right.
This is the shift we are excited about: observability and control built on a model's own mechanisms, not only its outputs. Grounding is the first step. As these signals move from offline analysis into the live serving path, interpretability becomes something you run in production — a practical layer for observing, debugging, and steering AI systems toward reliability.
