Resumen Técnico: Entrenamiento para el Modelo que se Devuelve

Planteamiento del Problema

Los procesos de implementación (pipelines) de los Modelos de Lenguaje (LM) modernos devuelven frecuentemente un modelo promediado —típicamente un Promedio Móvil Exponencial (EMA) de las iteraciones de entrenamiento— en lugar de la iteración final misma. Si bien el promedio de iteraciones es conocido por estabilizar la optimización y mejorar la generalización, una pregunta fundamental permanece: Dado que el resultado final será un promedio de iteraciones, ¿cómo debe modificarse el proceso de entrenamiento mismo para optimizar el rendimiento de este estimador específico?

El trabajo previo ha abordado esto implícitamente al observar que el promedio permite tasas de aprendizaje más agresivas. Sin embargo, este artículo toma un enfoque directo: busca diseñar un algoritmo de optimización que controle explícitamente la trayectoria de entrenamiento para minimizar el error del estimador de promedio de iteraciones devuelto.

Metodología

Formulación Teórica

Los autores enmarcan el problema como un problema de control óptimo dentro de un modelo cuadrático estocástico de tiempo continuo.

• Dinámica: La trayectoria de los parámetros $\theta_t$ sigue una ecuación diferencial estocástica (SDE) que representa el descenso de gradiente con ruido cerca de un mínimo cuadrático (proceso de Ornstein–Uhlenbeck), aumentada por una entrada de control $u_t$.
• Objetivo: El objetivo es encontrar una estrategia de control $u_t$ que minimice el error cuadrático medio entre el estimador promediado en el tiempo $\hat{\theta}_{EMA, T}$ devuelto y el óptimo real $\mu^*$, sujeto a una penalización cuadrática sobre la magnitud de la intervención de control (parámetro de regularización $\lambda$).
• Derivación: Utilizando la teoría de control Lineal-Cuadrático-Gaussiano (LQG) y el principio de equivalencia de certeza, los autores derivan un controlador óptimo. La solución exacta depende de la media posterior del óptimo dado la trayectoria, lo cual es intratable para LMs de gran escala.

Aproximación Práctica: PACE

Para hacer la solución implementable, los autores introducen una serie de aproximaciones heurísticas:

1. Sustitución del Estimador: Reemplazar la media posterior desconocida del óptimo con el promedio de la iteración actual ($\hat{\theta}_{EMA, t}$).
2. Aproximación Temporal: Asumir que el tiempo actual $t$ es mucho menor que el tiempo total de entrenamiento $T$ ($t \ll T$), simplificando los coeficientes dependientes del tiempo.
3. Discretización y Precondicionamiento: Discretizar el control continuo y reemplazar la matriz Hessiana desconocida con el precondicionador diagonal utilizado en AdamW (la estimación del segundo momento $v_t$).

Esto produce PACE (Pullback Averaging Control for Efficient Optimization), un envoltorio (wrapper) ligero alrededor de AdamW. La regla de actualización añade un término de "retroceso" (pullback) que mueve los pesos vivos hacia los pesos del EMA:
$$\theta_{k+1} = \theta_k - \eta g_k + C(\hat{\theta}_{EMA, k} - \theta_k)$$
donde $C$ es una matriz diagonal derivada del precondicionador de Adam, y la fuerza del retroceso se recorta para asegurar que la actualización siga siendo una combinación convexa del paso estándar y el retroceso.

Contribuciones Clave

1. Formulación de Control Óptimo: El artículo proporciona la primera derivación fundamentada de un algoritmo de optimización diseñado específicamente para minimizar el error de un estimador de promedio de iteraciones, en lugar de la iteración final.
2. Garantías Teóricas:
   • Entorno Convexo: Los autores demuestran que una versión estilizada de PACE converge al mínimo de cualquier pérdida convexa con gradientes estocásticos insesgados a la tasa estándar de $O(1/\sqrt{T})$, salvo por un factor constante dependiente de la fuerza del EMA. Esto asegura que el método nunca sea significativamente peor que el SGD estándar.
   • Entorno Cuadrático: En el entorno cuadrático utilizado para la derivación, la dinámica controlada mejora estrictamente el error cuadrático límite del estimador de promedio de iteraciones en comparación con la dinámica no controlada (SGD estándar con EMA). La mejora puede ser arbitrariamente grande en ciertas instancias cuadráticas.
3. Validación Empírica: El método se evalúa en:
   • Ajuste Fino (Fine-tuning): Modelos SmolLM2-1.7B, Qwen3-1.7B y Gemma3-1B en el conjunto de datos smol-smoltalk.
   • Preentrenamiento: GPT-2 (124M) en FineWeb con el presupuesto de tokens de Chinchilla-óptimo.
   • Líneas Base (Baselines): PACE se compara contra AdamW, AdamW con evaluación de EMA y el algoritmo Schedule-Free.

Resultados

• Rendimiento: PACE mejora estrictamente tanto a AdamW vanilla como a AdamW con evaluación de EMA en todos los modelos probados (1B–2B parámetros) y conjuntos de datos.
• Robustez: La mejora es robusta a través de una amplia gama de tasas de aprendizaje, esquemas de decaimiento (constante, coseno, Warmup-Stable-Decay) e hiperparámetros (fuerza de retroceso $c$, potencia $\kappa$ del EMA).
• Comparación con Schedule-Free: PACE es competitivo con el algoritmo Schedule-Free (Defazio et al., 2024), que también utiliza el promedio de iteraciones para eliminar la necesidad de decaimiento de la tasa de aprendizaje. Sin embargo, PACE logra esto con una regla de actualización diferente.
• Decaimiento de la Tasa de Aprendizaje: PACE funciona bien incluso con tasas de aprendizaje constantes, lo que sugiere que puede estabilizar el entrenamiento en regímenes donde el decaimiento de la tasa de aprendizaje es típicamente requerido. En comparaciones de presupuesto de tokens, la trayectoria de LR constante de PACE sigue el "suelo del río" (river floor) del paisaje de pérdida tan eficazmente como las ramas WSD, y a menudo las supera.

Significado y Limitaciones

Significancia:
El artículo demuestra que optimizar para el modelo devuelto (el promedio) en lugar de la iteración final conduce a ganancias de rendimiento tangibles. Al tratar la trayectoria de entrenamiento como un problema de control, PACE proporciona un método con base teórica para estabilizar el entrenamiento y mejorar la generalización sin requerir un ajuste complejo de esquemas. Los autores señalan que PACE ofrece una convergencia teórica mejorada en el entorno cuadrático en comparación con métodos estándar, aunque esto conlleva una garantía de convergencia teórica ligeramente más lenta en el entorno convexo general debido al factor EMA.

Limitaciones:

• Supuestos Teóricos: La derivación se basa en un modelo cuadrático simplificado con ruido aditivo. Aunque la convergencia está probada para funciones convexas generales, la incorporación rigurosa del momentum y el precondicionamiento adaptativo (central en AdamW) en el análisis teórico sigue siendo una dirección abierta.
• Escala: Los experimentos se limitaron a modelos de hasta 2B de parámetros para ajuste fino y 124M para preentrenamiento. La eficacia a escalas mayores (ej. 70B+) no ha sido verificada.
• Sobrecarga de Memoria: PACE requiere mantener una copia adicional de los pesos del modelo (el estado EMA) durante el entrenamiento, lo que introduce una sobrecarga de memoria no trivial, aunque esto se mitiga con tamaños de lote grandes en la práctica.
• Sensibilidad a la Semilla: Debido a restricciones de cómputo, las ejecuciones principales utilizaron una sola semilla, aunque los estudios de robustez con múltiples semillas confirmaron la estabilidad del orden del método.

Los autores concluyen que PACE representa un paso prometedor hacia la alineación de los algoritmos de optimización con la estrategia de despliegue real de los LMs modernos, donde el modelo promediado es el producto final.