Repensando o Muon Além do Pré-treinamento: Falhas Espectrais e Remédios Passa-Alta para VLA e RLVR

Resumo

Muon é um otimizador ciente de matrizes que utiliza iterações de Newton-Schulz (NS) para impor ortogonalização espectral do gradiente, conduzindo todos os valores singulares da matriz de momento para 1. Embora esse branqueamento espectral uniforme melhore a exploração e supere o AdamW no pré-treinamento de LLMs, mostramos que ele pode levar a limitações fundamentais além do pré-treinamento em dois regimes: (i) treinamento visão-linguagem-ação (VLA) cross-modular, onde gradientes intrinsecamente de baixo posto do módulo de ação causam amplificação de direções de cauda ruidosas, e (ii) aprendizado por reforço com recompensas verificáveis (RLVR), onde gradientes de baixa SNR e a necessidade de preservar a especialização por cabeça do treinamento anterior tornam o branqueamento instável. Para enfrentar esses desafios, propomos o Pion, um substituto direto para o Muon que preserva sua eficiência computacional enquanto substitui o branqueamento espectral uniforme por um mecanismo de Promoção+Supressão em dois estágios, que chamamos de iteração NS passa-alta. Esse design induz um efeito de filtro passa-alta espectral acentuado, ancorando os valores singulares dominantes em 1 enquanto suprime componentes de cauda ruidosos em direção a 0, com intensidade de filtro controlável. Para preservar a heterogeneidade pré-treinada por cabeça, o Pion também suporta um modo por cabeça que aplica atualizações de forma independente entre cabeças de atenção por meio de um simples redimensionamento, sem custo adicional. No treinamento VLA em LIBERO e LIBERO-Plus, o Pion supera consistentemente ambas as bases de comparação em arquiteturas de regressão l₁ (VLA-Adapter) e de correspondência de fluxo (VLANeXt), por exemplo, alcançando 100% de taxa de sucesso no LIBERO Object após 1.500 passos de treinamento com VLA-Adapter, contra 97,0% para o Muon e apenas 32,2% para o AdamW. A vantagem do Pion se estende ainda a um robô real Franka Research 3 com uma espinha dorsal pi₀.₅ sob a configuração DROID em três tarefas de agarrar e colocar. No pós-treinamento RLVR no Qwen3-1.7B/4B com GRPO e GMPO, o Pion também supera o AdamW em MATH e GSM8K, enquanto o Muon colapsa para zero.

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.