# Learning to Attribute with Attention

arXiv: https://arxiv.org/abs/2504.13752 · HTML: https://mlsanity.com/paper/2504.13752/

## Key facts

- Cited by 13 papers in the map (data through 14 Oct 2026).
- Builds on 53 of its 67 references that are in the map.
- First posted on arXiv 18 Apr 2025.

## Builds on (most cited first, 30 of 53)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning](https://mlsanity.com/paper/2501.12948/) (2025)
- ["Why Should I Trust You?": Explaining the Predictions of Any Classifier](https://mlsanity.com/paper/1602.04938/) (2016)
- [Qwen2.5 Technical Report](https://mlsanity.com/paper/2412.15115/) (2024)
- [Neural Machine Translation by Jointly Learning to Align and Translate](https://mlsanity.com/paper/1409.0473/) (2014)
- [HotpotQA: A Dataset for Diverse, Explainable Multi-hop Question Answering](https://mlsanity.com/paper/1809.09600/) (2018)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [Lost in the Middle: How Language Models Use Long Contexts](https://mlsanity.com/paper/2307.03172/) (2023)
- [Deep Inside Convolutional Networks: Visualising Image Classification Models and Saliency Maps](https://mlsanity.com/paper/1312.6034/) (2013)
- [Phi-3 Technical Report: A Highly Capable Language Model Locally on Your Phone](https://mlsanity.com/paper/2404.14219/) (2024)
- [MS MARCO: A Human Generated MAchine Reading COmprehension Dataset](https://mlsanity.com/paper/1611.09268/) (2016)
- [SmoothGrad: removing noise by adding noise](https://mlsanity.com/paper/1706.03825/) (2017)
- [Attention is not Explanation](https://mlsanity.com/paper/1902.10186/) (2019)
- [DROP: A Reading Comprehension Benchmark Requiring Discrete Reasoning Over Paragraphs](https://mlsanity.com/paper/1903.00161/) (2019)
- [Abstractive Text Summarization Using Sequence-to-Sequence RNNs and Beyond](https://mlsanity.com/paper/1602.06023/) (2016)
- [Quantifying Attention Flow in Transformers](https://mlsanity.com/paper/2005.00928/) (2020)
- [AGIEval: A Human-Centric Benchmark for Evaluating Foundation Models](https://mlsanity.com/paper/2304.06364/) (2023)
- [Measuring Faithfulness in Chain-of-Thought Reasoning](https://mlsanity.com/paper/2307.13702/) (2023)
- [Attention is not not Explanation](https://mlsanity.com/paper/1908.04626/) (2019)
- [Enabling Large Language Models to Generate Text with Citations](https://mlsanity.com/paper/2305.14627/) (2023)
- [ERASER: A Benchmark to Evaluate Rationalized NLP Models](https://mlsanity.com/paper/1911.03429/) (2019)
- [Towards Measuring the Representation of Subjective Global Opinions in Language Models](https://mlsanity.com/paper/2306.16388/) (2023)
- [Is Attention Interpretable?](https://mlsanity.com/paper/1906.03731/) (2019)
- [Teaching language models to support answers with verified quotes](https://mlsanity.com/paper/2203.11147/) (2022)
- [RARR: Researching and Revising What Language Models Say, Using Language Models](https://mlsanity.com/paper/2210.08726/) (2022)
- [TRAK: Attributing Model Behavior at Scale](https://mlsanity.com/paper/2303.14186/) (2023)
- [Measuring Attribution in Natural Language Generation Models](https://mlsanity.com/paper/2112.12870/) (2021)
- [Retrieval Head Mechanistically Explains Long-Context Factuality](https://mlsanity.com/paper/2404.15574/) (2024)

## Built on by (most cited first, 13 of 13)

- [Attributing Response to Context: A Jensen-Shannon Divergence Driven Mechanistic Study of Context Attribution in Retrieval-Augmented Generation](https://mlsanity.com/paper/2505.16415/) (2025)
- [The Atlas of In-Context Learning: How Attention Heads Shape In-Context Retrieval Augmentation](https://mlsanity.com/paper/2505.15807/) (2025)
- [The Why Behind the Action: Unveiling Internal Drivers via Agentic Attribution](https://mlsanity.com/paper/2601.15075/) (2026)
- [AttnTrace: Contextual Attribution of Prompt Injection and Knowledge Corruption](https://mlsanity.com/paper/2508.03793/) (2025)
- [AgentWatcher: A Rule-based Prompt Injection Monitor](https://mlsanity.com/paper/2604.01194/) (2026)
- [Do LLMs Encode Functional Importance of Reasoning Tokens?](https://mlsanity.com/paper/2601.03066/) (2026)
- [MaxShapley: Towards Incentive-compatible Generative Search with Fair Context Attribution](https://mlsanity.com/paper/2512.05958/) (2025)
- [Explaining the Reasoning of Large Language Models Using Attribution Graphs](https://mlsanity.com/paper/2512.15663/) (2025)
- [Citation Failure: Definition, Analysis and Efficient Mitigation](https://mlsanity.com/paper/2510.20303/) (2025)
- [Real-Time Visual Attribution Streaming in Thinking Model](https://mlsanity.com/paper/2604.16587/) (2026)
- [FlashRT: Towards Computationally and Memory Efficient Red-Teaming for Prompt Injection and Knowledge Corruption](https://mlsanity.com/paper/2604.28157/) (2026)
- [When the Agent Becomes the Kernel: A Systematization of Security on the Path to AI-Native Operating Systems](https://mlsanity.com/paper/2609.23700/) (2026)
- [RIFT: Token Influence Follows the Path of Least Resistance](https://mlsanity.com/paper/2610.16349/) (2026)

## Abstract (arXiv)

Given a sequence of tokens generated by a language model, we may want to identify the preceding tokens that influence the model to generate this sequence. Performing such token attribution is expensive; a common approach is to ablate preceding tokens and directly measure their effects. To reduce the cost of token attribution, we revisit attention weights as a heuristic for how a language model uses previous tokens. Naive approaches to attribute model behavior with attention (e.g., averaging attention weights across attention heads to estimate a token's influence) have been found to be unreliable. To attain faithful attributions, we propose treating the attention weights of different attention heads as features. This way, we can learn how to effectively leverage attention weights for attribution (using signal from ablations). Our resulting method, Attribution with Attention (AT2), reliably performs on par with approaches that involve many ablations, while being significantly more efficient. To showcase the utility of AT2, we use it to prune less important parts of a provided context in a question answering setting, improving answer quality. We provide code for AT2 at https://github.com/MadryLab/AT2 .