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

TL;DR

This paper introduces AIR, a novel, unsupervised data selection method that uses attention head influence to improve reasoning capabilities in large language models through targeted post-training data selection.

Contribution

AIR provides a mechanistic, training-free framework that identifies critical reasoning steps and samples, enhancing reasoning distillation in LLMs beyond heuristic approaches.

Findings

AIR outperforms heuristic baselines in reasoning accuracy

It effectively isolates critical reasoning steps and samples

The method improves data efficiency in reasoning distillation

Abstract

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.