# Learning to Focus: Causal Attention Distillation via Gradient-Guided Token Pruning

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

## Key facts

- Cited by 13 papers in the map (data through 14 Oct 2026).
- Placed in KV cache optimization, one of 12 directions in Efficient LLM inference and serving.
- Builds on 49 of its 58 references that are in the map.
- First posted on arXiv 9 Jun 2025.

## Placed in

- Efficient LLM inference and serving / [KV cache optimization](https://mlsanity.com/inference/kv-cache-optimization/)

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

- [Distilling the Knowledge in a Neural Network](https://mlsanity.com/paper/1503.02531/) (2015)
- [Training Verifiers to Solve Math Word Problems](https://mlsanity.com/paper/2110.14168/) (2021)
- [Evaluating Large Language Models Trained on Code](https://mlsanity.com/paper/2107.03374/) (2021)
- [Self-Consistency Improves Chain of Thought Reasoning in Language Models](https://mlsanity.com/paper/2203.11171/) (2022)
- [BERTScore: Evaluating Text Generation with BERT](https://mlsanity.com/paper/1904.09675/) (2019)
- [Measuring Mathematical Problem Solving With the MATH Dataset](https://mlsanity.com/paper/2103.03874/) (2021)
- [HotpotQA: A Dataset for Diverse, Explainable Multi-hop Question Answering](https://mlsanity.com/paper/1809.09600/) (2018)
- [Axiomatic Attribution for Deep Networks](https://mlsanity.com/paper/1703.01365/) (2017)
- [Deep Inside Convolutional Networks: Visualising Image Classification Models and Saliency Maps](https://mlsanity.com/paper/1312.6034/) (2013)
- [LiveCodeBench: Holistic and Contamination Free Evaluation of Large Language Models for Code](https://mlsanity.com/paper/2403.07974/) (2024)
- [Is Your Code Generated by ChatGPT Really Correct? Rigorous Evaluation of Large Language Models for Code Generation](https://mlsanity.com/paper/2305.01210/) (2023)
- [RULER: What's the Real Context Size of Your Long-Context Language Models?](https://mlsanity.com/paper/2404.06654/) (2024)
- [Constructing A Multi-hop QA Dataset for Comprehensive Evaluation of Reasoning Steps](https://mlsanity.com/paper/2011.01060/) (2020)
- [OlympiadBench: A Challenging Benchmark for Promoting AGI with Olympiad-Level Bilingual Multimodal Scientific Problems](https://mlsanity.com/paper/2402.14008/) (2024)
- [When Not to Trust Language Models: Investigating Effectiveness of Parametric and Non-Parametric Memories](https://mlsanity.com/paper/2212.10511/) (2022)
- [Qwen2.5-Math Technical Report: Toward Mathematical Expert Model via Self-Improvement](https://mlsanity.com/paper/2409.12122/) (2024)
- [MuSiQue: Multihop Questions via Single-hop Question Composition](https://mlsanity.com/paper/2108.00573/) (2021)
- [KILT: a Benchmark for Knowledge Intensive Language Tasks](https://mlsanity.com/paper/2009.02252/) (2020)
- [SpinQuant: LLM quantization with learned rotations](https://mlsanity.com/paper/2405.16406/) (2024)
- [∞Bench: Extending Long Context Evaluation Beyond 100K Tokens](https://mlsanity.com/paper/2402.13718/) (2024)
- [Step-DPO: Step-wise Preference Optimization for Long-chain Reasoning of LLMs](https://mlsanity.com/paper/2406.18629/) (2024)
- [TokenSkip: Controllable Chain-of-Thought Compression in LLMs](https://mlsanity.com/paper/2502.12067/) (2025)
- [Token-Budget-Aware LLM Reasoning](https://mlsanity.com/paper/2412.18547/) (2024)
- [Using an LLM to Help With Code Understanding](https://mlsanity.com/paper/2307.08177/) (2023)
- [Survey of the State of the Art in Natural Language Generation: Core tasks, applications and evaluation](https://mlsanity.com/paper/1703.09902/) (2017)
- [The Carbon Footprint of Machine Learning Training Will Plateau, Then Shrink](https://mlsanity.com/paper/2204.05149/) (2022)
- [CoT-Valve: Length-Compressible Chain-of-Thought Tuning](https://mlsanity.com/paper/2502.09601/) (2025)
- [Distilling Reasoning Capabilities into Smaller Language Models](https://mlsanity.com/paper/2212.00193/) (2022)
- [MasakhaNER: Named Entity Recognition for African Languages](https://mlsanity.com/paper/2103.11811/) (2021)
- [Rho-1: Not All Tokens Are What You Need](https://mlsanity.com/paper/2404.07965/) (2024)

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

- [Learning beyond Teacher: Generalized On-Policy Distillation with Reward Extrapolation](https://mlsanity.com/paper/2602.12125/) (2026)
- [TIP: Token Importance in On-Policy Distillation](https://mlsanity.com/paper/2604.14084/) (2026)
- [Not All Disagreement Is Learnable: Token Teachability in On-Policy Distillation](https://mlsanity.com/paper/2605.26844/) (2026)
- [Uni-OPD: Unifying On-Policy Distillation with a Dual-Perspective Recipe](https://mlsanity.com/paper/2605.03677/) (2026)
- [Sculptor: Empowering LLMs with Cognitive Agency via Active Context Management](https://mlsanity.com/paper/2508.04664/) (2025)
- [SPOT: Sparse Probing and Outcome Calibration for On-Policy Distillation](https://mlsanity.com/paper/2608.04419/) (2026)
- [Revisiting Long-context Modeling from Context Denoising Perspective](https://mlsanity.com/paper/2510.05862/) (2025)
- [Less Noise, More Voice: Reinforcement Learning for Reasoning via Instruction Purification](https://mlsanity.com/paper/2601.21244/) (2026)
- [Token-Sparse Medical Multimodal Reasoning via Dual-Stream Reinforcement Learning](https://mlsanity.com/paper/2606.31599/) (2026)
- [Learning What Matters: Supervising Global Context Pruning with Causal Evidence Sets](https://mlsanity.com/paper/2607.21692/) (2026)
- [Teach Yourself Where to Look: On-Policy Attention Self-Distillation for Reasoning](https://mlsanity.com/paper/2609.33200/) (2026)
- [Unbiased Top-k Estimation for On-Policy Distillation](https://mlsanity.com/paper/2609.34447/) (2026)
- [MCD: Causal Distillation of Multimodal In-Context Learning in Large Vision-Language Models](https://mlsanity.com/paper/2609.39920/) (2026)

## Abstract (arXiv)

Large language models (LLMs) have demonstrated significant improvements in contextual understanding. However, their ability to attend to truly critical information during long-context reasoning and generation still falls behind the pace. Specifically, our preliminary experiments reveal that certain distracting patterns can misdirect the model's attention during inference, and removing these patterns substantially improves reasoning accuracy and generation quality. We attribute this phenomenon to spurious correlations in the training data, which obstruct the model's capacity to infer authentic causal instruction-response relationships. This phenomenon may induce redundant reasoning processes, potentially resulting in significant inference overhead and, more critically, the generation of erroneous or suboptimal responses. To mitigate this, we introduce a two-stage framework called Learning to Focus (LeaF) leveraging intervention-based inference to disentangle confounding factors. In the first stage, LeaF employs gradient-based comparisons with an advanced teacher to automatically identify confounding tokens based on causal relationships in the training corpus. Then, in the second stage, it prunes these tokens during distillation to enact intervention, aligning the student's attention with the teacher's focus distribution on truly critical context tokens. Experimental results demonstrate that LeaF not only achieves an absolute improvement in various mathematical reasoning, code generation and multi-hop question answering benchmarks but also effectively suppresses attention to confounding tokens during inference, yielding a more interpretable and reliable reasoning model.