Heroverweging van Muon voorbij pre-training: Spectrale fouten en hoogdoorlaatoplossingen voor VLA en RLVR

Samenvatting

Muon is een matrix-bewuste optimizer die Newton-Schulz (NS)-iteraties gebruikt om spectrale gradiëntorthogonalisatie af te dwingen door alle singuliere waarden van de momentummatrix naar 1 te sturen. Hoewel deze uniforme spectrale witmaking de exploratie verbetert en beter presteert dan AdamW bij LLM-pretraining, tonen we aan dat dit tot fundamentele beperkingen kan leiden buiten pretraining in twee regimes: (i) cross-modale visie-taal-actie (VLA)-training, waar inherent laag-rank actiemodulegradiënten versterking van ruisende staartrichtingen veroorzaken, en (ii) reinforcement learning met verifieerbare beloningen (RLVR), waar laag-SNR-gradiënten en de noodzaak om per-hoofdspecialisatie uit eerdere training te behouden, witmaking instabiel maken. Om deze uitdagingen aan te pakken, stellen we Pion voor, een directe vervanging voor Muon die de computationele efficiëntie behoudt terwijl uniforme spectrale witmaking wordt vervangen door een tweetraps Promotie+Suppressiemechanisme, dat we de hoogdoorlaat-NS-iteratie noemen. Dit ontwerp induceert een scherp hoogdoorlaat-spectraaleffect, waarbij dominante singuliere waarden op 1 worden verankerd terwijl ruisende staartcomponenten naar 0 worden onderdrukt, met beheersbare filtersterkte. Om de voorgetrainde per-hoofdheterogeniteit te behouden, ondersteunt Pion ook een per-hoofdmodus die updates onafhankelijk over aandachtshoofden toepast via een eenvoudige hervorming, zonder extra kosten. In VLA-training op LIBERO en LIBERO-Plus presteert Pion consequent beter dan beide basislijnen over l₁-regressie (VLA-Adapter) en flow-matching (VLANeXt)-architecturen, bijvoorbeeld door een succespercentage van 100% te bereiken op LIBERO Object na 1.500 trainingsstappen met VLA-Adapter, versus 97,0% voor Muon en slechts 32,2% voor AdamW. Het voordeel van Pion strekt zich verder uit tot een echte Franka Research 3-robot met een pi_0.5-backbone onder de DROID-opstelling op drie grijp-en-plaatstaken. In RLVR-natraining op Qwen3-1.7B/4B met GRPO en GMPO presteert Pion ook beter dan AdamW op MATH en GSM8K, terwijl Muon naar nul instort.

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.