Test-Time Scaling Makes Overtraining Compute-Optimal

Resumen Técnico: Test-Time Scaling Makes Overtraining Compute-Optimal

1. El Problema

Las leyes de escalado actuales para Grandes Modelos de Lenguaje (LLM) operan en dos regímenes desconectados:

• Escalado de Pre-entrenamiento (Chinchilla): Determina cómo asignar el cómputo para entrenar un modelo (tamaño del modelo $N$ vs. número de tokens de entrenamiento $D$) para minimizar la pérdida (NLL). Asume implícitamente una inferencia de un solo paso ($k=1$).
• Escalado en Tiempo de Ejecución (Test-Time Scaling): Investiga cómo mejorar el rendimiento mediante la generación de múltiples muestras (muestreo repetido, $k > 1$) durante la inferencia para aumentar la probabilidad de éxito (métrica pass@k).

La brecha crítica: No existe una ley de escalado unificada que optimice conjuntamente el entrenamiento y la inferencia. Las decisiones de pre-entrenamiento (tamaño del modelo y tokens) se toman sin considerar el presupuesto de inferencia futuro. En la práctica, los modelos modernos a menudo se "sobreentrenan" (entrenan modelos más pequeños con muchos más tokens de lo que dicta Chinchilla) para reducir el costo por consulta, pero no hay un marco teórico que justifique esto óptimamente cuando se considera el costo de generar múltiples muestras.

Pregunta de investigación: ¿Deberían cambiar las decisiones de pre-entrenamiento (tamaño $N$, tokens $D$) si conocemos el presupuesto de escalado en tiempo de ejecución (número de muestras $k$)?

2. Metodología: Leyes de Escalado Train-to-Test (T2)

Los autores proponen T2 (Train-to-Test), un marco que optimiza conjuntamente $N$, $D$ y $k$ bajo un presupuesto total de cómputo fijo que incluye tanto el entrenamiento ($C_{train} \approx 6ND$) como la inferencia ($C_{inf} \approx 2Nk$).

El problema se formula como una optimización conjunta:
$$\min_{N,D,k} L(N, D, k) \quad \text{sujeto a} \quad 6ND \leq C_{train} \text{ y } 2Nk \leq C_{inf}$$

Para resolver esto, proponen dos enfoques de modelado complementarios:

• Enfoque 1: Modelado Paramétrico de la Pérdida (NLL)

  • Extienden la ley de escalado de Chinchilla ($L(N,D)$) añadiendo un término de ley de potencia para el número de muestras $k$.
  • Asumen que la contribución de la pérdida negativa logarítmica del pass@k sigue una ley de potencia: $\hat{L}(N, D, k) = E + \frac{A}{N^\alpha} + \frac{B}{D^\beta} + \frac{G}{k^\gamma}$.
  • Esto permite modelar cómo la pérdida disminuye al aumentar las muestras, manteniendo la estructura de Chinchilla.

• Enfoque 2: Modelado Paramétrico de la Precisión (Pass@k)

  • Modelan directamente la métrica de precisión pass@k en lugar de la pérdida.
  • Utilizan una distribución Beta para modelar la probabilidad de éxito por pregunta ($p_i$), donde los parámetros de la Beta dependen de la pérdida estimada por Chinchilla.
  • Calculan el pass@k esperado integrando sobre esta distribución, lo que evita la sobreestimación que ocurriría al aplicar la fórmula de pass@k a la precisión media simple (debido a la concavidad).

Corrección de Costo de Inferencia:
Para comparar modelos equitativamente, fijan un presupuesto de inferencia por token ($C_{inf}$). Esto implica una relación inversa entre el tamaño del modelo y el número de muestras:
$$k = \frac{C_{inf}}{2N}$$
Los modelos más pequeños reciben más muestras ($k$ alto), mientras que los grandes reciben menos, manteniendo constante el costo de inferencia.

3. Contribuciones Clave

1. Optimización End-to-End: Formalizan el problema de escalado desde el entrenamiento hasta la prueba como una optimización conjunta bajo presupuestos de cómputo fijos.
2. Dos Enfoques de Modelado: Introducen formulaciones basadas en pérdida (NLL) y precisión (Pass@k) que incorporan explícitamente el costo de inferencia.
3. Validación Empírica en el Régimen de Sobreentrenamiento: Entrenan modelos desde cero en el régimen de sobreentrenamiento predicho por T2 (más allá de los checkpoints estándar de Chinchilla) para validar las predicciones.
4. Persistencia Post-Entrenamiento: Demuestran que las ventajas de los modelos sobreentrenados óptimos persisten después de técnicas de ajuste fino (Fine-Tuning y SFT).

4. Resultados Experimentales

Los autores evaluaron sus leyes en 8 tareas (4 reales: LAMBADA, ARC-Easy, SciQ, OpenBookQA; y 4 sintéticas) utilizando una cuadrícula de más de 100 checkpoints (desde 5M hasta 901M parámetros).

• RQ1: ¿Debería cambiar el pre-entrenamiento?

  • Sí. Cuando se considera el presupuesto de inferencia, las decisiones óptimas de pre-entrenamiento se desplazan radicalmente hacia el régimen de sobreentrenamiento.
  • En lugar de la regla de Chinchilla (20 tokens por parámetro), T2 recomienda modelos significativamente más pequeños entrenados con muchos más tokens (sobreentrenamiento extremo).
  • Los modelos óptimos bajo T2 son más pequeños y utilizan más muestras repetidas en inferencia, logrando un mejor rendimiento que los modelos "óptimos" de Chinchilla para el mismo presupuesto total.

• RQ2: ¿Extrapola T2 a checkpoints sobreentrenados?

  • Sí. Los modelos entrenados desde cero en la región de sobreentrenamiento predicha por T2 superaron consistentemente a los checkpoints óptimos de Chinchilla en todas las tareas.
  • El Enfoque 1 (NLL) extrapoló con un error relativo del 2.8%, mientras que el Enfoque 2 (Precisión) tuvo un error del 8.4%, ambos validando la predicción teórica.

• RQ3: ¿Sobrevive T2 al post-entrenamiento?

  • Sí. Tras aplicar Fine-Tuning (FT) y Supervised Fine-Tuning (SFT), los modelos pequeños y sobreentrenados continuaron superando a los modelos grandes óptimos de Chinchilla.
  • Aunque el sobreentrenamiento hace que el ajuste fino sea más difícil (como se ha observado en trabajos previos), la ventaja de cómputo óptimo se mantiene.

5. Significado e Impacto

• Reevaluación de la Estrategia de Entrenamiento: El artículo desafía la sabiduría convencional de Chinchilla para escenarios de despliegue modernos. Sugiere que para tareas que requieren razonamiento o alta precisión (donde se usa muestreo repetido), es más eficiente entrenar modelos pequeños con datos masivos en lugar de modelos grandes con datos moderados.
• Eficiencia de Recursos: Proporciona una "hoja de ruta" para los practicantes: si se sabe que se utilizará escalado en tiempo de ejecución (muestreo repetido), se debe entrenar un modelo más pequeño por más tiempo.
• Unificación de Regímenes: Cierra la brecha entre la teoría de escalado de entrenamiento y las estrategias de inferencia, demostrando que están fundamentalmente acopladas.
• Implicaciones para Modelos Frontera: A medida que los modelos avanzados (como los de razonamiento) requieren miles de muestras para resolver problemas complejos, las leyes T2 ofrecen un marco para diseñar arquitecturas y planes de entrenamiento que maximicen el rendimiento dentro de un presupuesto de cómputo realista.

En conclusión, el trabajo demuestra que el sobreentrenamiento no es un desperdicio, sino una estrategia óptima de cómputo cuando se integra con estrategias de inferencia escalada.