Resumen Técnico: Seesaw: Acelerando el Entrenamiento mediante el Equilibrio entre el Programador de la Tasa de Aprendizaje y el Tamaño del Lote

Declaración del Problema

El preentrenamiento de los Grandes Modelos de Lenguaje (LLM) está cada vez más limitado por el tiempo de reloj (wall-clock time), extendiéndose a menudo durante varios meses para los modelos de vanguardia. Una estrategia principal para reducir esta duración es aumentar el tamaño del lote (batch size), lo que aprovecha el cómputo paralelo para reducir el número total de pasos de optimización secuenciales. Sin embargo, simplemente aumentar el tamaño del lote indefinidamente conduce a rendimientos decrecientes más allá de un "tamaño de lote crítico" (CBS, por sus siglas en inglés), donde la eficiencia de muestreo disminuye.

Aunque ejecuciones de entrenamiento a gran escala recientes (por ejemplo, LLaMA, Nemotron, OLMo) han empleado programas de "rampa de lote" (batch ramp)—aumentando gradualmente el tamaño del lote a lo largo del entrenamiento—, estas estrategias se ajustan actualmente de forma heurística. Existe una falta de justificación teórica sobre cómo equilibrar óptimamente el decaimiento de la tasa de aprendizaje con el aumento de la rampa del tamaño del lote, particularmente para optimizadores adaptativos como Adam. La pregunta central abordada es: ¿Cuál es el programa de tamaño de lote óptimo para minimizar el tiempo de ejecución serial sin sacrificar el rendimiento del modelo?

Metodología

Marco Teórico

Los autores desarrollan un marco basado en la regresión lineal ruidosa para establecer una equivalencia entre el decaimiento de la tasa de aprendizaje y la rampa del tamaño del lote.

1. Equivalencia de SGD: El artículo proporciona la primera prueba de muestra finita (no asintótica) demostrando que, para el Descenso de Gradiente Estocástico (SGD) en la regresión lineal ruidosa, un proceso donde el tamaño del lote se duplica en cada fase mientras la tasa de aprendizaje permanece fija es equivalente (hasta un factor constante en el exceso de riesgo) a un proceso donde el tamaño del lote es fijo y la tasa de aprendizaje se reduce a la mitad en cada fase, siempre que ambos procesos consuman el mismo número total de puntos de datos.
2. Extensión a SGD Normalizado (NSGD): Para cerrar la breancia con los optimizadores adaptativos como Adam, los autores analizan el SGD Normalizado, un sustituto analítico tratable para Adam. Introducen la Asunción 2, que postula que en el régimen dominado por la varianza observado en la práctica, las normas del cuadrado del gradiente esperado están dominadas por el ruido aditivo (escalando como $O(1/B)$). Bajo esta asunción, el NSGD se reduce efectivamente a un SGD con una tasa de aprendizaje reescalada.
3. La Condición de Equivalencia: El análisis deriva una relación específica requerida para mantener la dinámica de la pérdida al ajustar ambos parámetros. Para un factor de decaimiento de la tasa de aprendizaje $\alpha$ y un factor de aumento del tamaño del lote $\beta$, los procesos permanecen equivalentes si el producto $\alpha\sqrt{\beta}$ se mantiene constante.

El Algoritmo Seesaw

Basándose en la equivalencia teórica, los autores proponen Seesaw, un programador que equilibra dinámicamente la tasa de aprendizaje y el tamaño del lote:

• Mecanismo: Siempre que un programador estándar (por ejemplo, decaimiento de coseno) reduciría la tasa de aprendizaje por un factor $\alpha$, Seesaw en su lugar reduce la tasa de aprendizaje por $\sqrt{\alpha}$ e incrementa el tamaño del lote por $\alpha$.
• Restricción: Los autores identifican un límite teórico para prevenir la divergencia. El esquema de rampa más agresivo que mantiene la estabilidad está definido por $\alpha = \sqrt{\beta}$. Exceder esto (es decir, aumentar el tamaño del lote de forma demasiado agresiva en relación con el decaimiento de la tasa de aprendizaje) conduce a la divergencia.
• Implementación: Seesaw está diseñado como un reemplazo directo de los programadores existentes. En la práctica, aproxima el decaimiento de coseno mediante el decaimiento por pasos, activando los aumentos del tamaño del lote en los conteos de tokens específicos donde la tasa de aprendizaje de otro modo sería recortada.

Principales Contribuciones

1. Equivalencia Teórica: El artículo establece la primera equivalencia no asintótica entre el decaimiento de la tasa de aprendizaje y la rampa del tamaño del lote para SGD y extiende esto a SGD Normalizado bajo un régimen dominado por la varianza.
2. El Programador Seesaw: Introducción de un algoritmo basado en principios que reemplaza el decaimiento estándar de la tasa de aprendizaje con un programa acoplado de tasa de aprendizaje y tamaño de lote, teóricamente fundamentado para preservar la dinámica de la pérdida.
3. Análisis de Aceleración: Los autores derivan un límite superior teórico para la reducción del tiempo de ejecución serial. Al utilizar la rampa de aumento más agresiva y estable ($\alpha = \sqrt{\beta}$) frente a un control base de decaimiento de coseno, se calcula la aceleración teórica máxima como $(1 - 2/\pi) \approx 36.3\%$.

Resultados Empíricos

Los autores evaluaron Seesaw en modelos de 150M, 300M y 600M de parámetros entrenados a escala Chinchilla ($D=20N$) utilizando el código de OLMo.

• Paridad de Rendimiento: Seesaw iguala la dinámica de la pérdida de validación de los programas de decaimiento de coseno estándar en todos los tamaños de modelo cuando se entrena en el tamaño de lote crítico.
• Reducción de Tiempo de Ejecución: Seesaw logra una reducción significativa en el tiempo de ejecución serial, reduciendo el tiempo de reloj en aproximadamente un 36%, acercándose estrechamente al límite teórico derivado en el análisis.
• Robustez: El método funciona eficazmente con AdamW y diversos ajustes de decaimiento de peso (weight decay).
• Límites de la Agresividad: Los experimentos variando los parámetros $\alpha$ y $\beta$ confirman la restricción teórica. Los programas que son más agresivos que $\alpha = \sqrt{\beta}$ (por ejemplo, duplicar el tamaño del lote manteniendo la tasa de aprendizaje constante o decayéndola demasiado lentamente) resultan en divergencia o un rendimiento subóptimo.
• Límite del Tamaño de Lote Crítico: Los autores observan que la estrategia Seesaw es efectiva hasta el tamaño de lote crítico. Más allá de este punto, donde el término de ruido ya no domina la norma del gradiente, la equivalencia se rompe y Seesaw rinde peor que el decaimiento de coseno estándar.

Significancia y Reivindicaciones

El artículo afirma proporcionar una base teórica rigurosa para la práctica de la "rampa de lote", moviéndola de un ajuste heurístico a una estrategia de optimización basada en principios. Al demostrar que el decaimiento de la tasa de aprendizaje y la rampa del tamaño del lote son intercambiables bajo condiciones específicas, Seesaw ofrece un método práctico para acelerar el preentrenamiento de LLMs sin comprometer la calidad del modelo.

Los autores enfatizan que Seesaw sirve como un reemplazo directo para los programadores existentes en optimizadores adaptativos. La principal significancia radica en la capacidad de reducir el tiempo de ejecución de entrenamiento de modelos grandes en aproximadamente un tercio, abordando el cuello de botella de hardware donde los requisitos de entrenamiento están superando las mejoras de hardware. El trabajo también aclara los límites de esta estrategia, señalando que depende del régimen dominado por la varianza y puede no aplicarse una vez que los tamaños de lote exceden el umbral crítico donde los términos de gradiente promedio dominan.