Data Mixing Agent: Learning to Re-weight Domains for Continual Pre-training

1. Problema: El Dilema del Olvido Catastrófico en el Pre-entrenamiento Continuo

El pre-entrenamiento continuo (continual pre-training) es una estrategia efectiva para adaptar Grandes Modelos de Lenguaje (LLMs) a nuevos campos específicos (como razonamiento matemático o generación de código). Sin embargo, adaptar directamente un modelo a datos de un campo objetivo suele provocar un olvido catastrófico de sus capacidades originales (campo fuente) debido a los cambios significativos en la distribución de los datos.

La solución común es mezclar datos del campo fuente y del campo objetivo. No obstante, las estrategias existentes para re-pesar los dominios (determinar la proporción de datos de cada dominio durante el entrenamiento) dependen de:

• Designaciones manuales basadas en intuición humana.
• Heurísticas empíricas limitadas.
• Métodos estáticos o dinámicos que no aprenden heurísticas generalizables complejas.

Esto resulta en un equilibrio subóptimo entre el rendimiento en el nuevo campo y la preservación de las capacidades generales.

2. Metodología: El Agente de Mezcla de Datos (Data Mixing Agent)

Los autores proponen Data Mixing Agent, el primer marco basado en modelos y de extremo a extremo que aprende a re-pesar dominios automáticamente. El enfoque trata el problema de re-pesaje de dominios como un Proceso de Decisión de Markov (MDP) y utiliza el Aprendizaje por Refuerzo (RL) para encontrar heurísticas óptimas.

Componentes Clave del Marco:

• Formulación MDP:
  • Estado ($S$): La distribución de datos de los pasos de re-pesaje anteriores.
  • Acción ($A$): La distribución de probabilidad actual sobre los dominios (qué porcentaje de datos de cada dominio usar).
  • Recompensa ($R$): Evaluación del rendimiento del modelo en un entorno de evaluación ligero (benchmarks específicos).
• Muestreo de Trayectorias:
  • Se generan miles de trayectorias de mezcla de datos aleatorias.
  • Se entrenan modelos proxy pequeños (50M parámetros) en estas trayectorias desde cero.
  • En cada paso de re-pesaje, estos modelos se evalúan en un entorno ligero (usando subconjuntos de MMLU para capacidades generales y MATH para razonamiento matemático) para obtener retroalimentación.
• Entrenamiento del Agente (RL):
  • Arquitectura: Un modelo Transformer ligero (2.1M parámetros) que toma la historia de la trayectoria y la retroalimentación del entorno como entrada para predecir la siguiente distribución de dominios.
  • Algoritmo: Se utiliza Conservative Q-Learning (CQL). A diferencia del ajuste fino supervisado (SFT) que solo imita las "buenas" trayectorias, CQL permite aprender tanto de las trayectorias que mejoran el rendimiento como de las que lo degradan, evitando la sobreestimación de acciones fuera de distribución.
  • Fase de Calentamiento: Se realiza un SFT inicial para reducir el espacio de búsqueda antes de aplicar CQL.
• Generalización: Una vez entrenado, el agente puede predecir la distribución de datos óptima "en vuelo" (on-the-fly) durante el pre-entrenamiento continuo del modelo objetivo real, sin necesidad de reentrenar el agente para nuevos escenarios.

3. Contribuciones Principales

1. Primera Metodología Basada en Modelos: Introducen el primer agente que aprende heurísticas generalizables para el re-pesaje de dominios en pre-entrenamiento continuo, eliminando la dependencia de reglas manuales.
2. Superioridad Empírica: El agente supera significativamente a los métodos state-of-the-art (como RegMix y DBL) en benchmarks de razonamiento matemático y generación de código, logrando un mejor equilibrio entre capacidades generales y específicas.
3. Generalización Fuerte: El agente demuestra capacidad de generalización sin reentrenamiento a través de:
   • Diferentes modelos objetivo (LLaMA-3B, Pythia, etc.).
   • Diferentes distribuciones de datos fuente.
   • Espacios de dominios de alta dimensión (52 dimensiones) y baja dimensión (2 dimensiones).
   • Nuevos campos objetivo (aplicado a generación de código tras entrenarse en matemáticas).
4. Eficiencia de Datos: Logra un rendimiento superior utilizando menos tokens del campo fuente, reduciendo el costo computacional total.

4. Resultados Experimentales

Los experimentos se centraron en el pre-entrenamiento continuo para razonamiento matemático y generación de código.

• Rendimiento en Matemáticas:
  • DataAgentRL alcanzó un promedio de 47.03% en benchmarks generales y matemáticos, superando a RegMix en un 3.02% y a DBL en un 3.53%.
  • Logró un rendimiento superior al modelo base en capacidades generales y superior al entrenamiento ingenuo (solo datos objetivo) en razonamiento matemático, mitigando efectivamente el olvido catastrófico.
  • Redujo la varianza de rendimiento en más de 200 puntos en comparación con el modelo base, indicando una estabilidad superior.
• Generalización a Código:
  • Al aplicar el agente entrenado en matemáticas directamente a la generación de código (sin reentrenar), superó a los métodos optimizados específicamente para código (RegMix y DBL) en un promedio de 1.45% - 2.78%.
• Eficiencia Computacional:
  • El agente requiere menos horas de GPU que los métodos baselines (ahorro promedio de ~106 horas en matemáticas y ~264 horas en código) porque aprende a reducir la dependencia de datos del campo fuente innecesarios.
  • El costo de entrenamiento del agente (muestreo de trayectorias) es una inversión única que se amortiza al generalizar a múltiples escenarios.

5. Significado e Impacto

• Automatización de la Curación de Datos: Este trabajo avanza hacia la automatización de la estrategia de mezcla de datos, un paso crítico para el desarrollo eficiente de LLMs especializados.
• Descubrimiento de Heurísticas: El agente descubre heurísticas que coinciden con la intuición humana (ej. aumentar datos de "Ciencia" para MMLU) pero también encuentra patrones más sutiles y complejos que los humanos no habían explotado (ej. estrategias de tres etapas: calentamiento, aumento rápido de datos objetivo y reintroducción de datos fuente).
• Escalabilidad: Al ser un método agnóstico al modelo y a los datos, ofrece una solución escalable para adaptar LLMs a nuevos dominios sin necesidad de costosos procesos de búsqueda de hiperparámetros manuales para cada nuevo caso de uso.
• Eficiencia de Recursos: Demuestra que es posible lograr un rendimiento de vanguardia con menos datos de entrenamiento, lo cual es crucial dado el costo exponencial del entrenamiento de modelos grandes.

En resumen, Data Mixing Agent representa un cambio de paradigma desde las reglas heurísticas manuales hacia un enfoque de aprendizaje automático que optimiza dinámicamente la mezcla de datos, logrando un equilibrio superior entre la especialización y la preservación de capacidades generales en el pre-entrenamiento continuo.