RCStat: A Statistical Framework for using Relative Contextualization in Transformers

Debabrata Mahapatra; Shubham Agarwal; Apoorv Saxena; Subrata Mitra

cs.CLJun 24, 2025

RCStat: A Statistical Framework for using Relative Contextualization in Transformers

Authors: Debabrata Mahapatra, Shubham Agarwal, Apoorv Saxena, Subrata Mitra

Organizations: Adobe Research, India

Abstract

Identifying which tokens and activations truly influence model predictions is critical for both the efficiency and interpretability of large auto-regressive language models. Yet in practice, token importance is typically inferred from attention distributions that entangle contextual relevance with softmax-normalization effects, obscuring relative influence across tokens and heads. We propose $\textbf{RCStat}$ , a statistical framework for quantifying contextual influence in attention mechanisms. At its core $\textbf{Relative Contextualization (RC)}$ is a random variable that measures how strongly one subset of tokens contributes to another under the model's internal scoring scheme. RCStat admits computationally efficient bounds on expected influence that can be estimated at inference time. No retraining is required. We apply RCStat to two tasks. For $\textbf{attribution}$ , attention heads with high expected RC accurately identify the prompt spans that drive generation, enabling reliable span-level explanations. For $\textbf{key-value cache eviction}$ , RC-based adaptive thresholding selectively evicts low-impact KV entries, substantially reducing cache size while preserving generation quality. Across question answering, summarization, and attribution benchmarks, RCStat achieves consistent gains, improving generation quality by 15-40% with upto 36% error reduction for KV eviction, and attribution accuracy by 2-16%. These results demonstrate that explicitly modeling contextual influence provides a principled and practical alternative to attention-based heuristics.

Figures & tables

Figure 1: Attribution between generated text (right) and input prompt (left) is analyzed using pre-softmax attention logits vs. post-softmax values. The bottom histogram highlights pre-softmax logits, separating prompt tokens (cross-

\langle q,k\rangle

) from generated tokens (self-

\langle q,k\rangle

), at the

13^{\text{th}}

layer and

23^{\text{th}}

head of Llama-3B-instruct. Tokens in the overlapping region signify common content, which is suppressed in the post-softmax histogram above. Using logit distributions, we attribute (§ 4.1 ): (1) prompt parts not contextualized during generation, (2) generation parts uninfluenced by the prompt, and (3) generation heavily contextualized by the prompt (red dashed arrows).

Panel of Figure 2 — Figure 2: Illustration of the pre-softmax attention logits $Q^{T}K\in\mathbb{R}^{|s|\times|s|}$ , where $s$ is sequence $s$ with prompt $p$ and generated tokens $g$ , as they appear in the last $|g|$ rows the the $Q^{T}K$ matrix. Fig. 2(a) shows how the logits in the last $|g|$ rows are partitioned to construct the cross-contextualization (CC) and self-contextualization (SC) random variables. Fig. 2(b) and 2(c) show the logits that construct conditional CC and conditional SC. Fig. 2(d) shows their respective probability density function.

Panel of Figure 3 — Figure 3: Illustrations of upper and lower bounds ( Theorem 3.4 ) for two types of RC: $\max(X-Y,0)$ and $\max(Y-X,0)$ .

Panel of Figure 4 — Figure 4: Using LLaMA-3B (28 layers, 24 heads) on the listed datasets. Fig. 4(b) shows $\mathbb{E}[RC]$ ( 4 ) and its upper bound ( 2 ) (overlap area in Fig. 3(b) ) for all heads. Fig. 4(c) and 4(d) show head RC scores across datasets, sorted by median. Fig. 4(a) shows the anti-correlation between head compression ratio and RC score (each point is a head).

Panel of Figure 6 — Figure 6: KV-cache eviction performance on QMSum using LLaMA-3.2-3B and LLaMA-3.1-8B. We report Value Error Rate (VER↓) and generation quality (RL-F1↑) across different strategies.

Panel of Figure 7 — Figure 7: KV-cache eviction results on 2WikiMultiHop and SQuAD v2.0 using LLaMA-3.1-8B.

Appendix figures & tables21 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 8: Example 1: Relevant context

Figure 9: Example 2: Irrelevant context

Panel of Figure 10 — Figure 10: All plots are generated using the mentioned datasets with LLaMA-3B model: 28 layers, 24 heads in each layer. All plots show the RC score of the heads (upper bound of expected RC).

Panel of Figure 11 — Figure 11: All plots are generated using the mentioned datasets with LLaMA-8B model: 32 layers, 32 heads in each layer. All plots show the RC score of the heads (upper bound of expected RC).

Panel of Figure 12 — Figure 12: All plots are generated using the mentioned datasets with Qwen3-8B model: 36 layers, 32 heads in each layer. All plots show the RC score of the heads (upper bound of expected RC).

Panel of Figure 13 — Figure 13: VER scores for 3B model on QMSum, 2WikiMultiHop, and SQuAD datasets.

Panel of Figure 14 — Figure 14: VER scores for 8B model on QMSum, 2WikiMultiHop, and SQuAD datasets.

Panel of Figure 15 — Figure 15: ROUGE-1 F1 scores for 3B model on QMSum, 2WikiMultiHop, and SQuAD datasets.

Panel of Figure 16 — Figure 16: ROUGE-L F1 scores for 3B model on QMSum, 2WikiMultiHop, and SQuAD datasets.

Panel of Figure 17 — Figure 17: ROUGE-1 F1 scores for 8B model on QMSum, 2WikiMultiHop, and SQuAD datasets.

Panel of Figure 18 — Figure 18: ROUGE-L F1 scores for 8B model on QMSum, 2WikiMultiHop, and SQuAD datasets.

Model	Method	50%	60%	70%
3B	TOVA	0.2408	0.3103	0.3905
	RCStat (IOT)	0.1956	0.2648	0.3402
	RCStat (Non-IID)	0.1571	0.2295	0.3066
8B	TOVA	0.2177	0.2859	0.3639
	RCStat (IOT)	0.1615	0.2290	0.3007
	RCStat (Non-IID)	0.1043	0.2034	0.2836

Appendix

Table 2: Value Error Rate (VER) on the QMSum dataset across different compression ratios (50%, 60%, 70%) for LLaMA-3.2-3B and LLaMA-3.1-8B Instruct models. RCStat (IOT) assumes Independent Output Tokens , while RCStat (Non-IOT) does not assume any independence. Here, lower is better.

	QMSum					2Wiki					SQuAD v2
Method CR	0.40	0.50	0.60	0.70	0.80	0.40	0.50	0.60	0.70	0.80	0.40	0.50	0.60	0.70	0.80
RCStat ( $w=16$ )	0.12	0.16	0.22	0.31	0.47	0.13	0.18	0.24	0.32	0.46	0.09	0.13	0.18	0.26	0.39
RCStat ( $w=8$ )	0.13	0.16	0.21	0.29	0.45	0.14	0.19	0.25	0.33	0.48	0.10	0.14	0.19	0.27	0.41
TOVA	0.19	0.21	0.25	0.33	0.49	0.20	0.24	0.29	0.37	0.53	0.16	0.19	0.23	0.30	0.46
KNorm	0.39	0.41	0.48	0.61	0.80	0.40	0.45	0.52	0.65	0.85	0.32	0.37	0.44	0.58	0.77
SnapKV	0.45	0.55	0.72	1.34	2.33	0.47	0.58	0.75	1.30	2.25	0.41	0.52	0.69	1.20	2.10

Appendix

Table 3: Value Error Rate (VER) with Qwen3-8B for different KV compression methods across three datasets. RCStat consistently shows strong performance across Compression Ratios (CR).

	QMSum						SQuAD v2						2WikiMultiHopQA
$w$	0.4	0.5	0.6	0.7	0.8	0.9	0.4	0.5	0.6	0.7	0.8	0.9	0.4	0.5	0.6	0.7	0.8	0.9
1	-	-	-	-	-	0.74	-	-	-	-	-	1.00	-	-	-	-	-	0.92
2	-	-	-	-	0.60	0.63	-	-	-	-	0.51	0.63	-	-	-	-	-	0.48
4	-	-	0.27	0.36	0.42	0.52	-	-	0.25	0.29	0.44	0.60	-	-	0.14	0.19	0.24	0.36
8	0.12	0.15	0.20	0.28	0.40	0.51	-	0.13	0.19	0.28	0.42	0.59	0.05	0.07	0.11	0.14	0.20	0.35
16	0.11	0.16	0.22	0.28	0.42	0.54	0.10	0.15	0.22	0.33	0.45	0.61	0.04	0.08	0.12	0.17	0.27	0.42

Appendix

Table 4: VER under varying window sizes (

w

) and compression ratios (CR) for KV cache compression. Bold numbers indicate Pareto-optimal results.

Model	QuoteSum	Veri-Gran
L3.1-8B (all heads)	90.54	77.91
L3.1-8B (post-softmax least RC, $k{=}20$ )	35.72	4.69
L3.1-8B (pre-softmax least RC, $k{=}20$ )	29.49	2.81
L3.1-8B (post-softmax most RC, $k{=}20$ )	90.03	71.25
L3.1-8B ( pre-softmax most RC , $k{=}20$ )	93.91	79.37

Appendix

Table 5: Chunk-level accuracy for attributing extractive spans on QuoteSum and Veri-Gran (“L3.1” = LLaMa-3.1, “L3.2” = LLaMa-3.2). For LLaMa-3.2-3B and Qwen3-8B, values are pre-softmax. “HS” = best single hidden-state layer baseline.

Figure 19: (see on screen) The distribution of RC upper bound (overlap area) for QmSum dataset. The first (last) row corresponds to heads in the first (last) layer.

Figure 20: (see on screen) The distribution of RC upper bound (overlap area) for Squad v2 dataset. The first (last) row corresponds to heads in the first (last) layer.

Figure 21: (see on screen) The distribution of RC upper bound (overlap area) for 2WikiMultiHop dataset. The first (last) row corresponds to heads in the first (last) layer.

Panel of Figure 22 — Figure 22: (see on screen) Percentiles of RC upper bound (overlap area) across the QmSum dataset

Panel of Figure 24 — Figure 24: (see on screen) Percentiles of RC upper bound (overlap area) across the 2WikiMultiHop dataset

Panel of Figure 25 — Figure 25: Attention heatmaps from Layer 15 using three attribution strategies. The top head yields focused and accurate attribution, the mean head shows diluted but somewhat relevant attention, and the worst head highlights largely irrelevant regions.

Attention Mechanisms

Attention mechanisms let neural networks weight information from other tokens or regions according to relevance, enabling context-dependent processing.

See all 512 papers

Oct 12, 2026

RIFT: Token Influence Follows the Path of Least Resistance

Which context tokens drive a language model's answer? Graph-based tools such as attention rollout and attention flow answer this from a single forward pass, yet their faithfulness in decoder-only models is rarely tested against interventions. We view the layers of a decoder as a causal network whose conductances are attention weights, and read a token's influence on the answer from the resistance between them. We propose RIFT (Resistance-based Information Flow in Transformers), which scores a token by the series resistance of its path of least resistance to the answer, found exactly by a classic shortest-path dynamic program. On synthetic in-context retrieval tasks and extractive question answering over Wikipedia paragraphs, RIFT predicts the effect of replacing a token with a Spearman correlation of 0.52-0.68; predicting a single replacement, it reaches 77-96% of the agreement between two different random replacements. It needs no gradients and no extra forward passes, scoring every token of a prompt in 0.9-3.4 ms on one CPU thread, whereas the replacement experiment needs a forward pass per token. On 576 held-out prompts from four models of two families (Qwen3 and Llama 3.2), it outperforms token distance, attention rollout and attention flow in every pre-specified comparison, by 0.23 to 0.88 in mean Spearman correlation, and it is the best or tied-best of all attention-based scores we test. Aggregates that add up all routes (attention flow, effective resistance and rollout) instead mostly encode token position: as a simple proposition shows, early tokens have many more parallel routes to the answer.

Tingan Jin, Shuhang Dong, Haosong Li

University of California, Los Angeles Los Angeles, CA, USA · Independent Researcher

Transformer Interpretability LLM Interpretability Causal Interventions in Language Models

Ask PDF

May 10, 2026

Make Each Token Count: Towards Improving Long-Context Performance with KV Cache Eviction

The key-value (KV) cache is a major bottleneck in long-context inference, where memory and computation grow with sequence length. Existing KV eviction methods reduce this cost but typically degrade performance relative to full-cache inference. Our key insight is that full-cache attention is not always optimal: in long contexts, irrelevant tokens can dilute attention away from useful evidence, so selective, learnable eviction can improve generation rather than merely approximate the full cache. We introduce a global retention-based KV eviction method that learns each token's future utility under a unified memory budget. Lightweight retention gates assign utility scores to cached KV entries, and a shared final scoring projection calibrates these scores across all layers and heads. This enables a single global eviction policy in which tokens from different layers, heads, and modalities compete directly for cache capacity. We further provide theoretical analysis showing that preferentially retaining useful tokens reduces attention dilution, and we justify geometric retention as a query-agnostic proxy for future utility. Across diverse long-context language and vision-language reasoning, and multi-turn dialogue benchmarks, our method substantially reduces KV memory while matching or surpassing full-cache inference. These results suggest that learned, globally calibrated KV eviction is not only a compression technique, but also a mechanism for improving long-context reasoning.

Ngoc Bui, Hieu Trung Nguyen, Arman Cohan +1

Department of Computer Science, Yale University · The Chinese University of Hong Kong

Self-Attention KV Caching KV-Cache Eviction Long-Context Inference KV-Cache Pruning

Ask PDF

Oct 19, 2026

DeltaTrace: Efficient Signed Attribution for Reasoning Language Models

We study how words in the context support or oppose a language model's response and present DeltaTrace, a novel attribution framework that explains the effect of input tokens on a fixed model response. DeltaTrace measures how the response's log-probability changes when the context tokens being explained are masked. To distribute this difference among source tokens, it constructs operator-specific local rules for finite changes and composes them backward through the computational graph using effective secant slopes. Each token is assigned a signed attribution as the inner product of its embedding displacement and the corresponding backward-propagated slope. Local conservation makes these contributions sum to the measured score difference and preserves their meaning when aggregated over sentences or passages. DeltaTrace supports both attention and linear-attention models, with auxiliary memory growing linearly with sequence length at fixed model dimensions and chunk size. On the evaluated attention model, DeltaTrace achieves the lowest deletion-faithfulness scores (RISE and MAS) on most retrieval and reasoning tasks and the highest needle-in-a-haystack Recall, exceeding FlashTrace by 5.67-24.33 percentage points. On the evaluated hybrid attention and linear-attention model, DeltaTrace lowers RISE on 9 tasks and MAS on 12 relative to FlashTrace.

Chen Liang, Daoping Wu, Donghua Yang +4

Harbin Institute of Technology · Zhejiang Lab

Feature Attribution LLM Interpretability Faithfulness of Language Model Explanations

Ask PDF

Abstract

Attention Mechanisms

Explore similar work

RIFT: Token Influence Follows the Path of Least Resistance

Make Each Token Count: Towards Improving Long-Context Performance with KV Cache Eviction

DeltaTrace: Efficient Signed Attribution for Reasoning Language Models