💬 NLP

Attributing Response to Context: A Jensen-Shannon Divergence Driven Mechanistic Study of Context Attribution in Retrieval-Augmented Generation

This paper introduces ARC-JSD, a computationally efficient method that uses Jensen-Shannon Divergence to accurately attribute generated responses to specific context segments in RAG systems without requiring fine-tuning, while also providing mechanistic insights into the specific attention heads and MLP layers responsible for this attribution.

Original authors: Ruizhe Li, Chen Chen, Yuchen Hu, Yanjun Gao, Xi Wang, Emine Yilmaz

Published 2026-02-12

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Ruizhe Li, Chen Chen, Yuchen Hu, Yanjun Gao, Xi Wang, Emine Yilmaz

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are a detective trying to figure out exactly which witness in a crowded room provided the one crucial piece of evidence that solved a crime.

In the world of Artificial Intelligence, specifically Retrieval-Augmented Generation (RAG), this is a massive problem. RAG is a technique where an AI (like ChatGPT) doesn't just rely on its memory; it "looks up" information in a library of documents first, then writes an answer.

The Problem: The "Messy Library" Effect
When the AI gives you an answer, it’s often looking at dozens of different sentences at once. If the AI says, "The capital of France is Paris," and it was looking at a massive textbook, how do we know which specific sentence actually taught it that?

Current methods to find this "smoking gun" sentence are either:

Too slow: They require the AI to "re-think" the entire answer hundreds of times, changing one word at a time to see what happens.
Too expensive: They require humans to sit there and manually label every sentence.

The Solution: ARC-JSD (The "Information Echo" Method)

The researchers created a new method called ARC-JSD. To understand how it works, let's use two analogies.

1. The "Missing Ingredient" Test (The Core Logic)

Imagine you are tasting a complex soup. You want to know if the garlic is the reason it tastes so savory.

The Old Way: You'd have to cook 100 different batches of soup, each with a slightly different amount of garlic, and compare them all. (This is what current AI methods do—it's exhausting!)
The ARC-JSD Way: You take a spoonful of the soup, then you take a spoonful of the soup without the garlic. You compare the "flavor profile" (the probability distribution) of the two. If the flavor changes drastically, you’ve found your ingredient. If the flavor stays almost the same, the garlic wasn't important.

2. The "Jensen-Shannon Divergence" (The Measuring Tape)

How do you mathematically measure that "change in flavor"? They use something called Jensen-Shannon Divergence (JSD).

Think of JSD as a highly sensitive musical tuner. If you play a chord (the full context) and then play the same chord but remove one note (the sentence we are testing), JSD measures how much the "harmony" of the music shifts.

If the harmony stays perfect, that note wasn't part of the song.
If the harmony suddenly sounds "off" or completely different, that note was essential.

Why is this a big deal?

1. It’s Fast and Efficient (The Speed Demon)
Because they aren't re-training the AI or doing hundreds of "re-tastes," this method is up to three times faster than previous methods. It’s like finding the missing ingredient by a quick sniff rather than cooking 100 pots of soup.

2. It Peeks Under the Hood (The X-Ray Vision)
The researchers didn't just find the sentences; they used this method to look inside the AI's "brain." They discovered that certain "gears" (called Attention Heads and MLP layers) are specifically responsible for handling this retrieved information. It’s like discovering that a car has a specific gear that only engages when you're using GPS.

3. It Stops "Hallucinations" (The Lie Detector)
Because we can now identify exactly which sentence the AI is using, we can "gate" the AI. If the AI tries to say something that doesn't have a strong "echo" from the provided text, we can catch it. This helps prevent the AI from making things up (hallucinating).

Summary in a Nutshell

This paper provides a faster, cheaper, and more accurate way to prove that an AI is actually telling the truth based on the facts it was given, rather than just making up a convincing story. It turns the AI from a "black box" that just spits out answers into a transparent researcher that can point its finger and say, "I know this because of THIS sentence right here."

Technical Summary: Attributing Response to Context (ARC-JSD)

Published as a conference paper at ICLR 2026

1. Problem Statement

Retrieval-Augmented Generation (RAG) improves Large Language Model (LLM) accuracy by providing external context. A critical requirement for trustworthy RAG is context attribution: the ability to identify which specific segments of the retrieved context were actually used to generate a response.

Current attribution methods suffer from two main drawbacks:

Computational Intensity: Methods like gradient-based feature attribution or training linear surrogate models (e.g., ContextCite) require extensive fine-tuning, back-propagation, or hundreds of forward passes, making them slow and expensive.
Complexity/Inaccuracy: Existing inference-time methods often rely on complex modeling that may fail to capture non-linearities in how models process information.

2. Methodology: ARC-JSD

The authors propose ARC-JSD (Attribute Response to Context via Jensen–Shannon Divergence), a novel, training-free, and inference-only method.

Core Mechanism: The method is based on the "inference-attribution assumption": if removing a context segment significantly alters the probability distribution of the generated response, that segment is essential for grounding.
Metric Selection (JSD): Instead of using KL Divergence (which is asymmetric and numerically unstable when probabilities approach zero) or Total Variation (which ignores where probability mass moves), the authors use Jensen–Shannon Divergence (JSD). JSD is symmetric, bounded $[0, \log 2]$ , and scale-free, making it ideal for comparing divergence across different layers and models.
The Process:
1. Generate a response $R$ using the full context $C$ and query $Q$ .
2. For each sentence $c_i$ in the context, perform a single ablation (remove $c_i$ ).
3. Calculate the JSD between the original response distribution $P(R|C, Q)$ and the ablated distribution $P(R|C_{ablated}, Q)$ .
4. Rank sentences by their JSD scores; the highest score identifies the most critical context segment.
Mechanistic Integration: The authors extend this to the model's internal architecture using the Logit Lens. By applying JSD to the outputs of individual attention heads and MLP layers, they can pinpoint exactly which internal components are responsible for the attribution behavior.

3. Key Contributions

Efficient Attribution Algorithm: Developed ARC-JSD, which requires no fine-tuning, no gradients, and no surrogate modeling, allowing for direct integration into existing RAG pipelines.
Mechanistic Interpretability: Conducted a systematic study uncovering that context attribution is primarily driven by attention heads and MLP layers in the intermediate and higher layers of the transformer.
Hallucination Mitigation: Demonstrated that by using JSD scores as "confidence gates" to mask unreliable attention heads/MLPs, the hallucination rate can be significantly reduced.
Mathematical Validation: Provided a rigorous comparison showing why JSD outperforms KL, Wasserstein, Total Variation, and MMD in the specific context of LLM decoding behavior.

4. Results

Accuracy: ARC-JSD demonstrated superior accuracy across multiple benchmarks (TyDi QA, Hotpot QA, and MuSiQue), outperforming previous baselines by an average of 10.7% in context attribution accuracy.
Efficiency: The method achieved up to a 3-fold speedup compared to existing surrogate-based and gradient-based methods.
Robustness: Evaluations on various model scales (Qwen2-1.5B to Gemma2-9B) and larger models (LLaMA-3.1-8B) confirmed the method's generalizability.
Hallucination Control: In experiments on Hotpot QA, using JSD-based gating reduced the hallucination rate by approximately 39% while maintaining factual F1 scores.
Mechanistic Verification: The JSD-based location of critical components showed a high Spearman correlation with "semantic gain" (a metric measuring the direct contribution of a layer to the correct answer), validating the mechanistic findings.

5. Significance

This work moves RAG from a "black box" toward a transparent and controllable system. By providing a lightweight way to audit why a model said something, ARC-JSD enhances the reliability and interpretability of AI in high-stakes domains like medicine and law. Furthermore, the ability to manipulate internal layers to suppress hallucinations provides a new pathway for "steering" LLMs during inference without the massive cost of retraining.