Resumen Técnico: Repensar Muon más allá del Preentrenamiento

Enunciado del Problema

El artículo identifica limitaciones fundamentales de Muon (MomentUm Orthogonalized by Newton–Schulz), un optimizador consciente de la matriz que ha logrado éxitos en el preentrenamiento de Modelos de Lenguaje Grandes (LLM), cuando se aplica a dos regímenes de entrenamiento emergentes: el entrenamiento Visión-Lenguaje-Acción (VLA) y el Aprendizaje por Refuerzo con Recompensas Verificables (RLVR).

Muon opera impulsando todos los valores singulares de la matriz de momento hacia 1 mediante iteraciones de Newton–Schulz (NS), realizando efectivamente un blanqueamiento espectral uniforme. Aunque esto promueve la exploración en el preentrenamiento, los autores argumentan que induce un "desajuste espectral" en tareas posteriores:

1. En el Entrenamiento VLA: Los gradientes del módulo de acción son inherentemente de bajo rango (bajo rango efectivo, o erank). El blanqueamiento uniforme de Muon amplifica las direcciones de la cola ruidosa (que dominan el espectro de bajo rango) hasta la misma magnitud que los valores singulares principales informativos, corrompiendo la actualización.
2. En el Post-entrenamiento RLVR: Los gradientes de la política exhiben una baja Relación Señal-Ruido (SNR) debido a recompensas a nivel de trayectoria y mecanismos de estabilización (por ejemplo, recorte, muestreo por importancia). El blanqueamiento uniforme de Muon amplifica este ruido estocástico, lo que lleva al colapso de la política. Además, Muon trata las proyecciones de atención como bloques únicos, ignorando la heterogeneidad por cabeza (escalas de actualización especializadas) heredada del preentrenamiento.

Variantes existentes como Low-Rank Muon (LRMuon) abordan problemas de rango pero dependen de SVD o esquematización computacionalmente costosa, fallando en igualar la escalabilidad de Muon.

Metodología: Pion

Los autores proponen Pion (sPectral hIgh-pass Optimization on momeNtum), un reemplazo directo para Muon que mantiene su eficiencia computacional (costo por paso idéntico a Muon) pero sustituye el blanqueamiento uniforme por un filtro paso alto espectral.

Mecanismo Central: Iteración NS Paso Alto

Pion reimagina la iteración NS como un mapa polinómico escalar $f(\sigma)$ que actúa sobre valores singulares normalizados $\sigma \in [0, 1]$. En lugar de un único polinomio que impulse todos los $\sigma \to 1$, Pion divide los $k=5$ pasos de NS en una secuencia de dos etapas Promoción + Supresión:

1. Etapa de Promoción ($f_p$): Aplicada durante $k_p$ pasos. Este polinomio está diseñado para amplificar monótonamente los valores singulares, elevando las direcciones informativas por encima de un umbral de supresión mientras preserva su orden relativo. Los coeficientes se derivan para maximizar la pendiente en el origen ($a_p = 1.875$) mientras se ancla $\sigma=1$.
2. Etapa de Supresión ($f_s$): Aplicada durante $k_s = k - k_p$ pasos. Este polinomio ancla los valores singulares principales en 1 mientras impulsa los valores singulares pequeños y ruidosos hacia 0. Impone una derivada cero en el origen ($f'_s(0)=0$) para eliminar la amplificación lineal del ruido.

El mapa compuesto resultante crea un filtro paso alto agudo: los valores singulares dominantes se preservan (anclados cerca de 1), mientras que la cola ruidosa se suprime (contraída hacia 0). La fuerza de corte se controla mediante el hiperparámetro $k_p$ (número de pasos de promoción).

Modo por Cabeza

Para abordar la limitación respecto a la especialización de las cabezas de atención, Pion introduce un modo por cabeza. En lugar de aplicar la iteración NS a toda la matriz de pesos, remodela las proyecciones de atención a lo largo de la dimensión de la cabeza y aplica el NS paso alto independientemente a cada cabeza. Esto preserva las escalas de actualización heterogéneas requeridas entre las cabezas sin incurrir en costos computacionales adicionales más allá de un simple remodelado.

Contribuciones Clave

1. Identificación de Fallos Espectrales: El artículo es el primero en demostrar sistemáticamente que el blanqueamiento espectral uniforme de Muon es perjudicial en VLA (debido a gradientes de acción de bajo rango) y RLVR (debido a gradientes de baja SNR), lo que conduce a un rendimiento subóptimo o al colapso del modelo.
2. Optimizador Pion: La propuesta de Pion, que rediseña las iteraciones NS en un mecanismo de dos etapas de Promoción+Supresión. Logra un efecto paso alto espectral al mismo costo computacional que Muon, evitando operaciones SVD costosas.
3. Adaptación por Cabeza: La introducción de un modo de aplicación por cabeza que respeta la heterogeneidad de las cabezas de atención preentrenadas, crucial para un post-entrenamiento RLVR estable.
4. Validación Empírica Exhaustiva: Experimentos extensos que muestran que Pion supera consistentemente tanto a AdamW como a Muon en diversos escenarios, incluyendo simulación y despliegue en robots reales.

Resultados Experimentales

Entrenamiento VLA (LIBERO & LIBERO-Plus)

• Configuración: Entrenamiento de VLA-Adapter (regresión $\ell_1$) y VLANeXt (ajuste de flujo) en las suites LIBERO.
• Hallazgos:
  • Pion supera consistentemente a Muon y AdamW. En la tarea de Objetos de LIBERO, Pion alcanza una tasa de éxito del 100% después de 1.500 pasos, en comparación con el 97,0% de Muon y el 32,2% de AdamW.
  • Pion demuestra una convergencia más rápida y una mayor robustez bajo desplazamientos de distribución (por ejemplo, ruido, perturbaciones del robot) en la prueba de LIBERO-Plus.
  • Evaluación en Robot Real: En un robot Franka Research 3 utilizando una base $\pi0.5$, Pion logró una tasa de éxito promedio del 85,6% en tareas de agarrar y colocar, superando significativamente a Muon (38,9%) y AdamW (31,1%).

Post-entrenamiento RLVR (Qwen3-1.7B/4B)

• Configuración: Ajuste fino de modelos Qwen en MATH y GSM8K utilizando algoritmos GRPO y GMPO.
• Hallazgos:
  • Colapso de Muon: Muon falla completamente en configuraciones RLVR, con una precisión cayendo a casi cero debido a la amplificación del ruido.
  • Éxito de Pion: Pion recupera una señal de entrenamiento significativa, superando a AdamW en todos los tamaños de modelo y pruebas. Por ejemplo, en Qwen3-1.7B con GRPO en MATH, Pion logra un crecimiento estable de la precisión mientras Muon colapsa.
  • Ablación: Una variante "Low-pass Muon" (LPMuon), que filtra en la dirección opuesta (manteniendo valores pequeños, descartando los grandes), falla en entrenar, confirmando que la naturaleza paso alto de Pion es el factor crítico para el éxito.

Significado y Afirmaciones

El artículo afirma que Pion ofrece una solución directa y computacionalmente eficiente a los problemas de desajuste espectral inherentes a la aplicación de Muon en regímenes distintos al preentrenamiento. Al reemplazar el blanqueamiento uniforme con un filtro paso alto controlable, Pion estabiliza el entrenamiento en entornos ruidosos (RLVR) y de bajo rango (VLA) donde Muon falla.

Los autores enfatizan que las ventajas de Pion no son meramente arquitectónicas, sino que provienen de una corrección fundamental del sesgo inductivo espectral del optimizador. El trabajo sugiere que, si bien la exploración espectral uniforme es beneficiosa para el preentrenamiento, la selectividad espectral (preservar la señal, suprimir el ruido) es esencial para una adaptación y razonamiento posteriores efectivos. Los resultados en hardware de robots reales validan además que estas mejoras algorítmicas se traducen en ganancias tangibles en el rendimiento de tareas físicas.