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

ICLR 2026 Conference Submission12520 Authors

Keywords: data selection, mechanistic interpretability, influence measure, unsupervised learning, large language model

TL;DR: We propose AttentionInfluence, a training-free method that uses a small pretrained model to select reasoning-intensive data by masking attention heads, boosting larger LLMs’ performance.

Abstract: Recently, there has been growing interest in collecting reasoning-intensive pretraining data to improve the reasoning ability of LLMs. Prior approaches typically rely on supervised classifiers to identify such data, requiring labeling by humans or LLMs, often introducing domain-specific biases. Since attention heads are crucial to in-context reasoning, we propose \textbf{AttentionInfluence}, a simple yet effective, \textbf{training-free} method \textbf{without supervision signal}. Our approach enables a \textbf{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 incurred by masking them. We apply AttentionInfluence to a 1.3B-parameter dense model to conduct data selection on the SmolLM corpus of 241B tokens, and mix the corpus with the selected subset comprising 73B tokens to pretrain a 7B-parameter dense model using 1T training tokens and the Warmup-Stable-Decay (WSD) learning rate schedule. Experimental results demonstrate substantial improvements, ranging from \textbf{1.4pp} to \textbf{3.5pp}, across several knowledge-intensive and reasoning-heavy benchmarks (i.e., MMLU, MMLU-Pro, AGIEval-en, GSM8K, and HumanEval). This demonstrates an effective \textbf{Weak-to-Strong} scaling property, with small models improving the performance of larger models---offering a promising and scalable path for reasoning-centric data selection.

Primary Area: foundation or frontier models, including LLMs

Submission Number: 12520