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

TL;DR

AttentionInfluence introduces a training-free data selection method that leverages attention head influence by masking retrieval heads to create a weak-to-strong model pair. It identifies high-utility pretraining samples from a 241B token corpus by scoring delta losses between a base 1.3B model and a degraded reference, selecting top 20% (≈73B tokens) for pretraining a 7B model. Pretraining with the selected data yields consistent improvements across MMLU, MMLU-Pro, AGIEval-en, GSM8K, HumanEval, and other benchmarks, particularly on reasoning-intensive tasks, demonstrating a weak-to-strong scaling effect. The results suggest that intrinsic model mechanisms, rather than supervision signals, can guide scalable data selection with broad distribution and improved reasoning capabilities.

Abstract

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.

Figures (19)

• Figure 1: (a) Performance evolution on comprehensive benchmark evaluations during pretraining. The first 750 billion tokens correspond to the pretraining phase, represented by solid lines, while the subsequent 250 billion tokens represent the learning rate annealing phase, represented by dashed lines, using the same dataset. After around 100 billion tokens, AttentionInfluence-1.3B consistently outperforms the baseline across a wide range of tasks on average, including the annealing phase. (b) Training Loss during pretraining. AttentionInfluence-1.3B consistently achieves a lower loss than the baseline.
• Figure 2: The evolution of retrieval heads in a 1.3B dense model.
• Figure 3: The illustration of AttentionInfluence.
• Figure 4: The statistics of clustering. The left is the clustering result of AttentionInfluence, the right part is that of FineWeb-Edu Classifier.
• Figure 5: Visualization of data selected by AttentionInfluence and FineWeb-Edu Classifier.