Resumen Técnico: Sobre la No Linealidad del Escalamiento de la Tasa de Aprendizaje para el Entrenamiento de LLM

Declaración del Problema

El entrenamiento eficiente de modelos de lenguaje de gran tamaño (LLM) requiere un ajuste cuidadoso de los hiperparámetros, pero realizar barridos directos de tasas de aprendizaje en escalas objetivo suele ser prohibitivamente costoso. En consecuencia, los profesionales recurren a la transferencia de hiperparámetros, prediciendo configuraciones óptimas para modelos a gran escala basándose en ejecuciones a menor escala. Un supuesto prevalente en este dominio es que la tasa de aprendizaje óptima ($\eta^*$) sigue una relación de escalamiento log-lineal (ley de potencia) con respecto a la escala de datos ($D$) y el tamaño del modelo ($N$):
$$\log \eta^*(D, N) = a \log D + b \log N + c$$
Sin embargo, estudios existentes a menudo no evalúan rigurosamente si este supuesto log-lineal se cumple globalmente o solo localmente, y frecuentemente omiten el costo computacional requerido para ajustar las leyes de escalamiento. Este artículo investiga la validez de esta ley de escalamiento y la eficiencia de diferentes estrategias de extrapolación (eje de datos vs. eje de modelo) y parametrizaciones (tasa de aprendizaje vs. tasa de aprendizaje efectiva).

Metodología

Los autores realizaron un estudio empírico sistemático utilizando arquitecturas tipo GPT-2 (que van desde 22M hasta 707M de parámetros) entrenadas con el conjunto de datos FineWeb-100B con presupuestos de tokens de 5B a 100B.

1. Diseño Experimental:

   • Los modelos fueron entrenados usando AdamW con un tamaño de lote fijo (0.52M de tokens) y un esquema de calentamiento-estacionario-decaimiento (WSD).
   • Las tasas de aprendizaje se barrieron sobre una cuadrícula logarítmica ($\log_2 \eta \in [-16, -7]$).
   • Tasa de Aprendizaje Óptima ($\eta^*$): Estimada mediante el ajuste de un polinomio cúbico a la pérdida de validación como función de $\log \eta$.
   • Tasa de Aprendizaje Efectiva ($\eta_{eff}$): Definida como el tamaño del paso en el espacio de pesos normalizado: $\eta_{eff}(t) = \|\hat{w}_{t+1} - \hat{w}_t\|_2$, donde $\hat{w} = w/\|w\|_2$. La tasa de aprendizaje efectiva óptima ($\eta^*_{eff}$) fue estimada de manera similar.

2. Métricas de Evaluación:

   • $R^2$ Fuera de Distribución ($R^2_{OOD}$): Mide la precisión predictiva de las leyes de escalamiento en escalas no vistas.
   • Relación de Computación Extra (ECR): Cuantifica el costo computacional adicional incurrido al usar una tasa de aprendizaje predicha frente a la verdadera óptima para alcanzar la misma pérdida de validación.
   • Extrapolación Restringida por Computación: Los autores compararon la extrapolación a lo largo del eje de datos (variando $D$ con $N$ fijo y grande) frente al eje del modelo (variando $N$ con $D$ fijo y grande) bajo presupuestos de computación idénticos.

3. Análisis Teórico:

   • El artículo analiza la dinámica implícita de AdamW, específicamente la relación entre la tasa de aprendizaje bruta $\eta$ y la tasa de aprendizaje efectiva $\eta_{eff}$, considerando el comportamiento transitorio de las normas de los pesos.
   • Los experimentos se repitieron usando AdamH, un optimizador que controla explícitamente $\eta_{eff}$, para aislar los efectos de la dinámica de la norma de los pesos.

Resultados Clave

1. No Linealidad de la Tasa de Aprendizaje Óptima

El estudio encuentra que la tasa de aprendizaje óptima $\eta^*$ no sigue una ley de escalamiento log-lineal global.

• Curvatura Ascendente: A medida que la escala de datos $D$ y el tamaño del modelo $N$ aumentan, $\eta^*$ exhibe una curvatura ascendente sistemática. La extrapolación log-lineal desde escalas menores consistentemente subestima la tasa de aprendizaje óptima para escalas mayores.
• Linealidad de la Tasa de Aprendizaje Efectiva: En contraste, la tasa de aprendizaje efectiva óptima ($\eta^*_{eff}$) exhibe un comportamiento de escalamiento sustancialmente más lineal a través de los mismos rangos, con coeficientes de correlación de Pearson más altos.

2. Superioridad de la Extrapolación en el Eje de Datos

Al comparar estrategias de extrapolación bajo presupuestos de computación fijos:

• Eje de Datos vs. Eje de Modelo: El ajuste de las leyes de escalamiento a lo largo del eje del tamaño de datos ($D$) produce predicciones significativamente más precisas que a lo largo del eje del tamaño del modelo ($N$).
• Eficiencia Computacional: El uso de la extrapolación en el eje de datos con la tasa de aprendizaje efectiva incurre en solo un ~2% de computación adicional en relación con el uso de la verdadera tasa de aprendizaje óptima. Esto es menos del 20% de la computación requerida para un barrido completo de hiperparámetros en la escala objetivo.

3. Explicación Mecanística

El artículo atribuye la no linealidad en $\eta^*$ a la dinámica transitoria de las normas de los pesos en AdamW:

• Dos Regímenes: La relación entre $\eta$ y $\eta_{eff}$ depende de si la norma del peso $\|w_t\|$ ha alcanzado el equilibrio.
  • Regimen de Equilibrio: $\eta_{eff} \propto \sqrt{\eta}$ (pendiente 2 en el espacio log-log).
  • Regimen Pre-equilibrio: $\eta_{eff} \propto \eta$ (pendiente 1 en el espacio log-log).
• Impacto de la Escala: A medida que la tasa de aprendizaje óptima disminuye para modelos/datos más grandes, el horizonte de entrenamiento a menudo permanece dentro del regimen de pre-equilibrio (donde las normas de los pesos aún no se han relajado). Esto desplaza la dinámica del tamaño de paso efectivo, causando la curvatura ascendente observada en $\eta^*$.
• Validación vía AdamH: Al usar AdamH para controlar explícitamente $\eta_{eff}$, la ley de escalamiento se vuelve estrictamente log-lineal, confirmando que la no linealidad en el AdamW estándar surge del mapeo implícito y variante en el tiempo entre $\eta$ y $\eta_{eff}$.

Significancia y Reivindicaciones

El artículo reivindica tres contribuciones principales:

1. Corrección Empírica: Demuestra que la ley de escalamiento log-lineal estándar para las tasas de aprendizaje es solo una propiedad local y falla en escalas mayores debido a la curvatura ascendente, lo que conduce a una extrapolación inexacta.
2. Estrategias de Transferencia Prácticas: Identifica dos elecciones específicas que mejoran sustancialmente la transferencia de hiperparámetros:
   • Parametrizar la ley de escalamiento usando la tasa de aprendizaje efectiva ($\eta_{eff}$) en lugar de la tasa de aprendizaje bruta.
   • Extrapolar a lo largo del eje de datos ($D$) en lugar del eje del tamaño del modelo ($N$).
3. Perspectiva Teórica: Proporciona una explicación mecanística para la no linealidad observada, vinculándola al comportamiento transitorio de las normas de los pesos. Los autores argumentan que $\eta_{eff}$ es la cantidad más fundamental que gobierna la dinámica de optimización en arquitecturas invariantes de escala.

Los autores señalan limitaciones, incluyendo el enfoque en arquitecturas GPT-2 y AdamW, la dependencia del ajuste cúbico para la estimación de $\eta^*$, y que el análisis está limitado a escalas de hasta 100B de tokens y 707M de parámetros. No afirman que estos hallazgos se generalicen automáticamente a todas las arquitecturas u optimizadores sin una verificación adicional.