Repensando Muon más allá del Preentrenamiento: Fallos Espectrales y Remedios de Paso Alto para VLA y RLVR

Resumen

Muon es un optimizador consciente de la matriz que aprovecha las iteraciones de Newton-Schulz (NS) para imponer la ortogonalización espectral del gradiente llevando todos los valores singulares de la matriz de momento hacia 1. Si bien este blanqueo espectral uniforme mejora la exploración y supera a AdamW en el preentrenamiento de LLMs, mostramos que puede conducir a limitaciones fundamentales más allá del preentrenamiento en dos regímenes: (i) entrenamiento multimodal visión-lenguaje-acción (VLA), donde los gradientes inherentemente de bajo rango del módulo de acción causan amplificación de direcciones de cola ruidosas, y (ii) aprendizaje por refuerzo con recompensas verificables (RLVR), donde los gradientes de baja SNR y la necesidad de preservar la especialización por cabeza del entrenamiento previo hacen que el blanqueo sea inestable. Para abordar estos desafíos, proponemos Pion, un reemplazo directo para Muon que preserva su eficiencia computacional mientras sustituye el blanqueo espectral uniforme por un mecanismo de dos etapas de Promoción+Supresión, al que denominamos iteración NS de paso alto. Este diseño induce un fuerte efecto espectral de paso alto, anclando los valores singulares dominantes en 1 mientras suprime los componentes de cola ruidosos hacia 0, con una fuerza de filtro controlable. Para preservar la heterogeneidad preentrenada por cabeza, Pion también admite un modo por cabeza que aplica actualizaciones de forma independiente entre las cabezas de atención mediante una simple reorganización, sin costo adicional. En el entrenamiento VLA en LIBERO y LIBERO-Plus, Pion supera consistentemente a ambas líneas base en las arquitecturas de regresión \(l_1\) (VLA-Adapter) y de emparejamiento de flujo (VLANeXt), por ejemplo, alcanzando una tasa de éxito del 100% en LIBERO Object después de 1500 pasos de entrenamiento con VLA-Adapter, frente al 97,0% para Muon y solo el 32,2% para AdamW. La ventaja de Pion se extiende además a un robot real Franka Research 3 con un backbone \(\pi_{0.5}\) bajo la configuración DROID en tres tareas de agarre y colocación. En el posentrenamiento RLVR en Qwen3-1.7B/4B con GRPO y GMPO, Pion también supera a AdamW en MATH y GSM8K mientras que Muon colapsa a cero.

English

Muon is a matrix-aware optimizer that leverages Newton-Schulz (NS) iterations to enforce spectral gradient orthogonalization by driving all singular values of the momentum matrix toward 1. While this uniform spectral whitening enhances exploration and outperforms AdamW in LLM pretraining, we show it could lead to fundamental limitations beyond pretraining in two regimes: (i) cross-modality vision-language-action (VLA) training, where inherently low-rank action-module gradients cause amplification of noisy tail directions, and (ii) reinforcement learning with verifiable rewards (RLVR), where low-SNR gradients and the need to preserve per-head specialization from prior training make whitening unstable. To address these challenges, we propose Pion, a drop-in replacement for Muon that preserves its computational efficiency while replacing uniform spectral whitening with a two-stage Promotion+Suppression mechanism, which we call the high-pass NS iteration. This design induces a sharp spectral high-pass effect, anchoring dominant singular values at 1 while suppressing noisy tail components toward 0, with controllable filter strength. To preserve pretrained per-head heterogeneity, Pion also supports a per-head mode that applies updates independently across attention heads via a simple reshape, at no extra cost. In VLA training on LIBERO and LIBERO-Plus, Pion consistently outperforms both baselines across l_1-regression (VLA-Adapter) and flow-matching (VLANeXt) architectures, e.g., reaching 100% success rate on LIBERO Object after 1,500 training steps with VLA-Adapter, vs. 97.0% for Muon and only 32.2% for AdamW. The advantage of Pion further extends to a real Franka Research 3 robot with a pi_0.5 backbone under the DROID setup on three grasp-and-place tasks. In RLVR post-training on Qwen3-1.7B/4B with GRPO and GMPO, Pion also outperforms AdamW on MATH and GSM8K while Muon collapses to zero.