AttentionInfluence：基于注意力头影响力的弱到强预训练数据选择
计算与语言
2025-05-13 v1
摘要
最近，越来越多的人关注收集富有推理能力的预训练数据以提高 LLM 的复杂推理能力。以往的方法通常依赖监督分类器来识别此类数据，这需要人类或 LLM 的标记，常常引入特定于领域的偏见。鉴于注意力头对 in-context 推理至关重要，我们提出了 AttentionInfluence，这是一种简单而有效、无需监督信号的 training-free 方法。我们的方法使小型预训练语言模型能够通过简单的注意力头掩码操作充当强大的数据选择器。具体而言，我们识别了 retrieval heads 并计算掩盖这些 heads 时的 loss difference。我们将 AttentionInfluence 应用于 13 亿参数的稠密模型，对 2410 亿 token 的 SmolLM 语料库进行数据选择，并将 SmolLM 语料库与包含 73 亿 token 的所选子集混合，以 1T training tokens 和 WSD 学习率调度预训练 7B 参数稠密模型。我们的实验结果在多个 knowledge-intensive 和 reasoning-heavy 基准测试（即 MMLU、MMLU-Pro、AGIEval-en、GSM8K 和 HumanEval）上均实现了 1.4pp 到 3.5pp 的显著提升。这表明了小模型提升大模型最终性能的有效的弱到强 scaling property，为推进以推理为中心的数据选择提供了有前景且可扩展的路径。
引用
@article{arxiv.2505.07293,
title = {AttentionInfluence: Adopting Attention Head Influence for Weak-to-Strong Pretraining Data Selection},
author = {Kai Hua and Steven Wu and Ge Zhang and Ke Shen},
journal= {arXiv preprint arXiv:2505.07293},
year = {2025}
}
备注
28 pages, 19 figures