RCStat: A Statistical Framework for using Relative Contextualization in Transformers

arXiv:2506.19549 · cs.CL, cs.AI, cs.LG · Submitted 2025-06-24 · Read on arXiv

cs.CL, cs.AI, cs.LG

Submitted: 2025-06-24

Updated: 2026-10-13

Comments: Accepted to NeuirIPS 2026

Code: https://github.com/NVIDIA/kvpress

License: http://creativecommons.org/licenses/by-nc-sa/4.0/

The gist: Identifying which tokens and activations truly influence model predictions is critical for both the efficiency and interpretability of large auto-regressive language models.

Terminology

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 RCStat, a statistical framework for quantifying contextual influence in attention mechanisms. At its core 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 attribution, attention heads with high expected RC accurately identify the prompt spans that drive generation, enabling reliable span-level explanations. For 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.

Sources

Related papers