Resumen Técnico: ¿Puede Muon afinar modelos preentrenados con Adam?

Enunciado del Problema

Muon (MomentUm Orthogonalizado por Newton-Schulz) ha surgido como un optimizador altamente eficiente para el preentrenamiento de modelos de lenguaje grandes (LLM), ofreciendo aproximadamente un $2\times$ de eficiencia computacional sobre Adam y una reducción en el uso de memoria al eliminar la estimación del segundo momento. Sin embargo, su aplicación práctica en la afinación (fine-tuning) sigue siendo limitada. El principal obstáculo es el problema de incompatibilidad de optimizadores: la mayoría de los modelos abiertos disponibles públicamente están preentrenados con Adam, y cambiar ingenuamente a Muon para la afinación conduce a un rendimiento degradado. Por el contrario, utilizar Adam para afinar modelos preentrenados con Muon también produce resultados subóptimos. Esta incompatibilidad restringe severamente la adopción de Muon en el paradigma dominante de afinación, a pesar de sus éxitos en el preentrenamiento.

Metodología y Análisis

Los autores investigan esta incompatibilidad mediante experimentos controlados y análisis teórico, centrándose en los sesgos implícitos distintos de Adam y Muon.

1. Caracterización del Sesgo Implícito:

   • Adam utiliza un precondicionamiento elemento a elemento basado en los primeros y segundos momentos, convergiendo a soluciones que maximizan el margen de norma máxima ($\|W\|_{\max}$).
   • Muon aplica un precondicionamiento a nivel de matriz mediante ortogonalización de Newton-Schulz, convergiendo a soluciones que maximizan el margen de norma espectral ($\|W\|_2$).
   • El análisis teórico sobre un problema de regresión lineal simplificado demuestra que estos optimizadores convergen a soluciones fundamentalmente diferentes ($W^*_{\max}$ vs. $W^*_2$). Empíricamente, los pesos entrenados con Muon exhiben un rango estable más alto y distribuciones de valores singulares más uniformes en comparación con los pesos entrenados con Adam.

2. Mecanismo de Degradación:

   • La incompatibilidad altera el conocimiento preentrenado. Cuando se aplica un optimizador con un sesgo implícito diferente, fuerza a los pesos del modelo hacia una dirección estructural incompatible con la fase de preentrenamiento.
   • Los experimentos muestran que la afinación con incompatibilidad aumenta la sensibilidad a la fuerza de las actualizaciones. Actualizaciones más fuertes causan una mayor alteración de la estructura preentrenada, lo que conduce a una mayor perplejidad y un olvido catastrófico más severo.

3. Mitigación Propuesta: LoRA:

   • Los autores hipotetizan que restringir la magnitud de las actualizaciones puede mitigar esta incompatibilidad. Emplean Adaptación de Bajo Rango (LoRA), que congela los pesos preentrenados ($W_0$) y restringe las actualizaciones a un subespacio de bajo rango ($\Delta W = BA$).
   • Este enfoque limita la magnitud de los cambios en los pesos, evitando que el optimizador altere agresivamente la estructura preentrenada. El análisis teórico sugiere que la "inflación de incompatibilidad" (la brecha entre el rendimiento coincidente y el no coincidente) está acotada por el rango de LoRA $r$, desapareciendo cuando $r \to 1$.

Contribuciones Clave

• Reproducción y Análisis: El artículo reproduce el fenómeno de incompatibilidad de optimizadores a escalas accesibles (NanoChat, 561 millones de parámetros) y lo vincula a los sesgos implícitos distintos (norma máxima vs. norma espectral) de Adam y Muon.
• LoRA como Solución: Los autores demuestran que LoRA mitiga eficazmente la incompatibilidad. Al restringir las actualizaciones, LoRA permite que Muon afinar modelos preentrenados con Adam con un rendimiento comparable o superior al de LoRA-Adam.
• Validación Exhaustiva: Los hallazgos se validan en:
  • Comprensión del Lenguaje Natural (NLU): Puntuación GLUE con T5-Base.
  • Generación de Lenguaje Natural (NLG): Ajuste de instrucciones de Llama 2-7B/13B en matemáticas, código y razonamiento de sentido común.
  • Visión por Computadora: Afinación de CLIP ViT-B/32 en seis tareas de clasificación de imágenes.
• Estudios de Rango y Olvido:
  • Sensibilidad al Rango: El beneficio de LoRA-Muon es más pronunciado en rangos bajos a moderados. A medida que el rango aumenta (acercándose a la afinación completa), la incompatibilidad reaparece y el rendimiento se degrada en relación con Adam.
  • Olvido Catastrófico: La afinación completa con Muon en modelos preentrenados con Adam causa un olvido más severo del conocimiento preentrenado que Adam. LoRA reduce significativamente este olvido, preservando la estructura preentrenada.
• Variantes de LoRA: El estudio evalúa variantes existentes de LoRA (por ejemplo, rsLoRA, PiSSA, LoRA-One). Descubre que las variantes diseñadas para hacer que las actualizaciones de LoRA se acerquen más a la afinación completa (como PiSSA) no necesariamente benefician a Muon e incluso pueden exacerbar la incompatibilidad, lo que sugiere la necesidad de adaptaciones específicas para Muon.

Resultados

• Reducción de la Brecha de Rendimiento: En la afinación completa, Muon rinde consistentemente menos que Adam en modelos preentrenados con Adam (por ejemplo, una brecha de perplejidad de +0.009 a +0.023 en WikiText-2). Bajo LoRA, esta brecha se reduce significativamente (un 39% para modelos preentrenados con Muon y un 78% para modelos preentrenados con Adam).
• Superioridad en LoRA: En tareas de NLU, NLG y visión, LoRA-Muon iguala o supera a LoRA-Adam. Por ejemplo, en la puntuación GLUE, LoRA-Muon-PE logró la precisión promedio más alta (89.20%), superando tanto a LoRA-Adam (88.93%) como a Full-Adam (89.14%).
• Eficiencia: Aunque Muon incurre en una sobrecarga modesta por paso debido a la iteración de Newton-Schulz (1.1–1.2$\times$ más lento que Adam en configuraciones LoRA), su eficiencia de memoria (almacenar solo un búfer de momento frente a dos para Adam) sigue siendo una ventaja práctica, ahorrando aproximadamente el 50% de la memoria del estado del optimizador.

Significado y Afirmaciones

El artículo afirma proporcionar el primer análisis en profundidad del problema de incompatibilidad de optimizadores entre Adam y Muon. Su principal significado radica en habilitar el uso práctico de Muon para la afinación del vasto ecosistema de modelos preentrenados con Adam.

Los autores concluyen que:

1. La incompatibilidad de optimizadores es estructural: Surge de sesgos implícitos conflictivos que alteran el conocimiento preentrenado cuando la fuerza de las actualizaciones no está restringida.
2. Las actualizaciones restringidas son la clave: Técnicas como LoRA, que limitan la magnitud de los cambios en los pesos, neutralizan eficazmente esta incompatibilidad.
3. Recomendación Práctica: Los practicantes pueden utilizar Muon como un reemplazo directo de Adam en la afinación con LoRA para obtener eficiencia de memoria y rendimiento competitivo sin sufrir las penalizaciones de incompatibilidad observadas en la afinación completa.

El artículo nota modestamente que, aunque LoRA mitiga el problema, las variantes existentes de LoRA optimizadas para Adam no se transfieren automáticamente a Muon, lo que sugiere que se necesita trabajo futuro para desarrollar adaptaciones específicas para Muon. Además, la caracterización teórica de la incompatibilidad sigue siendo una pregunta abierta.