AttentionInfluence: Übernahme der Aufmerksamkeitskopfeinflüsse für die schwache bis starke Vorauswahl von Trainingsdaten

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

May 12, 2025

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

cs.AI

Zusammenfassung

In letzter Zeit hat das Interesse an der Sammlung von vortrainierten Daten, die eine hohe Anforderung an das logische Denken stellen, zugenommen, um die komplexen Denkfähigkeiten von LLMs (Large Language Models) zu verbessern. Bisherige Ansätze stützen sich typischerweise auf überwachte Klassifikatoren, um solche Daten zu identifizieren, was eine Kennzeichnung durch Menschen oder LLMs erfordert und oft domänenspezifische Verzerrungen einführt. Da die Aufmerksamkeitsköpfe (Attention Heads) für das kontextbezogene Denken entscheidend sind, schlagen wir AttentionInfluence vor, eine einfache, aber effektive, trainingsfreie Methode ohne Überwachungssignal. Unser Ansatz ermöglicht es einem kleinen vortrainierten Sprachmodell, durch eine einfache Maskierung der Aufmerksamkeitsköpfe als starker Datenauswähler zu fungieren. Konkret identifizieren wir Retrieval-Köpfe und berechnen den Verlustunterschied, wenn diese Köpfe maskiert werden. Wir wenden AttentionInfluence auf ein dichtes Modell mit 1,3 Milliarden Parametern an, um eine Datenauswahl im SmolLM-Korpus mit 241 Milliarden Tokens durchzuführen, und mischen das SmolLM-Korpus mit der ausgewählten Teilmenge von 73 Milliarden Tokens, um ein dichtes Modell mit 7 Milliarden Parametern unter Verwendung von 1 Billion Trainings-Tokens und WSD-Lernratenplanung vortrainieren. Unsere experimentellen Ergebnisse zeigen erhebliche Verbesserungen im Bereich von 1,4 bis 3,5 Prozentpunkten über mehrere wissensintensive und denkintensive Benchmarks hinweg (d. h. MMLU, MMLU-Pro, AGIEval-en, GSM8K und HumanEval). Dies demonstriert eine effektive schwache-zu-starke Skalierungseigenschaft, bei der kleine Modelle die Endleistung größerer Modelle verbessern – ein vielversprechender und skalierbarer Weg für die datenzentrierte Auswahl im Hinblick auf das logische Denken.

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.