本文提出AttentionInfluence方法，通过无监督地利用预训练模型注意力头机制选择推理密集型数据，显著提升了7B参数模型在知识和推理任务上的性能，展现了弱到强的扩展潜力。
Large Language Model, Pre-training, Reasoning, Data Augmentation, Efficiency
Kai Hua, Steven Wu, Ge Zhang, Ke Shen
ByteDance Seed
Generated by grok-3
Background Problem
大型语言模型（LLMs）的预训练数据质量对模型性能至关重要，尤其是在提升复杂推理能力方面。传统方法通常依赖人工标注或大型模型生成的监督信号来训练分类器以筛选高质量数据，但这种方式成本高昂且可能引入领域偏见，限制了数据的多样性。本文提出一个关键问题：如何高效、可扩展地识别多样化的高质量预训练数据，特别是推理密集型数据？为此，作者提出了AttentionInfluence方法，试图通过利用预训练模型内部的注意力头机制，在无监督、无训练成本的情况下解决这一问题。
Method
AttentionInfluence是一种无需训练、无监督信号的预训练数据选择方法，其核心思想是利用预训练语言模型中注意力头的激活模式来评估数据的推理强度。具体步骤如下：
- 检测重要注意力头：基于现有研究，作者聚焦于检索头（Retrieval Heads），认为其与模型的检索和推理能力密切相关。采用一个合成测试数据集（包含800个样本的3-shot检索任务），计算每个注意力头的检索分数，并选择排名前5%的头作为重要头。
- 计算AttentionInfluence分数：通过对基础模型（Base Model）的重要注意力头进行掩码操作（即将注意力权重设为均匀分布），得到一个能力较弱的参考模型（Reference Model）。随后，计算每个数据样本在基础模型和参考模型上的损失差值，定义为AttentionInfluence分数，分数越高表示数据样本的推理强度越高。
- 数据选择：在同一领域内，根据AttentionInfluence分数对数据样本排序，选择分数最高的样本（例如前20%）作为高质量预训练数据。
关键创新：该方法无需额外训练或标注，依赖模型内部机制进行数据选择，降低了成本并提高了可扩展性。
潜在问题：首先，检索头的选择基于特定任务和假设，但论文未充分验证这一假设是否适用于所有推理场景或不同模型架构。其次，掩码操作可能导致模型行为不可预测，损失差值的解释力可能因领域不同而变化，论文未提供理论上的严谨支持。此外，方法对领域内数据的依赖可能限制其在跨领域数据选择中的表现。
Experiment
实验中，作者使用一个1.3B参数的LLaMA2类似模型对SmolLM-Corpus（241B token）进行数据选择，筛选出约73.1B token的推理密集型数据，并将其与完整SmolLM-Corpus混合，用于预训练一个7B参数模型（训练token总量为1T）。对比实验包括使用完整SmolLM-Corpus预训练的基线模型，以及使用FineWeb-Edu分类器选择数据的对比。
- 数据集与设置：SmolLM-Corpus已通过教育聚焦的分类器（FineWeb-Edu-Dedup）进行强质量过滤，本身已是高质量数据集。评估基准涵盖知识密集和推理重型任务（如MMLU、MMLU-Pro、GSM8K、HumanEval等），采用few-shot设置。
- 结果：AttentionInfluence在多个基准上显著优于基线，MMLU提升1.4pp，MMLU-Pro提升2.7pp，GSM8K提升2.7pp，HumanEval提升3.5pp，表明其在推理和知识任务上的有效性。训练动态显示，在约100B token后，方法即展现出持续优势。此外，使用7B模型进行数据选择进一步提升了性能，特别是在推理密集任务上。
- 数据分布分析：通过聚类和词频分析，AttentionInfluence选择的数据在类别分布上更均衡，且倾向于选择更长、更复杂的样本（如包含详细代码或公式推理的内容）。
- 评估合理性与局限：实验设置较为全面，涵盖了多个任务类型和训练阶段，但依赖GPT-4o进行数据质量评分可能引入主观偏差，且未与其他无监督数据选择方法（如困惑度过滤）进行对比。此外，实验未探讨方法在更大规模模型或长文本数据上的表现，计算成本的增加可能限制其实际应用。总体而言，实验结果支持了方法的有效性，但未能完全排除数据选择中的潜在偏见或过度拟合特定任务的可能性。
Further Thoughts
AttentionInfluence提供了一个有趣的视角，即利用模型内部机制进行数据选择，这可能启发未来在其他模型架构（如CNN或GNN）中探索类似的无监督数据筛选方法。然而，方法对注意力头的依赖可能使其对Transformer架构过于特化，限制了普适性。进一步研究可以探索结合多头注意力与MLP机制的综合影响，以更全面地捕捉数据质量信号。此外，论文中提到的与FineWeb-Edu分类器的互补性值得深入挖掘，或许可以通过混合策略（例如将AttentionInfluence与困惑度过滤结合）进一步提升数据选择的多样性和质量。另一个值得关注的领域是长文本处理，特别是在长上下文推理任务中，AttentionInfluence是否能有效识别长距离依赖数据仍是一个开放问题。最后，考虑到计算成本随模型规模增加而显著上升，未来的优化方向可能包括设计更高效的注意力头检测算法或探索分布式计算框架，以支持更大规模的预训练数据筛选。