AttentionInfluence: Adotando a Influência das Cabeças de Atenção para Seleção de Dados de Pré-treinamento de Fraco para Forte

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

May 12, 2025

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

cs.AI

Resumo

Recentemente, tem havido um interesse crescente em coletar dados de pré-treinamento intensivos em raciocínio para melhorar a capacidade de raciocínio complexo dos LLMs (Large Language Models). Abordagens anteriores geralmente dependem de classificadores supervisionados para identificar esses dados, o que requer rotulação por humanos ou LLMs, frequentemente introduzindo vieses específicos de domínio. Devido à importância das cabeças de atenção para o raciocínio em contexto, propomos o AttentionInfluence, um método simples, mas eficaz, sem sinal de supervisão e que não requer treinamento. Nossa abordagem permite que um pequeno modelo de linguagem pré-treinado atue como um forte seletor de dados por meio de uma simples operação de mascaramento das cabeças de atenção. Especificamente, identificamos as cabeças de recuperação e calculamos a diferença de perda ao mascarar essas cabeças. Aplicamos o AttentionInfluence a um modelo denso de 1,3 bilhão de parâmetros para realizar a seleção de dados no corpus SmolLM de 241 bilhões de tokens, e misturamos o corpus SmolLM com o subconjunto selecionado, composto por 73 bilhões de tokens, para pré-treinar um modelo denso de 7 bilhões de parâmetros usando 1 trilhão de tokens de treinamento e agendamento de taxa de aprendizado WSD. Nossos resultados experimentais demonstram melhorias substanciais, variando de 1,4pp a 3,5pp, em vários benchmarks intensivos em conhecimento e pesados em raciocínio (ou seja, MMLU, MMLU-Pro, AGIEval-en, GSM8K e HumanEval). Isso demonstra uma propriedade eficaz de escalonamento de fraco para forte, com modelos pequenos melhorando o desempenho final de modelos maiores, oferecendo um caminho promissor e escalável para a seleção de dados centrada em raciocínio.

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.