Predicting Training Re-evaluation Curves Enables Effective Data Curriculums for LLMs

1. El Problema

El entrenamiento de Modelos de Lenguaje Grandes (LLMs) ha adoptado la práctica de utilizar "currículos de datos", donde se insertan datos de alta calidad (HQ), específicos de dominio o recientes en las etapas finales del pre-entrenamiento (fases de annealing).

• Hipótesis actual: Se asume comúnmente que colocar estos datos al final del entrenamiento, cuando la tasa de aprendizaje (Learning Rate - LR) es cercana a cero, maximiza su efectividad.
• Limitación: No existen principios claros que gobiernen la colocación óptima de los datos. Las estrategias actuales se basan en heurísticas costosas o pruebas de ablación que requieren reentrenar modelos completos. Además, no se entiende completamente cómo la posición temporal de los datos durante el entrenamiento afecta su retención por parte del modelo final.

2. Metodología y Conceptos Clave

Los autores introducen una nueva métrica diagnóstica llamada Curva de Reevaluación de Entrenamiento (TREC - Training Re-evaluation Curve) y un marco predictivo basado en la dinámica del optimizador.

A. Definición de TREC

La TREC mide qué tan bien un modelo ya entrenado (con parámetros finales $\theta_T$) se desempeña en los lotes de datos individuales ($B_t$) vistos durante el entrenamiento, en el orden en que aparecieron.

• Cálculo: Se evalúa el modelo final sobre la secuencia de entrenamiento original.
• Interpretación: Si el modelo retuviera todos los datos por igual, la TREC sería plana. En la práctica, la curva muestra variaciones: los datos vistos en ciertos momentos se retienen mejor (menor pérdida de reevaluación) que otros.
• Hallazgo clave: Bajo esquemas de decaimiento de LR (como step decay), la TREC a menudo alcanza su mínimo (mejor retención) antes del final del entrenamiento, no al final.

B. Hipótesis de Colocación Óptima

Los autores proponen que la colocación óptima de datos de alta calidad (HQ) corresponde a los mínimos de la TREC (donde el modelo final retiene mejor la información), no necesariamente al final del entrenamiento.

• Experimento: Insertaron datos de código (HQ) en diferentes segmentos del entrenamiento de modelos de 610M parámetros.
• Resultado: La colocación en el mínimo de la TREC mejoró significativamente la pérdida de validación en comparación con colocarlos al final o distribuirlos uniformemente.

C. Predicción de TREC (Sin reentrenar)

Dado que calcular la TREC real requiere entrenar el modelo primero, los autores desarrollaron un método para predecirla antes de entrenar:

1. Relación con AdamW: Demuestran que la forma de la TREC está gobernada por la escala de tiempo implícita del optimizador AdamW. Los parámetros de AdamW pueden verse como una Media Móvil Exponencial (EMA) de las actualizaciones de pesos.
2. Escala de Tiempo ($\tau$): Se define una escala de tiempo normalizada $\tau = \frac{B}{\eta \lambda D}$ (donde $B$ es el tamaño de lote, $\eta$ la LR, $\lambda$ el decaimiento de peso y $D$ el total de tokens).
3. Modelo Predictivo: La forma de la TREC se puede predecir combinando los coeficientes de la EMA de AdamW con un término de ajuste basado en la fracción de entrenamiento ($\hat{t}^m$).
   • La ecuación propuesta es: $\hat{L}_{re}(\hat{t}) = 1 - c(\hat{t})^p \cdot \hat{t}^m$.
   • El exponente $m$ captura el "desplazamiento del minimizador" (minimizer drift): cómo el óptimo local de un lote cambia a medida que avanza el entrenamiento.
   • Se encontró que $m$ sigue una ley de potencia en función de los tokens por parámetro (TPP) y la escala de tiempo $\tau$.

3. Contribuciones Principales

1. Introducción de la TREC: Una herramienta diagnóstica que revela la dinámica de retención de datos en función de la posición temporal durante el entrenamiento.
2. Estudio a Gran Escala: Análisis de 600 curvas TREC en modelos de 111M a 3.9B parámetros y diferentes configuraciones de datos (20 a 1280 TPP).
3. Modelo Predictivo Analítico: Un marco matemático que permite predecir la forma de la TREC y, por tanto, la ubicación óptima de datos HQ, utilizando solo los hiperparámetros del optimizador (LR, peso de lote, decaimiento) y el tamaño del modelo/dataset.
4. Explicación de Resultados Previos: El modelo explica por qué ciertas estrategias de currículos (como el annealing en Llama-3 405B) fallaron: colocaron datos en regiones donde la TREC estaba subiendo (mala retención) debido a una LR demasiado baja o ventanas de tiempo demasiado cortas.
5. Mejora en Entrenamiento Continuo (CPT): Aplicación exitosa del marco para mejorar el entrenamiento continuo de un modelo de 3.9B parámetros, colocando datos matemáticos en el mínimo predicho de la TREC en lugar del final.

4. Resultados Clave

• Ubicación Óptima: Colocar datos de alta calidad en el mínimo de la TREC (a menudo en la fase de decaimiento de LR, pero antes de que la LR llegue a cero) supera consistentemente a la colocación al final del entrenamiento.
• Invarianza de Escala: La forma de la TREC se alinea a través de diferentes tamaños de modelo (de 111M a 3.3B) siempre que la escala de tiempo $\tau$ y los tokens por parámetro (TPP) se mantengan constantes.
• Predicción Precisa: El modelo predictivo basado en la ley de potencia para el exponente $m$ logra una correlación de Pearson ($r_p$) superior al 98% con las TREC reales en modelos grandes, incluso cuando se entrena en modelos pequeños (111M) y se extrapola.
• Análisis de MoE: En modelos de Mezcla de Expertos (MoE), la TREC refleja una "TPP efectiva" reducida por experto, lo que explica por qué los MoE tienden a memorizar más que a razonar si no se ajustan sus escalas.
• Fallo de Heurísticas Comunes: Se demostró que colocar datos al final del entrenamiento bajo un esquema de decaimiento a cero (D2Z) es subóptimo, ya que la TREC tiende a subir (peor retención) cuando la LR es casi nula.

5. Significado e Impacto

• Diseño Proactivo de Currículos: Permite a los ingenieros diseñar currículos de datos óptimos antes de entrenar modelos costosos, eliminando la necesidad de costosas pruebas de ablación y reentrenamientos.
• Eficiencia Computacional: Al identificar la ventana exacta para insertar datos críticos (como matemáticas o código), se maximiza el rendimiento del modelo con la misma cantidad de datos, reduciendo el desperdicio de cómputo.
• Comprensión Teórica: Proporciona una conexión teórica sólida entre la dinámica de optimización (AdamW, EMA, decaimiento de LR) y la retención de datos, desafiando la intuición de que "más tarde es mejor".
• Aplicabilidad General: El marco es aplicable a pre-entrenamiento estándar y a fases de entrenamiento continuo (CPT), ofreciendo una guía principada para la actualización de modelos existentes.

En resumen, el paper demuestra que la retención de datos en LLMs no es uniforme a lo largo del tiempo, sino que sigue una curva predecible gobernada por la dinámica del optimizador. Utilizar esta curva (TREC) para guiar la inserción de datos de alta calidad es una estrategia superior a las heurísticas actuales.