🤖 aprendizaje automático

MKOR: Momentum-Enabled Kronecker-Factor-Based Optimizer Using Rank-1 Updates

Este artículo presenta MKOR, un optimizador de segundo orden habilitado para el impulso que utiliza actualizaciones de rango 1 e inversión basada en Sherman-Morrison para reducir la complejidad computacional y de comunicación de cúbica a cuadrática y lineal respectivamente, acelerando así significativamente el entrenamiento y la convergencia de los grandes modelos de lenguaje en comparación con los métodos de primer y segundo orden de vanguardia.

Autores originales: Mohammad Mozaffari, Sikan Li, Zhao Zhang, Maryam Mehri Dehnavi

Publicado 2024-01-30✓ Author reviewed ⓘ

📖 7 min de lectura🧠 Análisis profundo

Autores originales: Mohammad Mozaffari, Sikan Li, Zhao Zhang, Maryam Mehri Dehnavi

✨ Esta es una explicación generada por IA del artículo a continuación. No ha sido escrita por los autores. Para mayor precisión técnica, consulte el artículo original. Leer descargo de responsabilidad completo

En el vasto panorama de la inteligencia artificial moderna, los sistemas más potentes se construyen sobre redes neuronales profundas, estructuras matemáticas intrincadas que aprenden a reconocer patrones, traducir idiomas y generar texto. Para enseñar a estos sistemas, los investigadores deben resolver un problema de optimización masivo: encontrar el conjunto perfecto de números internos, o pesos, que permitan a la red realizar predicciones precisas. Durante décadas, el enfoque estándar ha sido dar pequeños empujones a estos pesos en la dirección que reduce más rápidamente el error, un método similar al de un excursionista que se orienta bajando una montaña con niebla comprobando la pendiente bajo sus pies. Este método de primer orden es fiable y ampliamente utilizado, pero puede ser lento, requiriendo a menudo miles de pasos para alcanzar el fondo del valle.

Una estrategia más ambiciosa, conocida como optimización de segundo orden, intenta darle al excursionista un mapa de todo el terreno, incluyendo la curvatura del suelo, lo que permite descensos mucho más rápidos. Al comprender cómo cambia la pendiente, estos métodos pueden, teóricamente, alcanzar la solución óptima en muchísimos menos pasos. Sin embargo, este enfoque ha sido históricamente demasiado costoso computacionalmente para ser práctico en los modelos masivos que se utilizan hoy en día. Calcular la información de curvatura necesaria para una red grande requiere una potencia de procesamiento y memoria inmensas, lo que a menudo hace que el método sea más lento que el enfoque más simple que pretende reemplazar. El desafío ha sido capturar la velocidad del excursionista que lee el mapa sin la pesada carga de cargar con él.

Un equipo de investigadores de la Universidad de Toronto, la Universidad de Rutgers y el Centro de Computación Avanzada de Texas ha desarrollado una nueva herramienta llamada MKOR que resuelve este dilema. Crearon un optimizador con capacidad de impulso (momentum) que utiliza un atajo matemático específico para aproximar la forma del terreno sin calcular cada detalle. En lugar de intentar calcular la curvatura completa y compleja del paisaje de pérdida, MKOR utiliza una técnica llamada actualizaciones de rango 1 (rank-1 updates). Este enfoque simplifica el cálculo al centrarse en la dirección de cambio más significativa, reduciendo eficazmente el costo computacional de una relación cúbica a una cuadrática. En términos prácticos, esto significa que el método escala mucho mejor a medida que los modelos crecen, evitando los cuellos de botella que anteriormente impedían el uso de métodos de segundo orden en los grandes modelos de lenguaje.

Los investigadores probaron este nuevo optimizador en un gran modelo de lenguaje conocido como BERT-Large-Uncased, ejecutando experimentos en 64 procesadores gráficos de alto rendimiento. Compararon MKOR contra los mejores métodos de primer orden existentes y las implementaciones de segundo orden más avanzadas disponibles. Los resultados fueron sorprendentes. En una tarea de respuesta a preguntas, el método estándar de primer orden requirió 1,536 pasos para alcanzar un nivel específico de precisión. El nuevo método MKOR alcanzó la misma o mejor precisión en solo 600 pasos. Al medir el tiempo total requerido para completar el entrenamiento, MKOR terminó la tarea en poco más de tres horas, mientras que el método estándar tomó casi ocho horas. Esto representa una aceleración de 2.57 veces, reduciendo efectivamente el tiempo de entrenamiento a menos de la mitad.

Más allá de la velocidad pura, el estudio abordó un problema crítico de estabilidad que suele afectar a los métodos de segundo orden. Debido a que estos métodos ajustan el proceso de aprendizaje basándose en factores matemáticos complejos, son propensos a los "gradientes explosivos", un fenómeno donde los ajustes se vuelven tan grandes que el proceso de entrenamiento colapsa o diverge. Las soluciones anteriores a menudo requerían ralentizar significativamente la tasa de aprendizaje para mantener la seguridad, lo que anulaba los beneficios de velocidad. MKOR incluye un estabilizador integrado que monitorea el tamaño de estos ajustes. Si los valores crecen demasiado, el sistema los reescala automáticamente, evitando que el entrenamiento estalle mientras mantiene el ritmo agresivo del enfoque de segundo orden. Esto permite al optimizador moverse rápidamente sin el riesgo de inestabilidad.

Los investigadores también observaron que la ventaja de los métodos de segundo orden es más pronunciada en las etapas iniciales del entrenamiento, cuando el modelo está aprendiendo los contornos generales de los datos. A medida que el modelo se acerca a la solución óptima, los beneficios de los cálculos de curvatura complejos disminuyen. Para abordar esto, introdujeron una versión híbrida del optimizador, llamada MKOR-H. Esta versión comienza con el potente método de segundo orden para acelerar la fase inicial de aprendizaje, pero cambia automáticamente a un método de primer orden más simple una vez que la pérdida de entrenamiento deja de disminuir rápidamente. Este enfoque híbrido permitió al sistema converger en incluso menos pasos, logrando la misma precisión final que el método estándar en menos de la mitad del tiempo.

El estudio también examinó por qué este atajo matemático específico funciona tan bien. Los investigadores analizaron las estructuras de datos dentro de la red neuronal y descubrieron que las matrices de covarianza, que describen cómo interactúan las diferentes partes de la red, suelen ser de bajo rango. Esto significa que, a pesar de la alta dimensionalidad de los datos, la información está concentrada en unas pocas direcciones clave. Al aproximar estas matrices mediante actualizaciones de rango 1, MKOR captura la información esencial necesaria para la optimización sin la sobrecarga de los cálculos de matrices completas. El equipo verificó mediante pruebas exhaustivas que esta aproximación introduce muy poco error, confirmando que la simplificación no se produce a costa de la precisión.

Además de la velocidad y la estabilidad, el nuevo método reduce significativamente la cantidad de datos que deben transferirse entre diferentes procesadores durante el entrenamiento. En la computación distribuida, donde miles de procesadores trabajan juntos, la comunicación puede convertirse en un cuello de botella importante. MKOR reduce esta complejidad de comunicación de una relación cuadrática a una lineal, lo que significa que a medida que se añaden más procesadores, el sistema escala eficientemente sin verse frenado por los retrasos en la transferencia de datos. Esta eficiencia permite al método manejar los tamaños de lote masivos requeridos por los modelos de lenguaje grandes modernos, lo que era una limitación para los enfoques previos de segundo orden.

Las implicaciones de este trabajo se extienden más allá del entrenamiento de un modelo específico. Al hacer que la optimización de segundo orden sea práctica para sistemas a gran escala, MKOR ofrece un camino hacia un entrenamiento más eficiente de la inteligencia artificial. Los investigadores señalaron que su método podría ayudar a ahorrar cantidades significativas de energía y tiempo para los profesionales del aprendizaje automático. Si bien la tecnología en sí es una herramienta neutral, su aplicación en campos que van desde el diagnóstico médico hasta la modelización climática podría acelerar el descubrimiento científico. Sin embargo, los autores también reconocieron que los algoritmos de optimización eficientes pueden aplicarse a cualquier modelo, incluidos aquellos con impactos sociales negativos, enfatizando que el uso responsable de tal tecnología depende de las decisiones de las personas que la emplean.

En última vez, el artículo presenta una solución concreta a un problema de larga data en el aprendizaje profundo. Al combinar el impulso, las aproximaciones de rango 1 y un mecanismo de cambio híbrido, los investigadores han creado un optimizador que es más rápido, más estable y más escalable que sus predecesores. Los experimentos demuestran que es posible aprovechar el poder de los métodos de segundo orden sin los costos computacionales prohibitivos que los han mantenido al margen. A medida que el campo de la inteligencia artificial continúa avanzando hacia modelos más grandes y complejos, herramientas como MKOR proporcionan una forma de entrenar estos sistemas de manera más eficiente, convirtiendo lo que antes era una ventaja teórica en una realidad práctica.

Resumen Técnico: MKOR – Optimizador Basado en Factores de Kronecker con Impulso (Momentum) mediante Actualizaciones de Rango-1

1. Planteamiento del Problema

Los métodos de optimización de segundo orden, como el Gradiente Natural Descendente (NGD) y la Curvatura Aproximada Factorizada por Kronecker (KFAC), ofrecen tasas de convergencia superiores en comparación con los métodos de primer orden (por ejemplo, SGD, Adam, LAMB) al utilizar información de curvatura (aproximaciones de la matriz de información de Fisher o de la Hessiana). Sin embargo, su adopción en el entrenamiento de Redes Neuronales Profundas (DNN) a gran escala, particularmente en Transformers y Modelos de Lenguaje de Gran Tamaño (LLMs), se ve obstaculizada por tres cuellos de botella principales:

Complejidad Computacional: Los métodos KFAC tradicionales requieren invertir matrices de covarianza de tamaño $d \times d$ (donde $d$ es la dimensión de la capa), lo que resulta en una complejidad de $O(d^3)$ . Del mismo modo, los métodos basados en Sherman-Morrison-Woodbury (SNGD) suelen escalar con el tamaño del lote $b$ ( $O(b^3)$ ). En los Transformers, donde los tamaños de lote escalan con la longitud de la secuencia, estas complejidades se vuelven prohibitivas.
Sobrecarga de Comunicación: El entrenamiento distribuido requiere sincronizar la información de segundo orden (matrices de covarianza y sus inversas) entre los trabajadores. KFAC requiere sincronizar $O(d^2)$ datos, mientras que SNGD requiere $O(bd + b^2)$ , creando cuellos de botella de comunicación significativos que limitan la escalabilidad.
Frecuencia de Actualización y Estabilidad: Debido a los altos costos de inversión, las implementaciones de vanguardia (como KAISA) actualizan la información de segundo orden con poca frecuencia (por ejemplo, cada 100–1000 iteraciones) utilizando factores "obsoletos" (stale). Esto reduce la efectividad de la información de curvatura. Además, los métodos de segundo orden son propensos a la inestabilidad numérica (gradientes explosivos, matrices singulares) y a menudo requieren factores de amortiguación o aritmética de alta precisión, lo que aumenta aún más los costos computacionales.

2. Metodología: El Algoritmo MKOR

Los autores proponen MKOR (Optimizador Basado en Factores de Kronecker con Impulso mediante Actualizaciones de Rango-1), el cual aborda los problemas anteriores a través de cuatro mecanismos centrales:

Aproximación de Rango-1 e Inversión Basada en SM:
En lugar de calcular e invertir matrices de covarianza completas, MKOR aproxima la covarianza de las activaciones y los gradientes utilizando matrices de rango-1 derivadas del promedio de las muestras ( $E[aa^T] \approx \bar{a}\bar{a}^T$ ). Al aplicar la identidad de Sherman-Morrison a estas actualizaciones de rango-1, MKOR actualiza directamente la inversa de los factores de Kronecker. Esto reduce la complejidad computacional de la inversión de factores de $O(d^3)$ a $O(d^2)$ .
- Regla de Actualización: Los factores inversos se actualizan utilizando impulso ( $\gamma$ ) y la fórmula de Sherman-Morrison, permitiendo actualizaciones frecuentes sin el costo de la inversión de matriz completa.
Complejidad de Comunicación Lineal:
En entornos distribuidos, MKOR solo sincroniza los dos vectores de aproximación de rango-1 (tamaño $2d$ ) entre los trabajadores, en lugar de matrices de covarianza completas ( $4d^2$ ). Esto reduce la complejidad de comunicación de cuadrática $O(d^2)$ a lineal $O(d)$ . Además, MKOR utiliza media precisión (FP16) para estos vectores, reduciendo a la mitad los costos de comunicación, una proeza difícil para otros métodos de segundo orden debido a la sensibilidad numérica.
Mecanismos de Estabilización:
Para prevenir la inestabilidad numérica y los gradientes explosivos comunes en los métodos de segundo orden:
- Estabilizador Basado en la Norma: Monitorea la norma de infinito de los factores inversos. Si exceden un umbral, los factores se mezclan con la matriz identidad ( $\hat{L} = \zeta L + (1-\zeta)I$ ) para prevenir la divergencia.
- Reescalado de Gradiente: Los gradientes precondicionados se reescalan para coincidir con la norma de los gradientes originales, asegurando que los programadores de la tasa de aprendizaje (learning rate schedulers) sigan siendo efectivos.
- Estabilidad Teórica: Los autores demuestran (Lema 3.1) que los factores permanecen definidos positivos sin necesidad de factores de amortiguación artificiales ( $\mu I$ ), y que la división escalar única en la regla de actualización garantiza ser distinta de cero.
MKOR Híbrido (MKOR-H):
Reconociendo que las ventajas de segundo orden a menudo disminuyen a medida que el entrenamiento converge (donde el paisaje de pérdida se acerca a la identidad), MKOR-H introduce un mecanismo de conmutación. Monitorea la tasa de reducción de la pérdida y vuelve automáticamente a un optimizador de primer orden si las actualizaciones de segundo orden ya no proporcionan aceleración, optimizando el equilibrio entre costo y velocidad de convergencia.

3. Contribuciones Clave

Reducción de Complejidad: MKOR logra una complejidad computacional de $O(d^2)$ y una complejidad de comunicación de $O(d)$ para actualizaciones de segundo orden, aliviando los cuellos de botella de $O(d^3)$ y $O(d^2)$ encontrados en KFAC y $O(b^3)$ en SNGD.
Actualizaciones de Alta Frecuencia: Al reducir los costos de inversión, MKOR permite actualizaciones de segundo orden en cada iteración (o muy frecuentemente), eliminando la necesidad de factores "obsoletos" y explotando plenamente la información de curvatura.
Robustez: La introducción de la estabilización basada en la norma y el reescalado de gradiente permite que MKOR opere de manera estable sin factores de amortiguación o aritmética de doble precisión, permitiendo el uso eficiente de hardware de media precisión.
Estrategia Híbrida: La variante MKOR-H cambia dinámicamente entre métodos de segundo orden y de primer orden para mantener la eficiencia durante todo el ciclo de vida del entrenamiento.

4. Resultados Experimentales

Los autores evaluaron MKOR en BERT-Large-Uncased utilizando 64 GPUs A100, comparándolo con el optimizador de primer orden de vanguardia LAMB e implementaciones de segundo orden como KAISA (KFAC) y HyLo (SNGD), así como Eva.

Aceleración (Speedup): MKOR logró una aceleración de 2.57× sobre LAMB y una aceleración de 1.85× sobre KAISA en la tarea de pre-entrenamiento. Específicamente, redujo el tiempo de entrenamiento de BERT-Large-Uncased de 8 horas a 3 horas.
Convergencia: En la tarea de respuesta a preguntas SQuAD v1.1, MKOR-H convergió en 600 pasos (frente a 1,536 para LAMB) logrando puntuaciones F1 comparables o mejores (90.64 vs. 90.44).
Benchmark GLUE: MKOR alcanzó nuevas métricas de vanguardia en el conjunto de datos GLUE con 1,500 pasos. Notablemente, mientras que KAISA no logró converger al promedio de precisión base incluso después de 1,563 pasos, MKOR convergió con éxito.
Escalabilidad: A diferencia de HyLo, que divergió o no pudo ajustarse en la memoria de la GPU debido a la alta sobrecarga de memoria ( $O(bd + b^2)$ ), MKOR entrenó con éxito con tamaños de lote grandes (hasta un tamaño de lote efectivo de 64k mediante acumulación) sin problemas de memoria.
Validez de la Aproximación: Los experimentos en ResNet-50 y BERT mostraron que las aproximaciones de rango-1 de las matrices de covarianza incurren en un error insignificante, validando el supuesto de aproximación central.

5. Significado y Reivindicaciones

El artículo afirma que MKOR cierra con éxito la brecha entre los beneficios teóricos de la optimización de segundo orden y las restricciones prácticas del entrenamiento de DNN modernas y a gran escala. Al reducir la complejidad de las actualizaciones de segundo orden a un nivel comparable al de los métodos de primer orden (en términos de escalado con el tamaño del modelo) mientras mantiene las ventajas de convergencia, MKOR permite:

Entrenamiento más Rápido: Reducción significativa en el tiempo de entrenamiento real (wall-clock time) para LLMs.
Mejor Convergencia: La capacidad de converger a métricas de precisión más altas donde otros métodos de segundo orden (como KAISA) fallan o se estancan.
Escalabilidad: La capacidad de escalar a grandes tamaños de lote y dimensiones de modelo sin los cuellos de botella de comunicación y memoria que anteriormente limitaban a los métodos de segundo orden.

Los autores enfatizan que su trabajo hace que la optimización de segundo orden sea viable para la era de los "modelos grandes", abordando específicamente los problemas de escalabilidad introducidos por los tamaños de lote dependientes de la longitud de la secuencia en las arquitecturas Transformer. Concluyen que MKOR representa un nuevo estado del arte para el entrenamiento eficiente y efectivo de redes neuronales profundas.