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

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.

We teach artificial intelligence to solve complex reasoning tasks, like math or coding, by training them on data that details step-by-step thinking processes. However, it is traditionally difficult to know which specific step in a long reasoning chain is truly critical for the model's understanding. We discovered that AI models rely on specialized internal components called "retrieval heads" to look back at key information during reasoning. By intentionally disabling these heads, we created a weakened reference model. Comparing the performance gap between the normal and weakened models allowed us to precisely isolate which problems or specific reasoning steps are the most valuable. Surprisingly, training a model on just 1,000 high-value samples selected by our method outperformed a baseline trained blindly on 800,000 samples. To help other researchers, we are releasing this unsupervised data selection framework, called AIR, to make reasoning training much more efficient and accessible.