Resumen Técnico: M+Adam – Entrenamiento de Baja Precisión mediante Optimización Aditiva-Multiplicativa

Planteamiento del Problema

El entrenamiento de grandes modelos de lenguaje (LLMs) en precisión completa (FP32) incurre en costos significativos de memoria, cómputo y energía. Aunque los aceleradores modernos admiten formatos de precisión reducida como BF16, FP8 y FP4, los pipelines de entrenamiento estándar suelen retener "pesos maestros" de alta precisión (FP32) para garantizar la estabilidad. Entrenar directamente con pesos maestros de baja precisión suele derivar en una degradación de la precisión e inestabilidad.

Los autores identifican un modo de fallo específico en la optimización de baja precisión: la resolución de cuantización gruesa en magnitudes de peso grandes. En los sistemas de punto flotante, el espaciamiento entre los valores representables (bins de cuantización) crece con la magnitud del número. En consecuencia, los optimizadores aditivos estándar (como Adam) a menudo no logran progresar cuando las magnitudes de los pesos son grandes, ya que los pasos de actualización calculados caen por debajo del paso de cuantización local y se redondean a cero. Por el contrario, los optimizadores puramente multiplicativos (como Madam), que escalan los pesos en relación con su magnitud actual, fallan cerca de cero o cuando se requieren cambios de signo, ya que no pueden escapar de la inicialización en cero exacto ni cambiar los signos eficazmente.

Metodología: M+Adam

El artículo propone M+Adam, un nuevo optimizador que combina las fortalezas complementarias de los mecanismos de actualización aditiva y multiplicativa para superar las limitaciones de la aritmética de baja precisión.

Mecanismo Central

M+Adam mantiene dos ramas de actualización paralelas que se combinan en cada paso:

1. Rama Aditiva (estilo Adam): Opera directamente en el espacio de los pesos. Calcula los momentos de primer y segundo orden del gradiente ($g$) para producir una actualización aditiva $u^{(a)}$. Esta rama es efectiva para manejar magnitudes pequeñas, escapar de la inicialización en cero y facilitar cambios de signo.
2. Rama Multiplicativa (estilo Madam): Opera en la escala logarítmica (espacio del exponente). Calcula el gradiente con respecto al exponente, $g^{(e)} = (\ln 2) w_t g$, y utiliza momentos adaptativos para producir una actualización multiplicativa relativa $u^{(m)}$. Esta rama asegura el progreso en magnitudes grandes donde las actualizaciones aditivas serían redondeadas a cero.

La regla de actualización de los pesos se define como:
$$w_{t+1} = w_t + w_t u^{(m)}_t + u^{(a)}_t$$
Esta estructura permite que el optimizador realice actualizaciones relativas conscientes de la escala cuando las magnitudes son grandes (vía el término multiplicativo), mientras retiene la flexibilidad para cambiar signos y actualizar valores cercanos a cero (vía el término aditivo).

Análisis Teórico

Los autores proporcionan una garantía teórica para M+Adam bajo supuestos estándar de suavidad (smoothness). Modelan la actualización como una operación de "escala-luego-suma" en coordenadas locales. El Teorema 4.1 establece que si las tasas de aprendizaje y las constantes de suavidad satisfacen condiciones específicas, el método asegura un descenso monótono ($L(w_{t+1}) < L(w_t)$) a menos que el gradiente sea cero. Esto demuestra que la geometría combinada es una estrategia de optimización válida.

Detalles de Implementación

• Soporte de Precisión: El método está diseñado para el almacenamiento de pesos maestros en BF16, FP8 y NVFP4.
• Gestión de Estado: M+Adam introduce un estado adicional del optimizador (el acumulador del segundo momento para la rama multiplicativa). Los autores demuestran que este estado puede comprimirse mediante aproximaciones de bajo rango (específicamente Apollo-E, una aproximación de rango 4) para minimizar la sobrecarga de memoria, manteniendo la huella cercana a la de un AdamW estándar con estado BF16.
• Sin Redondeo Estocástico: A diferencia de otros enfoques previos, M+Adam logra estabilidad sin requerir redondeo estocástico (SR) o acumulación compensada tipo Kahan, aunque es compatible con ellos.

Contribuciones Clave

1. Identificación de Modos de Fallo Complementarios: El artículo caracteriza formalmente los modos de fallo de los optimizadores aditivos (estancamiento en magnitudes grandes debido al redondeo) y los optimizadores multiplicativos (estancamiento en cero e incapacidad de cambiar signos), mostrando que son complementarios.
2. Algoritmo M+Adam: Un nuevo optimizador que unifica estas dos geometrías, proporcionando un enfoque basado en principios para el entrenamiento de baja precisión sin pesos maestros en FP32.
3. Garantía de Descenso Teórico: Una prueba de descenso monótono para la regla de actualización combinada aditiva-multiplicativa bajo supuestos estándar de suavidad.
4. Validación Empírica a Escala: Experimentos extensos en modelos de estilo LLaMA que van desde 60M hasta 1B de parámetros a través de presupuestos de entrenamiento de 1 a 8× Chinchilla.

Resultados Experimentales

Los autores evaluaron M+Adam en varios regímenes de precisión: BF16/BF16, BF16/FP8, FP8/FP8 y NVFP4/FP8.

• Mejora Consistente: M+Adam superó consistentemente a las líneas base de AdamW ajustadas en todos los tamaños de modelo y regímenes de precisión.
• Reducciones de Perplejidad: En el régimen de baja precisión más agresivo (NVFP4/FP8), M+Adam redujo la perplejidad por márgenes significativos comparado con AdamW (por ejemplo, una mejora del 16.6% para el modelo de 350M). Incluso en el entrenamiento BF16/BF16, M+Adam redujo la perplejidad entre un 2.8% y un 13.3% dependiendo del tamaño del modelo.
• Robustez: El método se mantuvo estable sin redondeo estocástico. También superó a AdamW combinado con redondeo estocástico (AdamW+SR) y a la compensación de Kahan.
• Escalabilidad: Las ganancias de rendimiento persistieron a medida que aumentaba el presupuesto de entrenamiento (hasta 8× Chinchilla) y a medida que el tamaño del modelo escalaba a 1B de parámetros.
• Eficiencia de Memoria: Utilizando la técnica de compresión Apollo-E, la huella de memoria adicional requerida para la rama multiplicativa se redujo, resultando en un estado de optimizador solo ligeramente mayor que el de un AdamW estándar en BF16, mientras se retienen la mayoría de las ganancias en perplejidad.

Significancia y Reivindicaciones

El artículo afirma que las actualizaciones combinadas aditivas-multiplicativas son especialmente útiles cuando los pesos maestros de baja precisión se convierten en una fuente dominante de error de optimización.

La importancia de este trabajo radica en proporcionar una solución fundamentada a la inestabilidad numérica del entrenamiento de baja precisión. Al alinear la geometría de actualización del optimizador con la estructura no uniforme de los números de punto flotante, M+Adam mitiga los errores de redondeo que afectan a los optimizadores aditivos estándar en magnitudes grandes y los problemas de estancamiento de los optimizadores multiplicativos cerca de cero.

Los autores concluyen que M+Adam representa un paso hacia el entrenamiento de baja precisión de extremo a extremo que no depende de pesos maestros en FP32. Esta capacidad es crucial para aprovechar el hardware de IA moderno (tensor cores) optimizado para FP8 y FP4, lo que potencialmente reduce la sobrecarga de memoria y los costos de energía para el entrenamiento de modelos a gran escala. El trabajo sugiere que los futuros optimizadores deben diseñarse teniendo en cuenta las restricciones específicas de la representación de punto flotante, en lugar de tratar la baja precisión simplemente como un problema de cuantización que deba resolverse mediante técnicas de redondeo.