AttentionInfluence: Het Adopteren van Attention Head Invloed voor Zwak-naar-Sterk Voorbereidende Data Selectie

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

Samenvatting

Recentelijk is er een groeiende interesse ontstaan in het verzamelen van trainingsdata die intensief redeneren vereist, om het complexe redeneervermogen van grote taalmodelen (LLMs) te verbeteren. Eerdere benaderingen vertrouwen doorgaans op begeleide classificatiemodellen om dergelijke data te identificeren, wat labeling door mensen of LLMs vereist en vaak domeinspecifieke vooroordelen introduceert. Gezien de cruciale rol van aandachtskoppen (attention heads) bij in-context redeneren, stellen we AttentionInfluence voor, een eenvoudige maar effectieve, trainingsvrije methode zonder begeleidingssignaal. Onze aanpak stelt een klein voorgetraind taalmodel in staat om als een sterke dataselector te fungeren door middel van een eenvoudige maskeringsoperatie van aandachtskoppen. Specifiek identificeren we retrieval heads en berekenen we het verschil in verlies bij het maskeren van deze koppen. We passen AttentionInfluence toe op een dicht model met 1,3 miljard parameters om dataselectie uit te voeren op het SmolLM-corpus van 241 miljard tokens, en mengen het SmolLM-corpus met de geselecteerde subset van 73 miljard tokens om een dicht model met 7 miljard parameters voor te trainen met 1 biljoen trainings-tokens en WSD-leersnelheidsplanning. Onze experimentele resultaten laten aanzienlijke verbeteringen zien, variërend van 1,4 tot 3,5 procentpunt, over verschillende kennisintensieve en redeneerzware benchmarks (zoals MMLU, MMLU-Pro, AGIEval-en, GSM8K en HumanEval). Dit toont een effectieve zwak-naar-sterk schaaleigenschap aan, waarbij kleine modellen de uiteindelijke prestaties van grotere modellen verbeteren—wat een veelbelovende en schaalbare weg biedt voor redeneringsgerichte dataselectie.

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.