# AIR: Post-training Data Selection for Reasoning via Attention Head Influence

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

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Placed in Test-time compute and efficient thinking, one of 17 directions in LLM reasoning.
- Builds on 39 of its 41 references that are in the map.
- First posted on arXiv 15 Dec 2025.

## Placed in

- LLM reasoning / [Test-time compute and efficient thinking](https://mlsanity.com/lr/test-time-compute-and-efficient-thinking/)

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

- [Training language models to follow instructions with human feedback](https://mlsanity.com/paper/2203.02155/) (2022)
- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Qwen3 Technical Report](https://mlsanity.com/paper/2505.09388/) (2025)
- [Measuring Mathematical Problem Solving With the MATH Dataset](https://mlsanity.com/paper/2103.03874/) (2021)
- [Qwen2.5 Technical Report](https://mlsanity.com/paper/2412.15115/) (2024)
- [GPQA: A Graduate-Level Google-Proof Q&A Benchmark](https://mlsanity.com/paper/2311.12022/) (2023)
- [Fine-Tuning Language Models from Human Preferences](https://mlsanity.com/paper/1909.08593/) (2019)
- [s1: Simple test-time scaling](https://mlsanity.com/paper/2501.19393/) (2025)
- [Scaling Language Models: Methods, Analysis & Insights from Training Gopher](https://mlsanity.com/paper/2112.11446/) (2021)
- [In-context Learning and Induction Heads](https://mlsanity.com/paper/2209.11895/) (2022)
- [Enhancing Chat Language Models by Scaling High-quality Instructional Conversations](https://mlsanity.com/paper/2305.14233/) (2023)
- [Beyond the 80/20 Rule: High-Entropy Minority Tokens Drive Effective Reinforcement Learning for LLM Reasoning](https://mlsanity.com/paper/2506.01939/) (2025)
- [LIMO: Less is More for Reasoning](https://mlsanity.com/paper/2502.03387/) (2025)
- [Data Selection for Language Models via Importance Resampling](https://mlsanity.com/paper/2302.03169/) (2023)
- [From Quantity to Quality: Boosting LLM Performance with Self-Guided Data Selection for Instruction Tuning](https://mlsanity.com/paper/2308.12032/) (2023)
- [OpenThoughts: Data Recipes for Reasoning Models](https://mlsanity.com/paper/2506.04178/) (2025)
- [Retrieval Head Mechanistically Explains Long-Context Factuality](https://mlsanity.com/paper/2404.15574/) (2024)
- [QuRating: Selecting High-Quality Data for Training Language Models](https://mlsanity.com/paper/2402.09739/) (2024)
- [LLM Post-Training: A Deep Dive into Reasoning Large Language Models](https://mlsanity.com/paper/2502.21321/) (2025)
- [Rho-1: Not All Tokens Are What You Need](https://mlsanity.com/paper/2404.07965/) (2024)
- [MATES: Model-Aware Data Selection for Efficient Pretraining with Data Influence Models](https://mlsanity.com/paper/2406.06046/) (2024)
- [Not All Heads Matter: A Head-Level KV Cache Compression Method with Integrated Retrieval and Reasoning](https://mlsanity.com/paper/2410.19258/) (2024)
- [Perplexed by Perplexity: Perplexity-Based Data Pruning With Small Reference Models](https://mlsanity.com/paper/2405.20541/) (2024)
- [Organize the Web: Constructing Domains Enhances Pre-Training Data Curation](https://mlsanity.com/paper/2502.10341/) (2025)
- [Procedural Knowledge in Pretraining Drives Reasoning in Large Language Models](https://mlsanity.com/paper/2411.12580/) (2024)
- [Stabilizing Knowledge, Promoting Reasoning: Dual-Token Constraints for RLVR](https://mlsanity.com/paper/2507.15778/) (2025)
- [NaturalThoughts: Selecting and Distilling Reasoning Traces for General Reasoning Tasks](https://mlsanity.com/paper/2507.01921/) (2025)
- [Reverse-Engineered Reasoning for Open-Ended Generation](https://mlsanity.com/paper/2509.06160/) (2025)
- [Interpreting Key Mechanisms of Factual Recall in Transformer-Based Language Models](https://mlsanity.com/paper/2403.19521/) (2024)
- [DataMan: Data Manager for Pre-training Large Language Models](https://mlsanity.com/paper/2502.19363/) (2025)

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

- [Which Reasoning Trajectories Teach Students to Reason Better? A Simple Metric of Informative Alignment](https://mlsanity.com/paper/2601.14249/) (2026)
- [DeltaPrompts: Escaping the Zero-Delta Trap in Multimodal Distillation](https://mlsanity.com/paper/2605.15532/) (2026)
- [SelectTTRL: Representation-Guided Data Selection for Efficient Test-Time Reinforcement Learning](https://mlsanity.com/paper/2610.16104/) (2026)

## Abstract (arXiv)

LLMs achieve remarkable multi-step reasoning capabilities, yet effectively transferring these skills via post-training distillation remains challenging. Existing data selection methods, ranging from manual curation to heuristics based on length, entropy, or overall loss, fail to capture the causal importance of individual reasoning steps, limiting distillation efficiency. To address this, we propose Attention Influence for Reasoning (AIR), a principled, unsupervised and training-free framework that leverages mechanistic insights of the retrieval head to select high-value post-training data. AIR first identifies reasoning-critical attention heads of an off-the-shelf model, then constructs a weakened reference model with disabled head influence, and finally quantifies the resulting loss divergence as the Attention Influence Score. This score enables fine-grained assessment at both the step and sample levels, supporting step-level weighted fine-tuning and global sample selection. Experiments across multiple reasoning benchmarks show that AIR consistently improves reasoning accuracy, surpassing heuristic baselines and effectively isolating the most critical steps and samples. Our work establishes a mechanism-driven, data-efficient approach for reasoning distillation in LLMs.