Repenser Muon au-delà du pré-entraînement : défaillances spectrales et remèdes passe-haut pour VLA et RLVR

Résumé

Muon est un optimiseur prenant en compte la structure matricielle qui exploite les itérations de Newton-Schulz (NS) pour imposer une orthogonalisation spectrale du gradient en ramenant toutes les valeurs singulières de la matrice de momentum vers 1. Bien que ce blanchiment spectral uniforme améliore l'exploration et surpasse AdamW dans le pré-entraînement des LLM, nous montrons qu'il peut conduire à des limitations fondamentales au-delà du pré-entraînement dans deux régimes : (i) l'entraînement cross-modal vision-langage-action (VLA), où les gradients intrinsèquement de faible rang du module d'action provoquent une amplification des directions de queue bruitées, et (ii) l'apprentissage par renforcement avec récompenses vérifiables (RLVR), où les gradients à faible SNR et la nécessité de préserver la spécialisation par tête issue de l'entraînement préalable rendent le blanchiment instable. Pour relever ces défis, nous proposons Pion, un remplacement direct de Muon qui préserve son efficacité computationnelle tout en substituant le blanchiment spectral uniforme par un mécanisme en deux étapes de Promotion+Suppression, que nous appelons itération NS passe-haut. Cette conception induit un effet spectral passe-haut marqué, ancrant les valeurs singulières dominantes à 1 tout en supprimant les composantes de queue bruitées vers 0, avec une force de filtrage contrôlable. Pour préserver l'hétérogénéité par tête issue du pré-entraînement, Pion supporte également un mode par tête qui applique les mises à jour indépendamment sur chaque tête d'attention via un simple redimensionnement, sans surcoût. Dans l'entraînement VLA sur LIBERO et LIBERO-Plus, Pion surpasse systématiquement les deux bases de comparaison sur les architectures de régression l₁ (VLA-Adapter) et de flow matching (VLANeXt), par exemple en atteignant un taux de succès de 100 % sur LIBERO Object après 1 500 étapes d'entraînement avec VLA-Adapter, contre 97,0 % pour Muon et seulement 32,2 % pour AdamW. L'avantage de Pion s'étend également à un vrai robot Franka Research 3 avec un backbone pi₀,₅ dans la configuration DROID sur trois tâches de saisie et placement. Dans le post-entraînement RLVR sur Qwen3-1.7B/4B avec GRPO et GMPO, Pion surpasse également AdamW sur MATH et GSM8K tandis que Muon s'effondre à zéro.

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.