Resumen Técnico: Adam-Decay con Puerta Coseno (CGAD) para DiLoCo Desacoplado

1. Declaración del Problema

En los sistemas de DiLoCo Desacoplado (DDiLoCo) asíncronos, los trabajadores realizan múltiples pasos de optimización local antes de intercambiar pseudo-gradiantes con un servidor de parámetros global. Un desafío crítico en esta arquitectura es la obsolescencia: los pseudo-gradiantes aplicados en la ronda externa $t$ pueden haberse calculado $\tau$ rondas antes debido a la latencia de la red, trabajadores rezagados o presupuestos de asincronía.

El optimizador externo estándar para DiLoCo es SGD con momento de Nesterov. Este método aplica un paso de actualización uniforme a cada pseudo-gradiante entrante sin condicionarlo a su edad ($\tau$). A medida que $\tau$ aumenta, el búfer de momento acumula una visión obsoleta del paisaje de pérdida. En consecuencia, el optimizador da pasos que sobrepasan direcciones que ya no descienden la pérdida actual, lo que lleva a inestabilidad o divergencia, particularmente a medida que aumenta la escala del modelo. Los análisis teóricos existentes de momentum-SGD asíncrono sugieren que las tasas de convergencia se degradan con un factor proporcional a $\tau_{max}^2$ (el cuadrado del retraso máximo).

2. Metodología: Adam-Decay con Puerta Coseno (CGAD)

Los autores proponen Adam-Decay con Puerta Coseno (CGAD), un reemplazo directo para el optimizador externo de Nesterov que condiciona explícitamente las actualizaciones a la edad del pseudo-gradiante.

Mecanismo Central

CGAD modifica la actualización estándar de Adam escalando cada pseudo-gradiante entrante $g_f$ para el fragmento $f$ por un factor $\sigma(\tau)$ antes de que ingrese a los búferes de primer y segundo momento de Adam:
$$\sigma(\tau) = \gamma(\tau) e^{-\alpha \tau}$$

Donde:

• Decaimiento Exponencial ($e^{-\alpha \tau}$): Modela el decaimiento de la información de los gradiantes obsoletos. A medida que $\tau$ crece, la contribución del gradiente a los búferes de momento se reduce exponencialmente.
• Puerta Coseno ($\gamma(\tau)$): Una función de corte suave definida como:
  $$\gamma(\tau) = \frac{1}{2} \left(1 + \cos\left(\frac{\pi \tau}{\tau_{cut}}\right)\right) \cdot \mathbb{I}[\tau \le \tau_{cut}]$$
  Esta puerta es igual a 1 para gradiantes frescos, decae suavemente a 0 en un corte elegido $\tau_{cut}$, y anula estrictamente las contribuciones para $\tau > \tau_{cut}$.

Decisiones Clave de Diseño

1. Puerta Pre-Búfer: El escalado $\sigma(\tau)$ se aplica al gradiente antes de que actualice los búferes de primer momento ($m$) y segundo momento ($v$). Esto evita que los gradiantes obsoletos inflen el segundo momento ($v$), lo que de otro modo atenuaría las actualizaciones frescas posteriores.
2. Variante por Fragmento (PA-CGAD): Para programadores de sincronización parcial (por ejemplo, Streaming DiLoCo), el método rastrea la edad de los fragmentos individuales ($a_f$) y puerta las actualizaciones basándose en $\max(\tau, a_f)$, permitiendo un modelo de obsolescencia unificado tanto para el retraso de red como para los retrasos de programación.
3. Hiperparámetros: El método introduce dos hiperparámetros, $\alpha$ (tasa de decaimiento) y $\tau_{cut}$ (corte), con valores predeterminados ($\alpha=0.2, \tau_{cut}=32$) que se transfieren a través de las escalas de modelos.

3. Contribuciones Teóricas

El artículo proporciona un límite de convergencia no asintótico para una actualización adaptativa con puerta idealizada en objetivos no convexos suaves.

• Independencia del Sesgo de Obsolescencia: El término de sesgo de obsolescencia del límite de convergencia depende del hiperparámetro de puerta $\alpha$ (específicamente $O(1/e^\alpha)$) en lugar del retraso máximo realizado $\tau_{max}$.
• Contraste con el Análisis Estándar: Los análisis estándar de momentum-SGD asíncrono producen un término de complejidad que escala con $\tau_{max}^2$. Al reemplazar esto con un término dependiente de $\alpha$, el mecanismo de puerta desacopla teóricamente la penalización de convergencia de la distribución de retraso del peor caso, siempre que $\tau_{cut} \ge 1/\alpha$.
• Alcance: Los autores notan que esto es una motivación para el mecanismo de puerta en lugar de un teorema de convergencia completo para Adam empírico, ya que asume pasos normalizados acotados en lugar de derivarlos de estadísticas empíricas.

4. Resultados Experimentales

Los experimentos se realizaron en el preentrenamiento de modelos de lenguaje estilo Llama en tres escalas: 25 M, 1 B y 7 B parámetros, utilizando retrasos controlados hasta $\tau=16$.

Estabilidad y Convergencia

• Escala 1 B: En $\tau=16$, la receta estándar de Nesterov divergió (pérdida $\approx 319$), mientras que CGAD logró una pérdida de 7.44, comparable a su rendimiento en $\tau=0$ (7.48). El método de Nesterov perdió la convergencia por completo en $\tau \ge 8$.
• Escala 25 M: CGAD y otras líneas base conscientes de la obsolescencia (Adam-Decay, SDM, Nesterov Retrasado) permanecieron estables, mientras que Nesterov divergió en $\tau \ge 4$. A esta escala, las diferencias entre los métodos conscientes de la obsolescencia fueron marginales.

Seguro de Escala y Robustez

El hallazgo más significativo surge en la escala 7 B bajo condiciones de memoria restringida (pesos bf16, AdamW de 8 bits):

• Varianza Semilla a Semilla: La línea base más cercana, Adam-Decay (CGAD sin el corte coseno), exhibió un aumento de 27 veces en la desviación estándar semilla a semilla ($\sigma$) a medida que el modelo escalaba de 25 M a 7 B (de 0.07 a 1.89 en $\tau=8$). En contraste, el $\sigma$ de CGAD se mantuvo plano (aprox. 0.38).
• Riesgo de Disparo Único: Usando "media + $\sigma$" como proxy para el riesgo de implementación de un solo disparo, el riesgo de Adam-Decay en 7 B / $\tau=8$ (13.19) cruzó la pérdida de nivel de azar (10.83) en 2.4 nats, implicando una alta probabilidad de fracaso en una sola ejecución de producción. El riesgo de CGAD (11.28) se mantuvo muy por debajo de este umbral.
• Mecanismo de Fallo: Los autores atribuyen el fallo de Adam-Decay en 7 B a la acumulación de ruido de redondeo al cuadrado en el búfer de segundo momento ($v$) proveniente de gradiantes obsoletos que nunca se anulan completamente. El corte coseno en CGAD previene esto anulando estrictamente las contribuciones más allá de $\tau_{cut}$.

Sensibilidad de Ajuste

• Nesterov vs. CGAD: Incluso cuando Nesterov se reajustó (eliminando el momento por completo), no pudo igualar el rendimiento de una configuración CGAD ajustada en retrasos altos. Nesterov ajustado requirió diferentes hiperparámetros para diferentes regímenes de retraso, mientras que CGAD mantuvo la robustez en todo el rango.
• Sensibilidad de Hiperparámetros: El $\alpha$ óptimo cambia ligeramente con la escala (0.2 para 25 M, 0.4 para 1 B en $\tau=8$), pero el $\alpha$ predeterminado de 0.2 permanece robusto en todo el espectro de retrasos.

5. Significado y Afirmaciones

El artículo afirma que CGAD es la opción de escalado más robusta para el entrenamiento asíncrono de DiLoCo, abordando específicamente la fragilidad del momento de Nesterov bajo grandes retrasos controlados.

• Seguro de Escala: La contribución principal es el corte coseno, que actúa como "seguro de escala". Mientras que el decaimiento exponencial por sí solo (Adam-Decay) es suficiente para modelos más pequeños o configuraciones de precisión completa, el corte duro es crítico para modelos de escala fronteriza (7 B+) que operan con estados de optimizador comprimidos (8 bits/bf16). Previene la acumulación de varianza obsoleta que conduce a una varianza catastrófica semilla a semilla y al fracaso de un solo disparo.
• Compatibilidad Directa: CGAD requiere solo una modificación de una línea en las implementaciones existentes de Adam y añade dos hiperparámetros con valores predeterminados que se transfieren a través de las escalas.
• Impacto Práctico: Para el entrenamiento asíncrono de escala fronteriza, el método ofrece la diferencia entre una sola ejecución de entrenamiento que se envía con éxito y una que falla (termina por encima de la pérdida de azar).

Los autores permanecen modestos con respecto al alcance teórico, reconociendo que el límite de convergencia se aplica a una actualización idealizada y que la superioridad empírica se demuestra específicamente dentro de la familia DiLoCo y los contextos de preentrenamiento estilo Llama. Notan que el recuento de pasos internos ($H=8$) utilizado en los experimentos es menor que las configuraciones de producción de DiLoCo ($H=100-500$), aunque los estudios de sensibilidad sugieren que la clasificación persiste en $H$ más altos.