Muon-OGD: Muon-based Spectral Orthogonal Gradient Projection for LLM Continual Learning

Resumen Técnico: Muon-OGD

Enunciado del Problema
El desafío central abordado en este trabajo es el olvido catastrófico en Modelos de Lenguaje Grande (LLM) durante el aprendizaje continuo. Cuando los LLM se ajustan finamente de forma secuencial en nuevas tareas, su rendimiento en tareas previamente aprendidas a menudo se degrada rápidamente. Los métodos existentes de aprendizaje continuo basados en proyección (por ejemplo, OGD, GEM, O-LoRA) mitigan esto restringiendo las actualizaciones de parámetros a subespacios ortogonales a las direcciones asociadas con tareas pasadas. Sin embargo, estos métodos se formulan típicamente bajo una geometría de parámetros euclidiana, donde las magnitudes de actualización y la ortogonalidad están gobernadas por la norma de Frobenius.

Los autores identifican una discrepancia fundamental: los optimizadores modernos para parámetros matriciales, como el recientemente propuesto optimizador Muon, operan bajo una geometría de norma espectral (norma de operador). Muon aplica actualizaciones matriciales ortogonalizadas (aproximadas mediante iteraciones de Newton–Schulz) y puede interpretarse como realizando un descenso de gradiente más pronunciado bajo una restricción de norma espectral. El artículo argumenta que confiar en la geometría de norma de Frobenius para las restricciones de aprendizaje continuo es subóptimo para los parámetros matriciales de los LLM, ya que no se alinea con la geometría de la dinámica de optimización subyacente.

Metodología: Muon-OGD
El artículo propone Muon-OGD (Proyección de Gradiente Ortogonal Espectral basada en Muon), un marco que integra la optimización consciente de la norma espectral con restricciones de aprendizaje continuo basadas en proyección.

1. Formulación: En lugar del problema estándar restringido por norma de Frobenius, Muon-OGD formula la actualización $\Delta$ como un problema de optimización restringido por norma espectral:
   $$\min_{\Delta \in \mathbb{R}^{m \times n}} \langle G, \Delta \rangle \quad \text{sujeto a} \quad \|\Delta\|_2 \leq \eta, \quad \langle C_i, \Delta \rangle = 0, \forall i=1,\dots,k$$
   Aquí, $G$ es la matriz de gradiente (o momento), $\|\cdot\|_2$ es la norma espectral, $\eta$ es el tamaño del paso, y $\{C_i\}$ son direcciones de restricción (por ejemplo, gradientes de tareas pasadas) que deben permanecer ortogonales a la actualización.

2. Optimización Dual: Dado que el problema carece de una solución de forma cerrada, los autores derivan una formulación dual. El objetivo dual implica minimizar la norma nuclear de una matriz desplazada $H = G + \sum \lambda_i C_i$. La actualización óptima se recupera mediante la función de signo de matriz: $\Delta^* = -\eta \cdot \text{msgn}(H)$.

3. Algoritmo: El método resuelve el problema dual iterativamente utilizando descenso de subgradiente sobre las variables duales $\lambda$.

   • Bucle Interno: Actualiza iterativamente $\lambda$ para hacer cumplir las restricciones de ortogonalidad.
   • Aproximación de Signo de Matriz: La función de signo de matriz $\text{msgn}(H)$ se aproxima eficientemente utilizando un pequeño número de iteraciones de Newton–Schulz (específicamente 5 iteraciones, denotadas como $NS_5$), evitando el costo computacional de la SVD explícita.
   • Instanciación de Bajo Rango: Para mayor eficiencia, los autores proponen una instanciación específica donde las restricciones se derivan de los $k$ vectores singulares principales de las tareas pasadas, permitiendo que las variables duales se organicen en una matriz de bajo rango $\Lambda$, reduciendo la dimensionalidad del problema dual.

Contribuciones Clave

1. Geometría de Norma Espectral para Aprendizaje Continuo: El artículo formula el aprendizaje continuo basado en proyección bajo geometría de norma espectral, alineando las restricciones de estabilidad del aprendizaje continuo con la geometría de norma de operador de optimizadores modernos como Muon.
2. Regla de Actualización Práctica: Deriva un algoritmo práctico (Muon-OGD) que combina restricciones de ortogonalidad de tareas pasadas con actualizaciones de signo de matriz al estilo de Muon, resolviendo el problema restringido de manera eficiente mediante una formulación dual y aproximaciones de Newton–Schulz.
3. Validación Empírica: El marco se instancia para el ajuste fino continuo de LLM y se evalúa en diversos puntos de referencia, demostrando mejoras consistentes sobre contrapartes basadas en norma de Frobenius.

Resultados Experimentales
Los autores evalúan Muon-OGD en:

• Puntos de Referencia Estándar: Puntos de referencia de aprendizaje continuo de 5 tareas y 15 tareas utilizando T5-Large.
• Punto de Referencia TRACE: Un punto de referencia para LLM ajustados con instrucciones (usando LLaMA-2-7B-Chat) que cubre diversas habilidades como programación, matemáticas y comprensión multilingüe.
• Currículos Específicos de Dominio: Aprendizaje secuencial en Programación (BigCodeBench) $\to$ Matemáticas (GSM8K) $\to$ Medicina (HuatuoGPT-o1) utilizando los modelos Qwen2.5 y LLaMA3.2.

Hallazgos:

• Rendimiento: Muon-OGD supera consistentemente al ajuste fino secuencial (SeqFT) y a líneas base competitivas de gradiente ortogonal (por ejemplo, O-LoRA, Esculpido de Subespacios, OSFT) en términos de Precisión Promedio (AA) y Transferencia hacia Atrás (BT).
  • En el punto de referencia de 15 tareas de T5-Large, Muon-OGD logró una precisión promedio del 72.0%, superando a O-LoRA (69.6%) y OSFT (71.3%).
  • En el punto de referencia TRACE, Muon-OGD logró la mayor AA (49.4%) y la menor tasa de olvido (-4.7%) entre los métodos comparados.
• Escalabilidad: El método permanece escalable computacionalmente, aprovechando la eficiencia de las iteraciones de Newton–Schulz y las representaciones de restricciones de bajo rango.
• Agnosticismo Arquitectónico: El enfoque es efectivo tanto en arquitecturas codificador-decodificador (T5) como en decodificador único (LLaMA, Qwen).

Significado y Afirmaciones
El artículo afirma que la geometría de actualización consciente de la norma espectral proporciona una alternativa práctica y efectiva a la proyección basada en norma de Frobenius para el aprendizaje continuo en LLM. Al cerrar la brecha entre la geometría de los optimizadores modernos (como Muon) y los requisitos de estabilidad del aprendizaje continuo, Muon-OGD mejora el compromiso entre estabilidad y plasticidad. Los autores sugieren que los métodos existentes basados en proyección pueden actualizarse sistemáticamente para operar bajo geometría de norma espectral, preservando mejor los beneficios de acondicionamiento y estabilidad observados en Muon mientras se previene el olvido catastrófico. El trabajo no afirma resolver todos los desafíos del aprendizaje continuo, sino que ofrece una alternativa geométrica principista que genera ganancias empíricas en retención y adaptación en diversos dominios y escalas de modelos.