Resumen Técnico: GEFEN: OPTIMIZADOR ESTOCÁSTICO OPTIMIZADO

Planteamiento del Problema

El aprendizaje profundo moderno depende en gran medida del optimizador AdamW debido a su rendimiento robusto en diversas arquitecturas. Sin embargo, AdamW mantiene medias móviles de los primeros y segundos momentos para cada parámetro, típicamente en FP32. Esto resulta en una huella de memoria del estado del optimizador de aproximadamente el doble del tamaño de los parámetros del modelo (aproximadamente 8 bytes por parámetro). Para modelos de gran escala (por ejemplo, 13B de parámetros), este estado por sí solo puede consumir ~97 GiB de memoria, lo que restringe severamente el tamaño de los modelos entrenables, los tamaños máximos de lote (batch sizes) y el rendimiento de entrenamiento en hardware con restricciones de memoria.

Las alternativas existentes de eficiencia de memoria enfrentan compromisos significativos:

• Optimizadores Cuantizados (ej. Adam8bit, Adam4bit): Reducen la memoria pero suelen depender de codebooks diseñados manualmente que son subóptimos, tamaños de bloque fijos (lo que introduce hiperparámetros) y pueden no ser universalmente aplicables (por ejemplo, evitando la cuantización de embeddings).
• Optimizadores Estructurados (ej. Adam-mini): Comparten estadísticas del segundo momento basándose en reglas manuales que imitan estructuras de bloque-diagonal de la Hessiana. Sin embargo, estos requieren metadatos específicos de la arquitectura (por ejemplo, el número de cabezales de atención) y carecen de una fuerte justificación teórica de por qué el agrupamiento alineado con la Hessiana es beneficioso.

El objetivo es diseñar un optimizador que reduzca significativamente la huella de memoria (más allá de Adam-mini) manteniendo un rendimiento de nivel AdamW, sin requerir configuración manual, hiperparámetros adicionales o reajuste de la tasa de aprendizaje.

Metodología

El método propuesto, Gefen, logra la reducción de memoria a través de dos mecanismos principales: el intercambio automático del segundo momento y la cuantización aprendida del primer momento.

1. Motivación Teórica

El enfoque se fundamenta en un resultado teórico (Teorema 4.1) que demuestra que, para parámetros con entradas de la Hessiana mixtas grandes, la relación de sus cuadrados de gradientes está restringida hacia uno. Esto implica que los parámetros con alta afinidad de la Hessiana poseen naturalmente magnitudes de gradiente al cuadrado similares, lo que los convierte en candidatos ideales para compartir estimaciones del segundo momento. Aunque computar la Hessiana directamente es intratable a escala, los autores observan que la estructura de los cuadrados de los gradientes en el primer paso de entrenamiento se correlaciona con esta estructura de la Hessiana.

2. Partición de Bloques Automática

En lugar de depender de reglas manuales o del cálculo de la Hessiana, Gefen infiere la estructura de bloque óptima a partir de los cuadrados de los gradientes del primer paso de optimización ($g_0$).

• Algoritmo: Para un tensor de parámetros aplanado, el algoritmo prueba todos los divisores propios $p$ del tamaño del tensor $n$. Reestructura los cuadrados de los gradientes en bloques de tamaño $p$ y calcula la heterogeneidad dentro del bloque $E(p)$ (la varianza de los cuadrados de los gradientes dentro de los bloques).
• Estrategia de Selección: El algoritmo selecciona el periodo $p$ que produce la primera caída pronunciada en $E(p)$ a medida que el tamaño del bloque aumenta, favoreciendo bloques más grandes para la eficiencia de memoria mientras asegura la similitud de los gradientes. Si no se encuentra una caída significativa o si el $p$ óptimo es $< 8$, por defecto establece $p=1$ (sin intercambio).
• Uso: Este esquema de partición único se utiliza tanto para el intercambio del segundo momento como para la cuantización del primer momento.

3. Cuantización Exacta mediante Programación Dinámica

Para reducir la memoria del primer momento (momentum), Gefen emplea un codebook de cuantización aprendible y exacto en lugar de esquemas fijos o heurísticos.

• Aprendizaje del Codebook: Utilizando los gradientes iniciales, el método construye un histograma de los valores de gradiente normalizados. Luego resuelve el problema de cuantización escalar óptima utilizando Programación Dinámica (DP) (inspirado en Wang & Song, 2011) para encontrar el codebook exacto para la distribución y el ancho de bits dados.
• Estabilidad: A diferencia de Lloyd-Max, que puede sufrir de "contracción" (reducción del rango de cuantización a lo largo del tiempo) y sensibilidad a la inicialización, Gefen fuerza los extremos del codebook a fijarse en $-1$y$1$. Esto asegura que el rango de cuantización se preserve durante todo el entrenamiento.
• Eficiencia: El codebook se aprende una sola vez al principio del entrenamiento (usando CUDA para el histograma y CPU/Numba para la DP) y se reutiliza durante toda la ejecución del entrenamiento.

4. Flujo del Algoritmo

Gefen opera como un reemplazo directo (drop-in replacement) para AdamW:

1. Inicialización: Particiona los parámetros basándose en $g_0$ y aprende el codebook de cuantización.
2. Paso de Actualización:
   • Computar gradientes.
   • Actualizar el primer momento (momentum) descuantizando los valores del codebook almacenado, aplicando el momentum, recuantizando y actualizando el máximo absoluto por bloque.
   • Actualizar el segundo momento computando la media de los cuadrados de los gradientes dentro de cada bloque y aplicando la media móvil exponencial.
   • Aplicar corrección de sesgo y actualizar los parámetros utilizando el segundo momento compartido y el primer momento cuantizado.

Principales Contribuciones

1. Eficiencia de Memoria: Gefen reduce la huella de memoria del optimizador aproximadamente 8× en comparación con AdamW (reduciendo de ~2.0× el tamaño de los parámetros a ~0.25×), lo que corresponde a una reducción de ~6.5 GiB por cada mil millones de parámetros.
2. Adaptación Automática: El método no requiere metadatos específicos de la arquitectura, reglas de agrupamiento manuales ni hiperparámetros adicionales más allá de los valores predeterminados de AdamW. Se adapta automáticamente a los datos infiriendo las estructuras de bloque.
3. Fundamentación Teórica: Proporciona un vínculo teórico entre las entradas de la Hessiana y las relaciones de los cuadrados de los gradientes, justificando el intercambio de estadísticas del segundo momento para parámetros alineados con la Hessiana.
4. Cuantización Óptima: Introduce un enfoque de programación dinámica basado en histogramas exactos para cuantizar el momentum, evitando la suboptimidad e inestabilidad de los métodos iterativos como Lloyd-Max.
5. Paridad de Rendimiento: Demuestra que estas reducciones no conllevan una pérdida en la calidad del modelo, igualando el rendimiento de AdamW en diversas tareas.

Resultados Experimentales

Los autores evaluaron Gefen frente a AdamW, Adam-mini, Adam8bit, Adam4bit, Adafactor y SM3 a través de modelos CNN, ResNet, GPT-2 y Llama 3.

• Huella de Memoria Pico: Gefen alcanzó consistentemente el uso de memoria pico más bajo entre los métodos comparados. Para un modelo de 1.5B de parámetros, Gefen redujo el estado del optimizador de 11.2 GiB (AdamW) a 1.5 GiB.
• Rendimiento (Calidad): En todas las configuraciones (OpenWebText, C4, ImageNet, MNIST), Gefen logró pérdidas de validación y precisiones comparables a AdamW. Notablemente, superó a Adam4bit y Adam-mini en varias métricas manteniendo el mismo esquema de tasa de aprendizaje.
• Rendimiento (Throughput): La reducción de la huella de memoria permitió microbatches más grandes en el entrenamiento distribuido:
  • FSDP (Llama 3-1.5B): Gefen permitió un tamaño de microbatch de 2 (vs. 1 en otros), lo que resultó en un incremento del 56% en el rendimiento (tokens/seg).
  • DDP (GPT-2-1B): Gefen permitió el entrenamiento con un tamaño de microbatch de 2, mientras que AdamW no podía siquiera ajustar un microbatch de 1. Esto resultó en una mejora del 21% en el rendimiento sobre Adam-mini.
• Estabilidad: Los estudios de ablación confirmaron que el codebook aprendido permanece estable durante todo el entrenamiento, validando la decisión de aprenderlo solo una vez en la inicialización.

Significado y Reivindicaciones

El artículo posiciona a Gefen como un reemplazo práctico y directo para AdamW que aborda el cuello de botella de la memoria en el entrenamiento a gran escala sin sacrificar el rendimiento ni requerir ajustes manuales.

• Impacto Práctico: Al reducir el estado del optimizador, Gefen permite entrenar modelos más grandes o utilizar tamaños de lote más grandes en hardware fijo, mejorando directamente el rendimiento del entrenamiento y la eficiencia de muestreo.
• Generalidad: A diferencia de otros métodos que dependen de heurísticas manuales o nombres de tensores específicos, Gefen es agnóstico a la arquitectura y se adapta automáticamente a los datos.
• Complementariedad: Los autores señalan que la compresión del estado del optimizador de Gefen es ortogonal a la compresión de comunicación de gradientes (ej. 1-bit Adam) y a las estrategias de entrenamiento distribuido (FSDP/DDP), lo que significa que puede combinarse con estas técnicas para obtener más ganancias.

Los autores proporcionan una implementación completa en Python con kernels de CUDA fusionados, enfatizando la reproducibilidad y la facilidad de integración en los flujos de trabajo de entrenamiento existentes.