Resumen Técnico: AMUSE (Muon en Cualquier Momento con Evaluación de Gradiente Estable)

Enunciado del Problema

La optimización moderna de aprendizaje profundo suele depender del optimizador AdamW combinado con programas de tasa de aprendizaje prescritos (por ejemplo, decaimiento coseno). Avances recientes han cuestionado estos componentes individualmente: la optimización Schedule-Free (SF) elimina los programas explícitos evaluando gradientes en puntos interpolados entre iteraciones actuales y promediadas, mientras que Muon mejora la geometría de la actualización para parámetros matriciales mediante la ortogonalización del momento.

A pesar del éxito empírico de Muon en el preentrenamiento de modelos de lenguaje grandes (LLM), su mecanismo subyacente no se comprende totalmente. El artículo identifica una inestabilidad específica en Muon: mientras que su ortogonalización acelera el progreso a lo largo de direcciones "planas" (el "río" del paisaje de pérdida), amplifica inadvertidamente el ruido en direcciones "empinadas" de alta curvatura (las "paredes del valle"), lo que conduce a trayectorias oscilatorias. Por el contrario, aunque los métodos Schedule-Free estabilizan las trayectorias manteniéndolas cerca del río, pueden carecer del progreso masivo rápido que ofrece la ortogonalización de Muon. El problema central es diseñar un optimizador que combine la rápida convergencia de Muon con la estabilidad del promediado Schedule-Free sin requerir programas manuales de tasa de aprendizaje.

Metodología: AMUSE

Los autores proponen AMUSE (Muon en Cualquier Momento con Evaluación de Gradiente Estable), que integra las actualizaciones matriciales ortogonalizadas de Muon con el marco Schedule-Free utilizando un coeficiente de interpolación variable en el tiempo.

Fundamento Teórico: Paisaje de Pérdida Río-Valle

El análisis se basa en la perspectiva del paisaje de pérdida río-valle:

• El Río: Un subespacio masivo de alta dimensión y baja curvatura donde ocurre el progreso útil del entrenamiento.
• Las Paredes del Valle: Un subespacio dominante de baja dimensión y alta curvatura que induce oscilaciones.
• Mecanismo de Muon: Ortogonalizar la matriz de momento aumenta el componente masivo relativo de las actualizaciones, acelerando el progreso a lo largo del río. Sin embargo, también impide que los componentes dominantes decaigan, lo que hace que el ruido pequeño en direcciones dominantes se amplifique hasta convertirse en oscilaciones a través de las paredes del valle.
• Mecanismo Schedule-Free: Evaluar gradientes en una trayectoria promediada (más cerca del río) filtra las fluctuaciones de alta curvatura antes del paso de actualización.

Diseño del Algoritmo

AMUSE emplea un coeficiente de interpolación variable en el tiempo $\beta_t$ para determinar el punto de evaluación del gradiente $Y_t$ entre la secuencia base rápida $Z_t$ (actualizaciones de Muon) y la secuencia promediada estable $X_t$:
$$Y_t = (1 - \beta_t)Z_t + \beta_t X_t$$

El programa para $\beta_t$ está diseñado para equilibrar la adaptación temprana y la estabilidad en etapas tardías:

1. Fase de Calentamiento ($t \le T_0$): $\beta_t$ se fija en un valor pequeño $\beta_1$. Esto permite que el optimizador evalúe gradientes cerca de la secuencia rápida de Muon ($Z_t$) para una adaptación inicial rápida.
2. Post-Calentamiento ($t > T_0$): $\beta_t$ aumenta gradualmente hacia 1 según:
   $$\beta_t = 1 - \left( \frac{T_0 - 1}{t - 1} \right)^\rho (1 - \beta_1)$$
   Aquí, $\rho \in [0, 1]$ controla la velocidad del cambio. A medida que $\beta_t$ aumenta, el punto de evaluación del gradiente se desplaza hacia la secuencia promediada $X_t$. Esto suprime los componentes de alta curvatura (dominantes) del gradiente antes de que sufran la ortogonalización de Muon, estabilizando así la trayectoria.

Crucialmente, este programa es independiente del horizonte total de entrenamiento, lo que permite el entrenamiento en cualquier momento (el entrenamiento puede detenerse en cualquier punto sin degradación del rendimiento). AMUSE no requiere un programa de decaimiento de la tasa de aprendizaje.

Contribuciones Clave

1. Análisis Geométrico de Muon: El artículo proporciona una explicación geométrica de la eficiencia e inestabilidad de Muon. Demuestra que la ortogonalización promueve actualizaciones orientadas al masivo (acelerando el progreso del río) pero amplifica el ruido en la dirección dominante (causando oscilaciones en las paredes del valle).
2. Estabilización mediante Interpolación: Los autores muestran que evaluar gradientes más cerca de la trayectoria promediada (mediante Schedule-Free) reduce los componentes dominantes antes de la ortogonalización, mitigando las oscilaciones.
3. Algoritmo AMUSE: Un optimizador novedoso que desplaza dinámicamente el punto de evaluación del gradiente desde la trayectoria rápida de Muon hacia la trayectoria promediada estable. Esto combina la rápida convergencia de Muon con la estabilidad de los métodos Schedule-Free.
4. Soporte para Entrenamiento en Cualquier Momento: Al utilizar un programa variable en el tiempo independiente de los pasos totales y evitar el decaimiento de la tasa de aprendizaje, AMUSE soporta horizontes de entrenamiento flexibles.

Resultados Experimentales

AMUSE se evaluó en diversas pruebas, incluyendo clasificación de imágenes (CIFAR-10/100, SVHN, ImageNet-1k), segmentación de imágenes (ISIC 2018) y preentrenamiento de LLM (modelos estilo Llama de 124M, 720M y 1B de parámetros en FineWeb).

• Frente de Pareto Rendimiento-Iteración: AMUSE supera consistentemente a AdamW (Schedule-Free) y Muon en todas las tareas.
• Ganancias de Eficiencia:
  • En el preentrenamiento de LLM de 720M, AMUSE alcanza el rendimiento final de Muon con 1.51$\times$ menos pasos.
  • En ImageNet con ResNet-50, logra el mismo rendimiento con 1.12$\times$ menos pasos.
  • En el ajuste fino de ViT en ImageNet, requiere 3.08$\times$ menos pasos.
• Estabilidad: A diferencia de Muon puro, que se beneficia significativamente del promediado exponencial de pesos (EWA) posterior o del decaimiento de la tasa de aprendizaje, AMUSE mantiene una trayectoria estable durante todo el entrenamiento, mostrando una mejora mínima cuando se aplican estas técnicas posteriores.
• Escalabilidad: AMUSE escala eficazmente a tamaños de lote grandes (hasta 2048) y escalas de modelo grandes (1B de parámetros), donde otras variantes Schedule-Free a veces tienen dificultades.

Significado y Afirmaciones

El artículo afirma que AMUSE representa un enfoque principiado para el diseño de optimizadores derivado de una comprensión sistemática de la dinámica de entrenamiento y la geometría del paisaje de pérdida. Al abordar explícitamente la compensación entre el progreso masivo rápido y la estabilidad de las paredes del valle, AMUSE logra un rendimiento superior sin necesidad de programación manual de la tasa de aprendizaje.

Los autores notan modestamente que, aunque AMUSE supera a los optimizadores existentes, incurre en una ligera sobrecarga de memoria (una copia adicional de estado para la secuencia base $Z_t$) en comparación con Muon puro, aunque sigue siendo comparable a AdamW y AdamW Schedule-Free. Sugieren que el trabajo futuro podría centrarse en lograr una estabilidad y rendimiento similares sin este costo adicional de memoria. El artículo enfatiza que la perspectiva "río-valle" ofrece un marco robusto para diseñar optimizadores que sean tanto eficientes como estables.