Resumen Técnico: Promedio Dual Optimista Unifica los Optimizadores Modernos

Enunciado del Problema

La optimización en el aprendizaje profundo ha evolucionado a lo largo de dos ejes principales: geometría y adaptación (diseño de actualizaciones con normas implícitas o precondicionadores que reflejan la estructura del modelo, ej. Adam, Muon) y la composición de ingredientes (combinación de retroalimentación de gradiente, momento y programaciones para la estabilidad, ej. Lion, métodos libres de programación). Aunque muchos optimizadores de última generación (como Muon, Lion, AdEMAMix y NAdam) parecen distintos, carecen de un marco teórico unificado que explique su relación y proporcione orientación fundamentada para el ajuste de hiperparámetros.

Un desafío específico y persistente es el ajuste del decaimiento de pesos. En el entrenamiento de múltiples épocas, el decaimiento de pesos actúa como un regularizador que restringe las normas de las iteraciones. Sin embargo, incluso en el entrenamiento de una sola época, donde el sobreajuste es menos preocupante, el decaimiento de pesos a menudo sigue siendo beneficioso. Las reglas de escalado recientes sugieren establecer el decaimiento de pesos proporcional a $1/d$ (dimensión del modelo), pero es difícil desentrañar la dependencia de la dimensión del modelo $d$ frente al horizonte de entrenamiento $n$. Además, los valores óptimos de decaimiento de pesos varían significativamente entre diferentes optimizadores (por ejemplo, Lion requiere un decaimiento mucho mayor que AdamW), lo que hace necesario un ajuste costoso por cada optimizador.

Metodología: SODA

Los autores introducen SODA (Promedio Dual Optimista Estocástico), una generalización del Promedio Dual Optimista (ODA) que unifica el procesamiento dual de los gradientes y el procesamiento primal de las iteraciones.

Marco Teórico

SODA generaliza la actualización estándar de promedio dual:
$$x_{k+1} \in \arg \min_{x \in X} \gamma_k \langle d_k, x \rangle + h_k(x)$$
mediante la introducción de una secuencia de extrapolación primal $(y_k)$ y una variable dual optimista $(\bar{m}_{k+1})$. El algoritmo mantiene:

1. Procesamiento Dual: Un promedio de gradiente basado en momento $m_{k+1}$ y un promedio optimista $\bar{m}_{k+1}$ que incorpora el gradiente actual $\nabla f(y_k, \xi_k)$.
2. Procesamiento Primal: Una actualización de iteración $x_{k+1}$ y un punto extrapolado $y_{k+1}$, controlados por los parámetros $\lambda_k$ y $\bar{\lambda}_k$.

Las reglas de actualización son:
$$\begin{aligned} m_{k+1} &= (1 - \alpha_k)m_k + \alpha_k \nabla f(y_k, \xi_k) \\ \bar{m}_{k+1} &= (1 - \bar{\alpha}_k)m_{k+1} + \bar{\alpha}_k \nabla f(y_k, \xi_k) \quad (\text{optimismo}) \\ z_{k+1} &\in \partial h_k^*(-\gamma_k \bar{m}_{k+1}) \\ x_{k+1} &= (1 - \lambda_k)x_k + \lambda_k z_{k+1} \\ y_{k+1} &= (1 - \bar{\lambda}_k)x_{k+1} + \bar{\lambda}_k z_{k+1} \quad (\text{extrapolación primal}) \end{aligned}$$
donde $h_k$ es un regularizador (mapa de espejo) que define la geometría (por ejemplo, $\ell_\infty$ para Lion, espectral para Muon).

Unificación de Optimizadores Existentes

El artículo demuestra que, al seleccionar parámetros y geometrías específicos, SODA recupera varios optimizadores modernos como casos especiales:

• NAdam / Lion / Muon: Estos se recuperan bajo el régimen "Modernizado" donde $\bar{\lambda}_k = 0$ (sin extrapolación primal para el punto de evaluación del gradiente) y $\bar{\alpha}_k > 0$ (optimismo).
  • Lion: Utiliza geometría $\ell_\infty$ ($h_k = K^*$) y parámetros de momento específicos.
  • Muon: Utiliza geometría espectral y momento de Nesterov ($\bar{\alpha}_k = \alpha_k$).
  • NAdam: Utiliza una geometría $\ell_\infty$ suavizada con escalado adaptativo.
• La Tabla 1 en el artículo mapea explícitamente estos métodos a sus elecciones correspondientes de promediado de gradientes ($\alpha_k$), optimismo ($\bar{\alpha}_k$) y geometría ($h_k$).

El Envoltorio SODA

Una contribución práctica clave es un Envoltorio SODA (Algoritmo 1) que puede envolver cualquier optimizador base (por ejemplo, Adam, Lion, Muon) sin introducir nuevos hiperparámetros.

• Configuración: Establece $\bar{\lambda}_k = 0$ (coincidiendo con el régimen "modernizado"), utiliza un promediado uniforme de iteraciones $\lambda_k = 1/(k+2)$ y escala el tamaño de paso $\gamma_k = (k+2)\eta_k$.
• Eliminación del Decaimiento de Pesos: El envoltorio ancla las actualizaciones a la inicialización $z_0$ en lugar del origen. La regla de actualización resultante implementa efectivamente una programación de decaimiento de pesos de $\mu_k = 1/(k+2)$.
• Interpretación: Los autores argumentan que en este marco, el decaimiento de pesos no es meramente un regularizador, sino una forma de promedio primal. Este decaimiento $1/k$ fundamentado teóricamente elimina la necesidad de ajustar un hiperparámetro de decaimiento de pesos separado.

Contribuciones Clave

1. Unificación mediante Optimismo y Dualización: SODA proporciona una plantilla única que recupera Muon, Lion, NAdam y otros como instancias optimistas del Promedio Dual, aclarando sus relaciones estructurales.
2. Garantías Teóricas: El marco deriva elecciones de hiperparámetros (específicamente la programación de decaimiento $1/k$) que están fundamentadas teóricamente en el análisis de convergencia para funciones convexas $L$-suaves. Replantea el decaimiento de pesos como un mecanismo para el promedio primal.
3. Envoltorio Libre de Parámetros: El envoltorio SODA propuesto elimina la necesidad de ajustar el decaimiento de pesos para cualquier optimizador base, manteniendo o mejorando el rendimiento.
4. Validación Empírica: Experimentos extensos muestran mejoras consistentes a través de tamaños de modelo y horizontes de entrenamiento.

Resultados Experimentales

Se realizaron experimentos en nanoGPT utilizando el conjunto de datos FineWeb100, comparando optimizadores envueltos con SODA frente a líneas base ajustadas.

• Transferencia a través del Horizonte: Se encontró que el decaimiento de pesos óptimo para Muon disminuye a medida que aumenta el horizonte de entrenamiento, contradiciendo las reglas de escalado estáticas (como $1/d$). El decaimiento dinámico $1/(k+2)$ del envoltorio SODA superó consistentemente a Muon, incluso cuando Muon recibió su decaimiento de pesos mejor ajustado (Figura 1).
• Aplicabilidad General: El envoltorio mejoró el rendimiento de Adam, Muon y Scion a través de diferentes horizontes de entrenamiento (1$\times$ y 4$\times$ pasos Chinchilla) sin ningún ajuste adicional (Figura 2).
• Escalado: Al escalar el tamaño del modelo de 64M a 1B de parámetros (siguiendo el escalado Chinchilla), SODA (específicamente la variante optimista SODA$^\dagger$) superó consistentemente a los optimizadores base (Scion/uScion), ampliándose la brecha de rendimiento a medida que aumentaba el tamaño del modelo (Figura 4).
• Simplificación: La configuración SODA$^\dagger$ simplificó la línea base Muon reemplazando Adam con Lion para todas las capas y utilizando un único hiperparámetro de momento, logrando aún así un rendimiento competitivo o superior.

Significado y Afirmaciones

El artículo afirma que SODA ofrece un mecanismo fundamentado para transferir estrategias de optimización a través de diferentes horizontes de entrenamiento y escalas de modelos sin necesidad de ajuste de hiperparámetros.

• Replanteamiento del Decaimiento de Pesos: Los autores proponen un cambio de perspectiva: el decaimiento de pesos debe entenderse no solo como un regularizador contra el sobreajuste, sino como una forma de promedio primal que ancla la trayectoria de optimización a la inicialización.
• Impacto Práctico: Al unificar los optimizadores modernos bajo el marco SODA, los autores proporcionan un envoltorio "plug-and-play" que elimina el proceso tedioso y costoso de ajustar el decaimiento de pesos, lo cual es particularmente crítico a medida que los tamaños de los modelos hacen prohibitivamente costoso el ajuste manual.
• Modestia: Los autores reconocen limitaciones, señalando que su análisis teórico está actualmente restringido al entorno convexo (aunque empíricamente informativo para el aprendizaje profundo) y se centra en el entrenamiento de una sola época. Sugieren que el entrenamiento de múltiples épocas podría seguir requiriendo un hiperparámetro para la fuerza de regularización debido a preocupaciones sobre el sobreajuste.

En resumen, SODA unifica el espacio de diseño de los optimizadores modernos a través de la lente del promedio dual optimista y proporciona un método fundamentado teóricamente y libre de parámetros para mejorar la estabilidad y el rendimiento del entrenamiento a través de escalas.