通过注意力头影响评估,精准挑选推理训练数据
AIR: Post-training Data Selection for Reasoning via Attention Head Influence
- 基于检索头机制识别关键推理步骤
- 损失差异量化得分,提升推理准确率12%以上
- 无需训练,适合高效微调与数据筛选
大语言模型具备出色的多步推理能力,但通过后训练蒸馏有效迁移这些能力仍具挑战。现有数据选择方法(如人工筛选或基于长度、熵、整体损失的启发式方法)无法捕捉单个推理步骤的因果重要性,限制了蒸馏效率。为此,我们提出注意力头影响推理(AIR),一种原理性、无监督且无需训练的框架,利用检索头的机制洞察来选择高价值后训练数据。AIR首先识别现成模型中对推理至关重要的注意力头,然后构建禁用头影响的弱化参考模型,最后以产生的损失差异作为注意力影响评分。该评分可在步骤和样本层面进行细粒度评估,支持步骤加权微调和全局样本选择。在多个推理基准测试中,AIR持续提升推理准确率,优于启发式基线,并能有效识别最关键的步骤与样本。本工作建立了一种机制驱动、数据高效的推理蒸馏方法。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。