Lion Optimizer/es

Article
Topic area	optimization
Prerequisites	Stochastic gradient descent, Momentum (optimization), Adam optimizer

Resumen

El optimizador Lion (abreviatura de EvoLved Sign Momentum) es un algoritmo de optimización estocástica de primer orden para entrenar redes neuronales, presentado por Chen et al. en 2023. A diferencia de la mayoría de los optimizadores más utilizados, que fueron diseñados a mano a partir de principios teóricos, Lion fue descubierto mediante búsqueda simbólica de programas en un espacio de programas optimizadores candidatos. El resultado es una regla de actualización conceptualmente más simple que Adam que, aun así, suele igualar o superar su exactitud en cargas de trabajo de aprendizaje profundo a gran escala, utilizando aproximadamente la mitad de la memoria de estado del optimizador.

La característica distintiva de Lion es que la dirección de actualización de los parámetros es el signo elemento a elemento de un término de momento interpolado. Esto hace que cada coordenada de la actualización tenga la misma magnitud (un tamaño de paso fijo escalado por la tasa de aprendizaje), a diferencia de Adam, cuyo escalado adaptativo por coordenada se rige por una media móvil exponencial de los gradientes al cuadrado. Lion ha sido adoptado en diversas tuberías de entrenamiento de gran escala para modelos de lenguaje y visión por computadora, y constituye un punto de referencia para la línea de investigación sobre optimizadores descubiertos automáticamente.

Antecedentes y motivación

A principios de la década de 2020, AdamW se había convertido en el optimizador de facto para el entrenamiento de transformadores, a pesar de limitaciones bien conocidas: almacena dos búferes de momentos por parámetro (primer y segundo momento), lo que duplica la memoria del optimizador más allá de los propios parámetros, y sus hiperparámetros son sensibles a la escala del modelo. Una larga línea de trabajos intentó reemplazar a Adam de forma manual, con resultados empíricos dispares: LAMB, Adafactor, Sophia, Shampoo y otros abordaron problemas específicos, pero no lograron desplazar a Adam como predeterminado.

Lion surgió a partir de una premisa distinta: en lugar de diseñar un optimizador, buscarlo. Los autores formalizaron un optimizador como un programa corto que opera sobre gradientes y estado, definieron un espacio de búsqueda de operaciones primitivas (sumas, multiplicaciones, signo, medias móviles exponenciales, recorte, etc.) y emplearon búsqueda evolutiva de programas con tareas de entrenamiento sustitutas al estilo de metaaprendizaje para evaluar candidatos. Lion es el programa de alto rendimiento más simple que emergió tras una poda sustancial de candidatos descubiertos de mayor tamaño.

Algoritmo

Sea $ g_t $ el gradiente de la función de pérdida respecto a los parámetros $ \theta_t $ en el paso $ t $, sea $ m_t $ un único búfer de momento inicializado a cero, y sea $ \eta $ la tasa de aprendizaje. Lion tiene dos coeficientes de interpolación $ \beta_1, \beta_2 \in [0,1) $ y un coeficiente de decaimiento de pesos $ \lambda \ge 0 $. La actualización en el paso $ t $ es:

$ {\displaystyle c_t = \beta_1 \, m_{t-1} + (1 - \beta_1) \, g_t} $

$ {\displaystyle \theta_t = \theta_{t-1} - \eta \, \big( \operatorname{sign}(c_t) + \lambda \, \theta_{t-1} \big)} $

$ {\displaystyle m_t = \beta_2 \, m_{t-1} + (1 - \beta_2) \, g_t} $

La primera línea forma una dirección de actualización $ c_t $ interpolando entre el momento anterior y el gradiente actual. La segunda línea aplica la actualización de parámetros utilizando solo el signo elemento a elemento de $ c_t $, más decaimiento de pesos desacoplado al estilo de AdamW. La tercera línea actualiza entonces el búfer de momento persistente $ m_t $ con un coeficiente de interpolación distinto. Crucialmente, el búfer empleado para la dirección del siguiente paso ($ m_t $) y el búfer empleado para la dirección actual (una mezcla que involucra a $ m_{t-1} $) se rigen por valores de $ \beta $ separados.

Los hiperparámetros predeterminados reportados por Chen et al. son $ \beta_1 = 0.9 $, $ \beta_2 = 0.99 $ y una tasa de aprendizaje aproximadamente un orden de magnitud menor que la de AdamW, con un decaimiento de pesos aproximadamente un orden de magnitud mayor.

Intuición

Tres observaciones ayudan a explicar por qué Lion funciona.

Primero, cada actualización tiene magnitud unitaria por coordenada. Dado que la actualización es $ \eta \cdot \operatorname{sign}(c_t) $, ninguna coordenada puede dar un paso mayor que $ \eta $. Este recorte implícito se asemeja al recorte de gradiente, pero se aplica a nivel de cada coordenada y en cada paso, aportando robustez frente a los gradientes atípicos habituales en el entrenamiento de modelos de lenguaje.

Segundo, la estructura de dos $ \beta $ desacopla el suavizado de la dirección de la acumulación en el búfer. Con $ \beta_1 < \beta_2 $, la dirección $ c_t $ reacciona más al gradiente actual que el búfer persistente $ m_t $. Esto dota a Lion de una anticipación intrínseca tipo Nesterov: el paso dado en el instante $ t $ refleja más a $ g_t $ de lo que lo haría el búfer de momento por sí solo.

Tercero, tomar el signo descarta la información de magnitud del gradiente, que Adam aprovecha mediante su segundo momento. El hecho empírico de que Lion siga entrenando bien sugiere que para muchas redes neuronales grandes, la dirección importa más que un escalado preciso de la magnitud, siempre que la tasa de aprendizaje y el decaimiento de pesos se ajusten de forma compatible.

Memoria y cómputo

Lion almacena un único búfer de momento $ m_t $ por parámetro. Adam y AdamW almacenan dos (el primer y el segundo momento). Para un modelo con $ P $ parámetros entrenado en bfloat16 o precisión de 32 bits, esto reduce a la mitad la memoria de estado del optimizador: de $ 2P $ a $ P $ valores. Para modelos con decenas o cientos de miles de millones de parámetros, este ahorro es significativo en el entrenamiento distribuido bajo presupuestos fijos de memoria de acelerador.

El cómputo por paso de Lion también es ligeramente más barato: una operación de signo y dos interpolaciones, sin división por la raíz cuadrada de una estimación del segundo momento. En la práctica, el tiempo por paso está dominado por los pases hacia adelante y hacia atrás, por lo que la aceleración práctica derivada de una aritmética del optimizador más barata es modesta, aunque no nula.

Ajuste de hiperparámetros

Los hiperparámetros de Lion no se transfieren directamente desde Adam. Los ajustes recomendados, observados empíricamente por los autores y corroborados en reproducciones independientes, son:

• Tasa de aprendizaje: aproximadamente entre $ 3\times $ y $ 10\times $ menor que el valor de AdamW que funciona para el mismo modelo. Dado que la actualización de Lion tiene magnitud unitaria por coordenada, el tamaño de paso efectivo es más predecible, pero menor en términos absolutos que el escalado adaptativo de Adam en coordenadas bien condicionadas.
• Decaimiento de pesos: aproximadamente entre $ 3\times $ y $ 10\times $ mayor que el valor de AdamW, para compensar la menor tasa de aprendizaje preservando la fuerza efectiva de regularización $ \eta \lambda $.
• Tamaño de lote: Lion se beneficia de tamaños de lote mayores más que Adam en algunos regímenes, posiblemente porque las actualizaciones basadas en el signo tienen mayor varianza por paso y promedian de forma más limpia con lotes más grandes.
• $ \beta_1, \beta_2 $: los valores predeterminados $ (0.9, 0.99) $ son robustos; el artículo original observó ganancias menores con ajuste, pero recomienda partir de los valores predeterminados.

Rendimiento empírico

En el artículo original, Lion fue evaluado en clasificación de imágenes sobre ImageNet, aprendizaje contrastivo imagen-texto (estilo CLIP), entrenamiento de modelos de difusión y preentrenamiento de modelos de lenguaje de hasta varios miles de millones de parámetros. En todos estos escenarios, Lion igualó o superó modestamente a AdamW en exactitud de validación o perplejidad, normalmente con un tiempo de entrenamiento de pared comparable.

Las reproducciones independientes en modelos más pequeños han confirmado en general que Lion es competitivo, pero no uniformemente superior. El resultado depende de la arquitectura del modelo, el tamaño de lote y el esfuerzo de ajuste. Algunos escenarios de ajuste fino, en particular aquellos con conjuntos de datos muy pequeños, han mostrado que Lion es más sensible que AdamW a la elección del programa de tasa de aprendizaje.

Variantes y trabajos posteriores

Han aparecido varias variantes y análisis desde 2023:

• Tiger aplica una idea similar de signo de momento con un único $ \beta $ y se ha reportado que es aún más eficiente en memoria, a costa de una pérdida modesta de exactitud.
• Lion-K sustituye el signo elemento a elemento por una no linealidad más general, con un análisis teórico que conecta a Lion con una interpretación regularizada de descenso espejo.
• Se han desarrollado análisis en tiempo continuo y de convergencia de Lion, que ven al algoritmo como una aproximación en tiempo discreto a una determinada ecuación diferencial ordinaria con no linealidad de signo, lo que proporciona una justificación teórica parcial de la robustez observada.

Estos trabajos posteriores coinciden en general en que el comportamiento de Lion está regido por la interacción entre la no linealidad de signo y la media móvil exponencial del momento, más que por cualquiera de los componentes por separado.

Comparación con optimizadores relacionados

• vs. signSGD: signSGD toma el signo del gradiente sin pasar por un búfer de momento. Lion es signSGD con un esquema particular de momento de dos $ \beta $; el momento es esencial para el rendimiento empírico de Lion.
• vs. Adam / AdamW: Lion utiliza un búfer en lugar de dos y descarta la información de magnitud. Adam adapta la escala por coordenada mediante el segundo momento; Lion no.
• vs. SGD con momento: SGD con momento utiliza el vector de momento sin transformar como dirección. Lion utiliza su signo, acotando la magnitud de actualización por coordenada.
• vs. Adafactor: Adafactor reduce la memoria del optimizador factorizando el segundo momento; Lion elimina el segundo momento por completo.

Limitaciones

Lion no es un reemplazo universal de Adam. Entre las limitaciones reportadas se incluyen:

• Sensibilidad a los hiperparámetros: dado que la actualización de Lion se basa en el signo, cambios sutiles en la tasa de aprendizaje, el decaimiento de pesos y el calentamiento pueden alterar el comportamiento de forma más abrupta que en Adam.
• Regímenes de lotes pequeños: las actualizaciones basadas en el signo tienen mayor varianza por paso; con lotes muy pequeños, el escalado adaptativo por coordenada de Adam puede actuar como un estabilizador del que Lion carece.
• Gradientes dispersos: en modelos con una estructura de gradiente disperso intrínseca (algunos sistemas de recomendación, tablas de embeddings), la operación de signo puede amplificar el ruido en coordenadas que rara vez se actualizan.
• Garantías teóricas limitadas: aunque existen análisis posteriores, Lion carece de la misma profundidad de teoría de convergencia convexa que SGD o las variantes de Adam.

Historia

Lion se presentó en el artículo de 2023 "Symbolic Discovery of Optimization Algorithms" de Chen et al. en Google. El optimizador fue el producto destacado de una metodología más amplia de búsqueda simbólica que, según los autores, podía aplicarse a otros componentes de las tuberías de entrenamiento de aprendizaje automático. Aparecieron implementaciones de código abierto en PyTorch, JAX y Optax a los pocos meses de su publicación, y el algoritmo se incorporó a varios sistemas de entrenamiento a gran escala, incluido el preentrenamiento de modelos de lenguaje y la generación de imágenes.