arXiv:2505.07293cs.CL2025-05被引 5

用注意力头影响筛选高质量推理数据,小模型也能提升大模型表现

AttentionInfluence: Adopting Attention Head Influence for Weak-to-Strong Pretraining Data Selection

  • 通过掩蔽注意力头计算损失差异,无监督识别关键推理头
  • 在2410亿词元语料中选出730亿词元优质数据,提升大模型性能1.4~3.5个百分点
  • 无需标注、训练自由,适合大规模推理型数据筛选场景

近期研究关注收集强推理型预训练数据以提升大模型的复杂推理能力。以往方法依赖有监督分类器识别此类数据,需人工或大模型标注,常引入领域偏见。由于注意力头对上下文推理至关重要,我们提出AttentionInfluence——一种无需监督信号、无需训练的简单有效方法。该方法使小型预训练语言模型通过简单的注意力头掩蔽操作,成为强数据选择器。具体地,我们识别检索型注意力头,并计算掩蔽这些头时的损失差异。我们将AttentionInfluence应用于13亿参数稠密模型,在包含2410亿词元的SmolLM语料上进行数据筛选,选取730亿词元子集,与原语料混合后,使用1万亿训练词元和WSD学习率调度策略,对70亿参数稠密模型进行预训练。实验结果表明,在多个知识密集型与推理型基准(MMLU、MMLU-Pro、AGIEval-en、GSM8K、HumanEval)上性能提升1.4~3.5个百分点。这验证了弱到强的扩展能力:小模型可显著提升大模型最终表现,为面向推理的数据筛选提供了一条高效且可扩展的新路径。

原文摘要 · Abstract (English)

Recently, there has been growing interest in collecting reasoning-intensive pretraining data to improve LLMs' complex reasoning ability. Prior approaches typically rely on supervised classifiers to identify such data, which requires labeling by humans or LLMs, often introducing domain-specific biases. Due to the attention heads being crucial to in-context reasoning, we propose AttentionInfluence, a simple yet effective, training-free method without supervision signal. Our approach enables a small pretrained language model to act as a strong data selector through a simple attention head masking operation. Specifically, we identify retrieval heads and compute the loss difference when masking these heads. We apply AttentionInfluence to a 1.3B-parameter dense model to conduct data selection on the SmolLM corpus of 241B tokens, and mix the SmolLM corpus with the selected subset comprising 73B tokens to pretrain a 7B-parameter dense model using 1T training tokens and WSD learning rate scheduling. Our experimental results demonstrate substantial improvements, ranging from 1.4pp to 3.5pp, across several knowledge-intensive and reasoning-heavy benchmarks (i.e., MMLU, MMLU-Pro, AGIEval-en, GSM8K, and HumanEval). This demonstrates an effective weak-to-strong scaling property, with small models improving the final performance of larger models-offering a promising and scalable path for reasoning-centric data selection.

数据筛选注意力机制大模型训练推理增强

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。