AttentionInfluence: Использование влияния голов внимания для выбора данных предварительного обучения от слабого к сильному

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

May 12, 2025

Авторы: Kai Hua, Steven Wu, Ge Zhang, Ke Shen

cs.AI

Аннотация

В последнее время наблюдается растущий интерес к сбору данных для предварительного обучения, требующих интенсивного рассуждения, с целью улучшения способности крупных языковых моделей (LLM) к сложным умозаключениям. Предыдущие подходы обычно полагались на контролируемые классификаторы для идентификации таких данных, что требует аннотирования людьми или LLM, что часто приводит к внесению доменно-специфических предубеждений. Поскольку механизмы внимания играют ключевую роль в контекстном рассуждении, мы предлагаем AttentionInfluence — простой, но эффективный метод, не требующий обучения и не использующий сигналы контроля. Наш подход позволяет небольшой предварительно обученной языковой модели выступать в роли мощного селектора данных с помощью простой операции маскирования механизмов внимания. В частности, мы идентифицируем механизмы извлечения и вычисляем разницу в потере при маскировании этих механизмов. Мы применяем AttentionInfluence к плотной модели с 1,3 миллиардами параметров для отбора данных в корпусе SmolLM, содержащем 241 миллиард токенов, и смешиваем корпус SmolLM с выбранным подмножеством из 73 миллиардов токенов для предварительного обучения плотной модели с 7 миллиардами параметров, используя 1 триллион токенов для обучения и планирование скорости обучения WSD. Наши экспериментальные результаты демонстрируют значительные улучшения, варьирующиеся от 1,4 до 3,5 процентных пунктов, на нескольких тестах, требующих интенсивных знаний и сложных рассуждений (например, MMLU, MMLU-Pro, AGIEval-en, GSM8K и HumanEval). Это демонстрирует эффективное свойство масштабирования от слабых к сильным моделям, где небольшие модели улучшают итоговую производительность более крупных моделей, предлагая перспективный и масштабируемый путь для отбора данных, ориентированного на рассуждения.

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.