AttentionInfluence : Adoption de l'influence des têtes d'attention pour la sélection de données d'entraînement faible à fort

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

May 12, 2025

Auteurs: Kai Hua, Steven Wu, Ge Zhang, Ke Shen

cs.AI

Résumé

Récemment, l’intérêt pour la collecte de données de pré-entraînement intensives en raisonnement afin d’améliorer les capacités de raisonnement complexe des modèles de langage de grande taille (LLMs) a considérablement augmenté. Les approches précédentes reposent généralement sur des classificateurs supervisés pour identifier ces données, ce qui nécessite un étiquetage par des humains ou des LLMs, introduisant souvent des biais spécifiques au domaine. Étant donné que les têtes d’attention sont cruciales pour le raisonnement en contexte, nous proposons AttentionInfluence, une méthode simple mais efficace, sans signal de supervision et ne nécessitant pas d’entraînement. Notre approche permet à un petit modèle de langage pré-entraîné d’agir comme un sélecteur de données performant grâce à une simple opération de masquage des têtes d’attention. Plus précisément, nous identifions les têtes de récupération et calculons la différence de perte lors du masquage de ces têtes. Nous appliquons AttentionInfluence à un modèle dense de 1,3 milliard de paramètres pour effectuer une sélection de données sur le corpus SmolLM de 241 milliards de tokens, puis mélangeons le corpus SmolLM avec le sous-ensemble sélectionné comprenant 73 milliards de tokens pour pré-entraîner un modèle dense de 7 milliards de paramètres en utilisant 1 trillion de tokens d’entraînement et une planification du taux d’apprentissage WSD. Nos résultats expérimentaux montrent des améliorations substantielles, allant de 1,4 à 3,5 points de pourcentage, sur plusieurs benchmarks intensifs en connaissances et en raisonnement (c’est-à-dire MMLU, MMLU-Pro, AGIEval-en, GSM8K et HumanEval). Cela démontre une propriété efficace de mise à l’échelle faible à forte, où les petits modèles améliorent les performances finales des modèles plus grands, offrant ainsi une voie prometteuse et évolutive pour la sélection de données centrée sur le raisonnement.

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.