Neubetrachtung von Muon jenseits des Vortrainings: Spektrale Fehler und Hochpass-Lösungen für VLA und RLVR

Zusammenfassung

Muon ist ein matrixbewusster Optimierer, der Newton-Schulz (NS)-Iterationen nutzt, um eine spektrale Gradientenorthogonalisierung zu erzwingen, indem alle Singulärwerte der Momentum-Matrix in Richtung 1 getrieben werden. Während dieses gleichmäßige spektrale Whitening die Exploration verbessert und AdamW beim LLM-Pretraining übertrifft, zeigen wir, dass es in zwei Bereichen über das Pretraining hinaus zu grundlegenden Einschränkungen führen kann: (i) modusübergreifendes Vision-Language-Action (VLA)-Training, bei dem inhärent niedrigrangige Gradienten des Aktionsmoduls eine Verstärkung verrauschter Ausläuferrichtungen verursachen, und (ii) Verstärkungslernen mit verifizierbaren Belohnungen (RLVR), bei dem Gradienten mit niedrigem SNR und die Notwendigkeit, die kopfweise Spezialisierung aus vorherigem Training zu bewahren, das Whitening instabil machen. Um diese Herausforderungen zu bewältigen, schlagen wir Pion vor, einen Drop-in-Ersatz für Muon, der dessen Recheneffizienz beibehält, während das gleichmäßige spektrale Whitening durch einen zweistufigen Promotion+Suppression-Mechanismus ersetzt wird, den wir Hochpass-NS-Iteration nennen. Dieses Design erzeugt einen scharfen spektralen Hochpasseffekt, der dominante Singulärwerte bei 1 verankert, während verrauschte Ausläuferkomponenten in Richtung 0 unterdrückt werden, mit kontrollierbarer Filterstärke. Um die vortrainierte kopfweise Heterogenität zu bewahren, unterstützt Pion auch einen Kopfmodus, der Aktualisierungen unabhängig über Aufmerksamkeitsköpfe hinweg durch eine einfache Umformung anwendet, ohne zusätzliche Kosten. Im VLA-Training auf LIBERO und LIBERO-Plus übertrifft Pion durchgängig beide Basislinien sowohl bei l_1-Regressions- (VLA-Adapter) als auch bei Flow-Matching-Architekturen (VLANeXt), z. B. erreicht es nach 1500 Trainingsschritten mit VLA-Adapter eine Erfolgsrate von 100 % auf LIBERO Object, gegenüber 97,0 % für Muon und nur 32,2 % für AdamW. Der Vorteil von Pion erstreckt sich weiterhin auf einen echten Franka Research 3 Roboter mit einem pi_0.5-Backbone unter dem DROID-Aufbau bei drei Greif- und Platzierungsaufgaben. Im RLVR-Posttraining auf Qwen3-1.7B/4B mit GRPO und GMPO übertrifft Pion ebenfalls AdamW auf MATH und GSM8K, während Muon auf Null zusammenbricht.

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.