AttentionInfluence: Adopción de la Influencia de Cabezas de Atención para la Selección de Datos de Pretrenamiento de Débil a Fuerte

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

Resumen

Recientemente, ha habido un creciente interés en recopilar datos de preentrenamiento intensivos en razonamiento para mejorar la capacidad de razonamiento complejo de los modelos de lenguaje grandes (LLMs). Los enfoques anteriores generalmente dependen de clasificadores supervisados para identificar dichos datos, lo que requiere etiquetado por humanos o LLMs, introduciendo a menudo sesgos específicos del dominio. Dado que las cabezas de atención son cruciales para el razonamiento en contexto, proponemos AttentionInfluence, un método simple pero efectivo, sin señal de supervisión y que no requiere entrenamiento. Nuestro enfoque permite que un pequeño modelo de lenguaje preentrenado actúe como un selector de datos robusto mediante una simple operación de enmascaramiento de cabezas de atención. Específicamente, identificamos las cabezas de recuperación y calculamos la diferencia de pérdida al enmascarar estas cabezas. Aplicamos AttentionInfluence a un modelo denso de 1.3B parámetros para realizar la selección de datos en el corpus SmolLM de 241B tokens, y mezclamos el corpus SmolLM con el subconjunto seleccionado que comprende 73B tokens para preentrenar un modelo denso de 7B parámetros utilizando 1T tokens de entrenamiento y programación de tasa de aprendizaje WSD. Nuestros resultados experimentales demuestran mejoras sustanciales, que van desde 1.4pp hasta 3.5pp, en varios puntos de referencia intensivos en conocimiento y razonamiento (es decir, MMLU, MMLU-Pro, AGIEval-en, GSM8K y HumanEval). Esto demuestra una propiedad efectiva de escalado de débil a fuerte, donde modelos pequeños mejoran el rendimiento final de modelos más grandes, ofreciendo un camino prometedor y escalable para la selección de datos centrada en el razonamiento.

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.