How to Allocate Your Tokens? Scaling Laws with Training Steps and Batch Size

Resumen Técnico: Leyes de Escalamiento con Pasos de Entrenamiento y Tamaño de Lote

Planteamiento del Problema
Las leyes de escalamiento existentes para los Modelos de Lenguaje Extensos (LLMs), como la ley Chinchilla, modelan la pérdida de prueba del modelo como una función del tamaño del modelo ($N$) y el total de datos de entrenamiento ($D$). Aunque son efectivas para determinar la asignación óptima de cómputo entre $N$ y $D$, estas leyes no consideran explícitamente los hiperparámetros de entrenamiento, específicamente el tamaño de lote o batch size ($b$) y el número de pasos de entrenamiento ($K$). Los intentos previos de modelar el escalamiento de hiperparámetros a menudo fallan al integrarse con las leyes de escalamiento basadas en la pérdida o dependen del concepto de "tamaño de lote crítico" que, en algunas formulaciones teóricas, implica un tamaño de lote óptimo de uno, lo cual es una contradicción con los hallazgos empíricos donde los tamaños de lote óptimos escalan con el presupuesto de tokens. Además, ajustar las leyes de escalamiento para hiperparámetros óptimos suele requerir barridos exhaustivos de tamaños de lote para varios presupuestos de tokens, lo que incurre en costos computacionales masivos.

Metodología
Los autores proponen una ley de escalamiento de tres términos (3TL) que descompone explícitamente el presupuesto de datos de entrenamiento $D$ en el tamaño de lote en tokens ($M = b \times s$, donde $s$ es la longitud de la secuencia) y los pasos de entrenamiento ($K$), de tal manera que $D = MK$. La forma funcional propuesta para la pérdida $L$ es:

$$L(N, M, K) = E + \frac{A}{N^\alpha} + \frac{B}{M^\beta} + \frac{C}{K^\gamma}$$

donde $(E, A, B, C, \alpha, \beta, \gamma)$ son parámetros ajustados a partir de ejecuciones de entrenamiento.

Los autores comparan esta 3TL contra un enfoque de ley de dos términos (2TL), donde el tamaño del modelo $N$ es fijo y la ley se ajusta independientemente para cada $N$ utilizando la forma $L(M, K) = E + B/M^\beta + C/K^\gamma$.

Los pasos metodológicos clave incluyen:

1. Derivación del Tamaño de Lote Óptimo: Al minimizar la 3TL con respecto a $M$ sujeto a la restricción $D=MK$, los autores derivan un tamaño de lote óptimo $M^\star$ que escala como una ley de potencia de $D$: $M^\star \propto D^{\frac{\gamma}{\beta+\gamma}}$.
2. Utilización de Datos: A diferencia de los métodos anteriores que requieren identificar el único tamaño de lote óptimo para cada configuración (mediante barridos completos), la 3TL se ajusta utilizando datos de tamaños de lote subóptimos. Esto permite que la ley se ajuste de manera robusta con significativamente menos ejecuciones de entrenamiento.
3. Modelado del Régimen Subóptimo: Para abordar las restricciones de hardware donde el tamaño de lote óptimo teórico es infactible, los autores proponen un procedimiento de ajuste de dos etapas. Primero, utilizan la 3TL para informar una forma funcional simplificada $L(b) = \tilde{E} + \tilde{A}b^{-\tilde{\alpha}} + \tilde{B}b^{\tilde{\alpha}}$ (asumiendo $\gamma \approx \beta$). Segundo, ajustan una ley de potencia al intervalo de tamaños de lote $\epsilon$-subóptimos (donde $\epsilon$ representa un umbral de desperdicio de cómputo del 5%).
4. Análisis del Tamaño de Lote Crítico: Los autores analizan el comportamiento del número de pasos $K_{\bar{L}}(b)$ necesarios para alcanzar una pérdida objetivo $\bar{L}$ como una función del tamaño de lote, comparando las predicciones de la 3TL con los hallazgos empíricos sobre el tamaño de lote crítico.

Resultados Clave
Los autores evaluaron su enfoque en dos conjuntos de datos: Li (Li et al., 2025a) y OpenEuroLLM.

• Escalamiento del Tamaño de Lote Óptimo: La 3TL recupera con éxito el escalamiento del tamaño de lote óptimo $M^\star$. En el conjunto de datos Li, el escalamiento derivado $M^\star \approx 0.667 \cdot D^{0.566}$ se alinea estrechamente con la "Ley de Pasos" (Step-Law) ($M^\star \approx 0.58 \cdot D^{0.571}$) de Li et al. (2025a), la cual fue derivada de un barrido completo de tamaño de lote.
• Eficiencia de Cómputo: Un hallazgo principal es que la 3TL puede ajustarse utilizando solo el 28% de las ejecuciones de entrenamiento requeridas por los métodos tradicionales. Al enmascarar el conjunto de datos para incluir solo 2 tamaños de lote por combinación de $(N, D)$ (en lugar de un barrido completo), la 3TL aún recupera la ley de escalamiento correcta para $M^\star$. En contraste, ajustar directamente una ley de potencia a los pocos tamaños de lote "mejores" disponibles (sin modelar la superficie de la pérdida) resulta en leyes de escalamiento altamente inestables e inexactas.
• Tamaños de Lote Subóptimos: El procedimiento de ajuste de dos etapas modela con éxito el intervalo de tamaños de lote subóptimos que incurren en un máximo de 5% de desperdicio de cómputo. Los autores encuentran que este intervalo tiene un ancho de aproximadamente $2^2$ (un factor de 4) en el tamaño de lote, y su comportamiento de escalamiento es consistente a través de diferentes tamaños de modelo y conjuntos de datos.
• Tamaño de Lote Crítico: La 3TL modela correctamente el fenómeno donde el número de pasos para alcanzar una pérdida objetivo se vuelve menos sensible al aumento del tamaño de lote más allá de un punto crítico. Crucialmente, el modelo muestra que este comportamiento crítico escala con el presupuesto de tokens $D$ pero es mayormente invariante al tamaño del modelo $N$, coincidiendo con las observaciones empíricas de Zhang et al. (2025).
• Comparación con Chinchilla: Cuando la 3TL se reduce a una forma Chinchilla estándar (al sustituir el $M^\star$ óptimo de vuelta en la ecuación de la pérdida), el escalamiento de cómputo óptimo para el tamaño del modelo $N$ implícito difiere de un ajuste directo de la ley Chinchilla. La 3TL sugiere un exponente de escalamiento menor para $D$, lo que resalta que la elección de la formulación de la ley de escalamiento impacta las configuraciones óptimas derivadas.

Significancia y Reivindicaciones
El artículo afirma que la ley de tres términos ofrece un marco unificado que cierra la brecha entre las leyes de escalamiento tipo Chinchilla (tamaño del modelo vs. datos) y las leyes de escalamiento de hiperparámetros (tamaño de lote). Su significancia radica en tres áreas principales:

1. Robustez con Datos Dispersos: La capacidad de ajustar la ley utilizando tamaños de lote subóptimos reduce drásticamente el costo computacional de derivar las leyes de escalamiento, haciendo que sea factible determinar los tamaños de lote óptimos sin necesidad de barridos exhaustivos.
2. Consistencia Teórica: La formulación es estructuralmente similar a los límites de convergencia en la optimización estocástica, proporcionando un puente teórico entre el análisis de escalamiento empírico y la teoría de la optimización.
3. Aplicabilidad Práctica: A diferencia de modelos anteriores que implican un tamaño de lote óptimo de uno o fallan al describir regímenes subóptimos, la 3TL proporciona un tamaño de lote óptimo no trivial y modela explícitamente el desempeño en regímenes subóptimos, lo cual es crítico para el entrenamiento práctico bajo restricciones de hardware.

Los autores señalan limitaciones, incluyendo posibles inconsistencias en los valores de los parámetros a través de diferentes conjuntos de datos (por ejemplo, el término de pérdida irreducible $E$) y el hecho de que la ley aún requiere tasas de aprendizaje óptimas para cada configuración. También reconocen que la 3TL difiere de los ajustes directos de Chinchilla en su escalamiento de cómputo óptimo, sugiriendo que puede no ser el único instrumento para determinar tamaños de modelo óptimos, pero es superior para comprender la interacción entre pasos, tamaño de lote y datos.