Target-Oriented Pretraining Data Selection via Neuron-Activated Graph

1. El Problema

El preentrenamiento de Modelos de Lenguaje Grandes (LLMs) es costoso y requiere grandes volúmenes de datos. Aunque la selección de datos de alta calidad es crucial para mejorar el rendimiento, los enfoques actuales presentan dos limitaciones principales:

• Selección basada en calidad general: Métodos existentes (como FineWeb-Edu) seleccionan datos basándose en heurísticas de calidad general o puntuaciones de "educatividad". Sin embargo, estos criterios a menudo no se alinean con las capacidades específicas que se desea adquirir (ej. razonamiento matemático vs. comprensión narrativa), dejando un vacío entre la calidad del dato y la utilidad para una tarea concreta.
• Selección orientada a objetivos "caja negra": Métodos recientes que intentan alinear los datos con tareas específicas (como BETR) utilizan representaciones internas del modelo (embeddings o pérdidas) distiladas en clasificadores auxiliares. Estos enfoques son opacos (caja negra), carecen de interpretabilidad y a menudo dependen de similitudes superficiales en las capas finales, lo que puede capturar estilo en lugar de la lógica subyacente de la tarea.

2. Metodología: NAG-based Ranking

Los autores proponen NAG-based Ranking (Ranking basado en Grafos Activados por Neuronas), un marco sin entrenamiento adicional (training-free) e interpretable para la selección de datos. La metodología se basa en la premisa de que diferentes tareas activan subconjuntos dispersos y específicos de neuronas dentro del modelo.

El proceso consta de tres pasos principales:

A. Cuantificación del Impacto de la Neurona

Para una entrada dada en un LLM existente (modelo "off-the-shelf"), se calcula el impacto de cada neurona.

• Se define una neurona como una columna de una matriz de proyección (en capas de Atención Q, K, V y en el módulo FFN UP/DOWN).
• El impacto se mide localmente desactivando la neurona (poniendo su columna en cero) y calculando la magnitud del cambio en la salida de esa capa específica:
  $$Imp(N_k | h_{in}) = ||h_{in}^T W - h_{in}^T (W \setminus N_k)||_2 = |h_{in}^T W_{:,k}|$$
• Las neuronas con mayor impacto se consideran "activadas" para esa entrada.

B. Construcción del Grafo Activado por Neuronas (NAG)

En lugar de usar embeddings densos, se construye una representación estructurada y dispersa:

• Para cada capa del modelo, se seleccionan las K neuronas de mayor impacto.
• El NAG de una entrada es la colección de estos índices de neuronas a través de todas las capas del modelo.
• Esto crea un perfil de "esqueleto funcional" que describe cómo el modelo procesa internamente la información.

C. Ranking y Selección de Datos

• Perfil Objetivo: Se toma un pequeño conjunto de ejemplos objetivo ($D_{target}$) y se agrega sus NAGs para crear un perfil de activación objetivo.
• Similitud: Se calcula la similitud entre el NAG de cada candidato en el pool de datos ($D_{pool}$) y el perfil objetivo. La similitud se basa en la superposición de neuronas activadas (similar al coeficiente de Dice).
• Selección: Se seleccionan los datos candidatos con mayor similitud al perfil objetivo para el preentrenamiento.

3. Contribuciones Clave

1. Marco Interpretable y Sin Entrenamiento: A diferencia de los métodos que requieren entrenar clasificadores auxiliares, NAG utiliza señales directas e interpretables de neuronas en modelos preexistentes.
2. Identificación del "Esqueleto Funcional": Demuestran que las tareas dependen de un subconjunto extremadamente disperso de neuronas (aprox. 0.12% del total). Desactivar estas neuronas seleccionadas por NAG provoca un colapso del rendimiento del 23.5%, mientras que desactivar neuronas aleatorias tiene un efecto negligible.
3. Algoritmo de Selección Multi-Objetivo Robusto: El método mantiene su eficacia en escenarios de múltiples objetivos (mezcla de tareas), donde otros métodos basados en embeddings tienden a degradarse.
4. Análisis de Arquitectura: Revelan que la señal relevante para la tarea está distribuida a través de todas las capas del modelo, no solo en la última. Restringir NAG a la última capa reduce el rendimiento en un 4.1% promedio. Además, las neuronas de la capa de proyección FFN UP son las más informativas para la selección.

4. Resultados Experimentales

Los autores evaluaron el método en seis benchmarks (ARC-Challenge, HellaSwag, TriviaQA, MMLU, XStoryCloze, XWinograd) utilizando modelos base de 1.2B a 7B parámetros.

• Rendimiento General: NAG mejora el preentrenamiento orientado a objetivos en un 4.9% en promedio en comparación con el muestreo aleatorio.
• Comparación con Baselines:
  • Supera a FineWeb-Edu (selección por calidad general) en un 2.4% promedio.
  • Supera a BETR (estado del arte en selección orientada a objetivos) en un 5.3% de precisión en HellaSwag y un 1% en promedio.
• Escenarios Multi-Objetivo: En configuraciones donde se mezclan múltiples tareas, NAG supera a los baselines en 1.1% y 4.1% respectivamente, demostrando una mejor generalización.
• Eficiencia Computacional: El método actúa como un "multiplicador de cómputo". Para alcanzar el mismo nivel de precisión, NAG requiere significativamente menos tokens de entrenamiento (hasta 3.7x más eficiente en XStoryCloze) que los métodos aleatorios o baselines.
• Robustez: Los resultados son consistentes independientemente del modelo "extractor" utilizado (Qwen, Llama, SmolLM) y funcionan incluso si el modelo extractor es más pequeño que el modelo que se va a entrenar.

5. Significado e Impacto

Este trabajo representa un cambio de paradigma en la selección de datos para LLMs:

• De lo superficial a lo funcional: Se aleja de la selección basada en similitud semántica superficial (embeddings) o heurísticas de calidad, acercándose a la mecánica interna del modelo.
• Interpretabilidad: Proporciona una ventana clara a qué partes del modelo son responsables de aprender ciertas capacidades, permitiendo una ingeniería de datos más precisa.
• Escalabilidad: Al ser un método "training-free" que solo requiere un paso forward en un modelo existente, es altamente escalable y económico de implementar en pipelines de preentrenamiento masivo.
• Validación Teórica: Confirma la hipótesis de que las capacidades de los LLMs están localizadas en subconjuntos de neuronas altamente específicos y dispersos, ofreciendo una nueva vía para "dirigir" (steering) el aprendizaje de modelos grandes mediante la curación de datos basada en la arquitectura neuronal.

En resumen, NAG-based Ranking demuestra que alinear los datos de preentrenamiento con la "firma neuronal" específica de una tarea objetivo es una estrategia superior para construir modelos expertos en dominios específicos, superando tanto a la selección aleatoria como a los métodos de calidad general o caja negra.