Resumen Técnico: LoRDO – Optimización de Bajo Rango Distribuida con Comunicación Infrecuente

1. Planteamiento del Problema

El entrenamiento distribuido de modelos fundacionales mediante Paralelismo de Datos Distribuido (DDP) está cada vez más limitado por el ancho de banda de interconexión y la sobrecarga de memoria requerida para almacenar los estados del optimizador (por ejemplo, los momentos de Adam). Si bien las estrategias que emplean comunicación infrecuente (actualizaciones locales) reducen la frecuencia de sincronización, siguen estando limitadas por los costes de comunicación y memoria de los estados del optimizador de rango completo.

Los optimizadores adaptativos de bajo rango (por ejemplo, GaLore, LDAdam) ofrecen una vía para aliviar estas restricciones proyectando los gradientes en un subespacio de bajo rango. Sin embargo, generalizar estos métodos al régimen de sincronización infrecuente presenta un desafío de optimización crítico:

1. Proyecciones Locales: Calcular proyecciones de bajo rango localmente en cada trabajador utilizando gradientes estocásticos introduce un ruido significativo debido a los tamaños de lote efectivos pequeños por trabajador. Esto conduce a subespacios de proyección subóptimos y degradación del rendimiento.
2. Proyecciones Globales: El uso de una matriz de proyección global derivada de pseudo-gradientes agregados (como en DDP) proporciona una base más estable. Sin embargo, en el régimen de actualización local, este enfoque restringe permanentemente la trayectoria de optimización a un subespacio fijo de bajo rango. Sin señales de rango completo, la optimización se estanca porque la base de proyección no logra "refrescarse" o explorar nuevas direcciones fuera del subespacio inicial.

La pregunta central abordada es: ¿Cómo podemos diseñar optimizadores de bajo rango de alto rendimiento para el entrenamiento con eficiencia de comunicación que eviten el estancamiento del subespacio manteniendo al mismo tiempo una baja sobrecarga de comunicación y memoria?

2. Metodología: LoRDO

Los autores proponen LoRDO (Optimización Distribuida de Bajo Rango), un marco que unifica la optimización de bajo rango con la sincronización infrecuente. La metodología aborda la tensión entre la estabilidad de la proyección y la exploración del subespacio mediante tres mecanismos clave:

A. Matrices de Proyección Global

En lugar de que cada trabajador determine su propia matriz de proyección $Q_m$ basada en gradientes estocásticos locales (lo que sufre de alta varianza), LoRDO emplea una matriz de proyección global $Q_t$. Esta matriz se calcula en el límite de sincronización utilizando el pseudo-gradiente agregado ($\Delta_t$), que representa el cambio total en los parámetros del modelo durante $K$ pasos locales.

• Beneficio: El pseudo-gradiente tiene un tamaño de lote efectivo de $M \times B$ (donde $M$ es el número de trabajadores), lo que genera una señal más estable con menor varianza en comparación con los gradientes locales. También incorpora información de curvatura integrada en el pseudo-gradiente, lo que conduce a una estimación de la base de mayor calidad.

B. Momento Cuasi-Hiperbólico de Rango Completo (QHM)

Para resolver el problema del "subespacio estancado" inherente a las proyecciones globales, LoRDO inyecta un término de momento cuasi-hiperbólico de rango completo en la actualización local.

• Mecanismo: La regla de actualización combina el momento de bajo rango con una señal de gradiente de rango completo escalada por el segundo momento.
  $$\text{Actualización} \propto (1-\omega)\frac{\hat{G}_t}{\mu(\sqrt{\hat{v}_t} + \epsilon)} + \omega Q_t \left[ \frac{u_t}{\sqrt{v_t} + \epsilon} \right]$$
• Efecto: Esta inyección evita que el pseudo-gradiente agregado permanezca atrapado dentro de un subespacio fijo de rango-$r$. Permite que el procedimiento de optimización agregue temporalmente subespacios de rango-$r$ de diferentes trabajadores, permitiendo eventualmente una exploración completa del subespacio. Esto restaura la capacidad de explorar el espacio de soluciones completo sin aumentar la sobrecarga de comunicación.

C. Retroalimentación de Error y Rotación de Momento

LoRDO incorpora retroalimentación de error localmente para dar cuenta de la información perdida durante la compresión de bajo rango. Además, tras el cálculo de una nueva matriz de proyección, los estados del optimizador (momentos) son rotados para asegurar que permanezcan acumulados en el mismo subespacio, una técnica adaptada de LDAdam.

3. Contribuciones Clave

1. Análisis de Modos de Fallo de Proyección: Los autores demuestran que las proyecciones locales perjudican el rendimiento debido a la alta varianza de los pequeños lotes de los trabajadores, mientras que las proyecciones globales sin señales de rango completo inducen el estancamiento del subespacio.
2. Restauración de la Exploración de Subespacio Completo: LoRDO introduce un optimizador de bajo rango que inyecta una señal de gradiente de rango completo (vía QHM) en las actualizaciones locales mientras mantiene una proyección global derivada de pseudo-gradientes agregados. Esto evita el estancamiento sin aumentar la sobrecarga de comunicación o memoria.
3. Paridad y Eficiencia con DDP: El marco logra una paridad casi total con el DDP de bajo rango sincrónico (brecha de perplexidad < 1%) mientras reduce la memoria del optimizador y la comunicación hasta en $8\times$–$12\times$. Bajo restricciones severas de memoria que requieren rangos bajos, LoRDO incluso supera a DDP en perplexidad.

4. Resultados Experimentales

Los experimentos se realizaron en transformadores de solo decodificador escalados a 16M, 125M y 720M de parámetros utilizando la mezcla de datos de SmolLM2.

• Estancamiento vs. Exploración: En modelos de 16M, las proyecciones globales sin QHM no lograron aprender (perplejidad 46.8) en comparación con DDP (36.6), confirmando el estancamiento del subespacio. La adición de QHM de rango completo restauró el rendimiento, permitiendo que LoRDO-Global superara a las variantes locales.
• Calidad de la Proyección: LoRDO-Global demostró una estabilidad y rendimiento superiores, particularmente en rangos bajos ($r=8, 16$), debido a la base de proyección de mayor calidad derivada del pseudo-gradiente agregado. Mostró menor sensibilidad a los cambios en el tamaño del lote local en comparación con los métodos de proyección local.
• Escalado a 720M: A la escala de 720M, LoRDO-Global ($r=256$) logró una paridad casi total con el DDP de bajo rango ($r=256$) con una brecha de perplexidad de menos del 1% (10.41 vs. 10.34). También igualó la precisión en tareas de downstream (ARC, HellaSwag, MMLU, PIQA) con DDP.
• Ganancias de Eficiencia:
  • Comunicación: LoRDO reduce la sobrecarga de comunicación aproximadamente un 10$\times$ en comparación con el DDP de bajo rango y un 25$\times$ en comparación con el DDP de rango completo.
  • Memoria: Reduce la sobrecarga de memoria del estado del optimizador en 8$\times$ en comparación con sus contrapartes de rango completo.
  • Tiempo de Ejecución (Wall-Clock Time): En entornos de bajo ancho de banda (modelado a 10 Gbps), LoRDO supera significativamente a DDP en tiempo de ejecución debido a la reducción en la frecuencia de sincronización.

5. Significado y Reivindicaciones

El artículo afirma que LoRDO resuelve la tensión fundamental entre la estabilidad de la proyección y la exploración del subespacio en la optimización eficiente de comunicación de bajo rango. Al utilizar pseudo-gradientes agregados para la estabilidad de la proyección y un momento cuasi-hiperbólico de rango completo para la exploración, LoRDO permite:

• Entrenamiento Escalable: Rendimiento competitivo a escala (hasta 720M de parámetros) con costes de comunicación y memoria significativamente reducidos.
• Eficiencia de Recursos: Viabilidad del pre-entrenamiento en hardware con recursos limitados (baja memoria/ancho de banda), donde los optimizadores de rango completo son intratables.
• Aprendizaje Descentralizado: Preparar el camino para un aprendizaje eficiente en entornos descentralizados al extender el entrenamiento de modelos más allá de un único centro de datos.

Los autores señalan limitaciones, incluyendo la falta de una prueba de convergencia formal y el enfoque actual en optimizadores adaptativos (Adam) en lugar de optimizadores basados en matrices (por ejemplo, Muon), dejando esto para trabajos futuros. Sin embargo, postulan que las tendencias observadas en modelos más pequeños se mantienen en escalas mayores.