AttentionInfluence: Adozione dell'Influenza delle Teste di Attenzione per la Selezione dei Dati di Pretraining da Debole a Forte

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

May 12, 2025

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

cs.AI

Abstract

Recentemente, c'è stato un crescente interesse nel raccogliere dati di pre-addestramento ad alta intensità di ragionamento per migliorare la capacità di ragionamento complesso dei modelli linguistici di grandi dimensioni (LLM). Gli approcci precedenti si basano tipicamente su classificatori supervisionati per identificare tali dati, il che richiede etichettature da parte di esseri umani o LLM, introducendo spesso bias specifici del dominio. Poiché le testine di attenzione sono cruciali per il ragionamento in contesto, proponiamo AttentionInfluence, un metodo semplice ma efficace, privo di supervisione e senza segnali di addestramento. Il nostro approccio consente a un piccolo modello linguistico pre-addestrato di agire come un forte selezionatore di dati attraverso una semplice operazione di mascheramento delle testine di attenzione. Nello specifico, identifichiamo le testine di recupero e calcoliamo la differenza di perdita quando queste testine vengono mascherate. Applichiamo AttentionInfluence a un modello denso da 1,3 miliardi di parametri per condurre la selezione dei dati sul corpus SmolLM di 241 miliardi di token, e mescoliamo il corpus SmolLM con il sottoinsieme selezionato composto da 73 miliardi di token per pre-addestrare un modello denso da 7 miliardi di parametri utilizzando 1 trilione di token di addestramento e una pianificazione del tasso di apprendimento WSD. I nostri risultati sperimentali dimostrano miglioramenti sostanziali, che vanno da 1,4pp a 3,5pp, su diversi benchmark ad alta intensità di conoscenza e ragionamento (ad esempio, MMLU, MMLU-Pro, AGIEval-en, GSM8K e HumanEval). Ciò dimostra una proprietà efficace di scalabilità da debole a forte, con modelli piccoli che migliorano le prestazioni finali di modelli più grandi, offrendo un percorso promettente e scalabile per la selezione di dati centrata sul ragionamento.

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.