Technical Summary: Rethinking Muon Beyond Pretraining

Problem Statement

The paper identifies fundamental limitations of Muon (MomentUm Orthogonalized by Newton–Schulz), a matrix-aware optimizer that has achieved success in Large Language Model (LLM) pretraining, when applied to two emerging training regimes: Vision-Language-Action (VLA) training and Reinforcement Learning with Verifiable Rewards (RLVR).

Muon operates by driving all singular values of the momentum matrix toward 1 via Newton–Schulz (NS) iterations, effectively performing uniform spectral whitening. While this promotes exploration in pretraining, the authors argue it induces a "spectral mismatch" in downstream tasks:

1. In VLA Training: Action-module gradients are inherently low-rank (low effective rank, or erank). Muon's uniform whitening amplifies the noisy tail directions (which dominate the low-rank spectrum) to the same magnitude as the informative leading singular values, corrupting the update.
2. In RLVR Post-Training: Policy gradients exhibit low Signal-to-Noise Ratio (SNR) due to trajectory-level rewards and stabilization mechanisms (e.g., clipping, importance sampling). Muon's uniform whitening amplifies this stochastic noise, leading to policy collapse. Additionally, Muon treats attention projections as single blocks, ignoring the per-head heterogeneity (specialized update scales) inherited from pretraining.

Existing variants like Low-Rank Muon (LRMuon) address rank issues but rely on computationally expensive SVD or sketching, failing to match Muon's scalability.

Methodology: Pion

The authors propose Pion (sPectral hIgh-pass Optimization on momeNtum), a drop-in replacement for Muon that retains its computational efficiency (per-step cost identical to Muon) but replaces uniform whitening with a spectral high-pass filter.

Core Mechanism: High-Pass NS Iteration

Pion reimagines the NS iteration as a scalar polynomial map $f(\sigma)$ acting on normalized singular values $\sigma \in [0, 1]$. Instead of a single polynomial driving all $\sigma \to 1$, Pion splits the $k=5$ NS steps into a two-stage Promotion + Suppression sequence:

1. Promotion Stage ($f_p$): Applied for $k_p$ steps. This polynomial is designed to monotonically amplify singular values, lifting informative directions above a suppression threshold while preserving their relative ordering. Coefficients are derived to maximize the slope at the origin ($a_p = 1.875$) while anchoring $\sigma=1$.
2. Suppression Stage ($f_s$): Applied for $k_s = k - k_p$ steps. This polynomial anchors leading singular values at 1 while driving small, noisy singular values toward 0. It enforces a zero derivative at the origin ($f'_s(0)=0$) to eliminate linear amplification of noise.

The resulting composite map creates a sharp high-pass filter: dominant singular values are preserved (anchored near 1), while the noisy tail is suppressed (contracted toward 0). The cutoff strength is controlled by the hyperparameter $k_p$ (number of promotion steps).

Per-Head Mode

To address the limitation regarding attention head specialization, Pion introduces a per-head mode. Instead of applying the NS iteration to the entire weight matrix, it reshapes attention projections along the head dimension and applies the high-pass NS independently to each head. This preserves the heterogeneous update scales required across heads without incurring additional computational cost beyond a simple reshape.

Key Contributions

1. Identification of Spectral Failures: The paper is the first to systematically demonstrate that Muon's uniform spectral whitening is detrimental in VLA (due to low-rank action gradients) and RLVR (due to low-SNR gradients), leading to suboptimal performance or model collapse.
2. Pion Optimizer: The proposal of Pion, which redesigns NS iterations into a two-stage Promotion+Suppression mechanism. It achieves a spectral high-pass effect at the same computational cost as Muon, avoiding expensive SVD operations.
3. Per-Head Adaptation: The introduction of a per-head application mode that respects the heterogeneity of pretrained attention heads, crucial for stable RLVR post-training.
4. Comprehensive Empirical Validation: Extensive experiments showing Pion consistently outperforms both AdamW and Muon across diverse settings, including simulation and real-robot deployment.

Experimental Results

VLA Training (LIBERO & LIBERO-Plus)

• Setup: Training VLA-Adapter ($\ell_1$-regression) and VLANeXt (flow-matching) on LIBERO suites.
• Findings:
  • Pion consistently outperforms Muon and AdamW. On the LIBERO Object task, Pion reaches 100% success rate after 1,500 steps, compared to 97.0% for Muon and 32.2% for AdamW.
  • Pion demonstrates faster convergence and higher robustness under distribution shifts (e.g., noise, robot perturbations) in the LIBERO-Plus benchmark.
  • Real-Robot Evaluation: On a Franka Research 3 robot using a $\pi0.5$ backbone, Pion achieved an average success rate of 85.6% on grasp-and-place tasks, significantly outperforming Muon (38.9%) and AdamW (31.1%).

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

• Setup: Fine-tuning Qwen models on MATH and GSM8K using GRPO and GMPO algorithms.
• Findings:
  • Muon Collapse: Muon fails completely in RLVR settings, with accuracy dropping to near zero due to noise amplification.
  • Pion Success: Pion recovers a meaningful training signal, outperforming AdamW across all model sizes and benchmarks. For instance, on Qwen3-1.7B with GRPO on MATH, Pion achieves stable accuracy growth while Muon collapses.
  • Ablation: A "Low-pass Muon" (LPMuon) variant, which filters in the opposite direction (keeping small values, discarding large ones), fails to train, confirming that the high-pass nature of Pion is the critical factor for success.

Significance and Claims

The paper claims that Pion offers a drop-in, computationally efficient solution to the spectral mismatch problems inherent in applying Muon to non-pretraining regimes. By replacing uniform whitening with a controllable high-pass filter, Pion stabilizes training in noisy (RLVR) and low-rank (VLA) environments where Muon fails.

The authors emphasize that Pion's advantages are not merely architectural but stem from a fundamental correction of the optimizer's spectral inductive bias. The work suggests that while uniform spectral exploration is beneficial for pretraining, spectral selectivity (preserving signal, suppressing noise) is essential for effective downstream adaptation and reasoning. The results on real-robot hardware further validate that these algorithmic improvements translate to tangible gains in physical task performance.