Resumo Técnico: Repensando Muon Além do Pré-treinamento

Declaração do Problema

O artigo identifica limitações fundamentais do Muon (MomentUm Orthogonalizado por Newton–Schulz), um otimizador consciente de matrizes que alcançou sucesso no pré-treinamento de Modelos de Linguagem de Grande Escala (LLM), quando aplicado a dois regimes de treinamento emergentes: treinamento Visão-Linguagem-Ação (VLA) e Aprendizado por Reforço com Recompensas Verificáveis (RLVR).

O Muon opera ao conduzir todos os valores singulares da matriz de momento para 1 via iterações Newton–Schulz (NS), realizando efetivamente um branqueamento espectral uniforme. Embora isso promova a exploração no pré-treinamento, os autores argumentam que induz um "desajuste espectral" em tarefas subsequentes:

1. No Treinamento VLA: Os gradientes do módulo de ação são inerentemente de baixo posto (baixo posto efetivo, ou erank). O branqueamento uniforme do Muon amplifica as direções de cauda ruidosas (que dominam o espectro de baixo posto) para a mesma magnitude dos valores singulares líderes informativos, corrompendo a atualização.
2. No Pós-treinamento RLVR: Os gradientes de política exibem baixa Relação Sinal-Ruído (SNR) devido a recompensas ao nível de trajetória e mecanismos de estabilização (por exemplo, clipping, amostragem por importância). O branqueamento uniforme do Muon amplifica esse ruído estocástico, levando ao colapso da política. Além disso, o Muon trata projeções de atenção como blocos únicos, ignorando a heterogeneidade por cabeça (escalas de atualização especializadas) herdada do pré-treinamento.

Variantes existentes, como o Muon de Baixo Posto (LRMuon), abordam questões de posto, mas dependem de SVD ou esboçamento computacionalmente caros, falhando em igualar a escalabilidade do Muon.

Metodologia: Pion

Os autores propõem o Pion (Otimização de Alta Passagem Espectral em Momento), um substituto direto para o Muon que mantém sua eficiência computacional (custo por passo idêntico ao do Muon), mas substitui o branqueamento uniforme por um filtro de alta passagem espectral.

Mecanismo Central: Iteração NS de Alta Passagem

O Pion reimagina a iteração NS como um mapa polinomial escalar $f(\sigma)$ atuando sobre valores singulares normalizados $\sigma \in [0, 1]$. Em vez de um único polinômio conduzindo todos os $\sigma \to 1$, o Pion divide os $k=5$ passos NS em uma sequência de duas etapas Promoção + Supressão:

1. Etapa de Promoção ($f_p$): Aplicada por $k_p$ passos. Este polinômio é projetado para amplificar monotonicamente os valores singulares, elevando direções informativas acima de um limiar de supressão enquanto preserva sua ordenação relativa. Os coeficientes são derivados para maximizar a inclinação na origem ($a_p = 1.875$) enquanto ancoram $\sigma=1$.
2. Etapa de Supressão ($f_s$): Aplicada por $k_s = k - k_p$ passos. Este polinômio ancora os valores singulares líderes em 1 enquanto conduz pequenos valores singulares ruidosos para 0. Ele impõe uma derivada zero na origem ($f'_s(0)=0$) para eliminar a amplificação linear de ruído.

O mapa composto resultante cria um filtro de alta passagem agudo: valores singulares dominantes são preservados (ancorados perto de 1), enquanto a cauda ruidosa é suprimida (contraída em direção a 0). A força do corte é controlada pelo hiperparâmetro $k_p$ (número de passos de promoção).

Modo por Cabeça

Para abordar a limitação referente à especialização das cabeças de atenção, o Pion introduz um modo por cabeça. Em vez de aplicar a iteração NS a toda a matriz de pesos, ele remodela as projeções de atenção ao longo da dimensão da cabeça e aplica a NS de alta passagem independentemente a cada cabeça. Isso preserva as escalas de atualização heterogêneas necessárias entre as cabeças sem incorrer em custo computacional adicional além de uma simples remodelagem.

Contribuições Chave

1. Identificação de Falhas Espectrais: O artigo é o primeiro a demonstrar sistematicamente que o branqueamento espectral uniforme do Muon é prejudicial no VLA (devido a gradientes de ação de baixo posto) e no RLVR (devido a gradientes de baixa SNR), levando a desempenho subótimo ou colapso do modelo.
2. Otimizador Pion: A proposta do Pion, que redesenha as iterações NS em um mecanismo de duas etapas Promoção+Supressão. Ele alcança um efeito de alta passagem espectral ao mesmo custo computacional do Muon, evitando operações de SVD caras.
3. Adaptação por Cabeça: A introdução de um modo de aplicação por cabeça que respeita a heterogeneidade das cabeças de atenção pré-treinadas, crucial para o pós-treinamento estável em RLVR.
4. Validação Empírica Abrangente: Experimentos extensivos mostrando que o Pion consistentemente supera tanto o AdamW quanto o Muon em configurações diversas, incluindo simulação e implantação em robôs reais.

Resultados Experimentais

Treinamento VLA (LIBERO & LIBERO-Plus)

• Configuração: Treinamento de VLA-Adapter (regressão $\ell_1$) e VLANeXt (correspondência de fluxo) nas suítes LIBERO.
• Descobertas:
  • O Pion supera consistentemente o Muon e o AdamW. Na tarefa de Objeto do LIBERO, o Pion atinge 100% de taxa de sucesso após 1.500 passos, comparado a 97,0% para o Muon e 32,2% para o AdamW.
  • O Pion demonstra convergência mais rápida e maior robustez sob deslocamentos de distribuição (por exemplo, ruído, perturbações do robô) no benchmark LIBERO-Plus.
  • Avaliação em Robô Real: Em um robô Franka Research 3 usando uma base $\pi0.5$, o Pion alcançou uma taxa de sucesso média de 85,6% em tarefas de agarrar e colocar, superando significativamente o Muon (38,9%) e o AdamW (31,1%).

Pós-treinamento RLVR (Qwen3-1.7B/4B)

• Configuração: Ajuste fino de modelos Qwen em MATH e GSM8K usando algoritmos GRPO e GMPO.
• Descobertas:
  • Colapso do Muon: O Muon falha completamente em configurações RLVR, com a precisão caindo para perto de zero devido à amplificação de ruído.
  • Sucesso do Pion: O Pion recupera um sinal de treinamento significativo, superando o AdamW em todos os tamanhos de modelo e benchmarks. Por exemplo, no Qwen3-1.7B com GRPO em MATH, o Pion alcança crescimento estável de precisão enquanto o Muon colapsa.
  • Ablação: Uma variante "Muon de Baixa Passagem" (LPMuon), que filtra na direção oposta (mantendo pequenos valores, descartando grandes), falha em treinar, confirmando que a natureza de alta passagem do Pion é o fator crítico para o sucesso.

Significado e Alegações

O artigo afirma que o Pion oferece uma solução direta e computacionalmente eficiente para os problemas de desajuste espectral inerentes à aplicação do Muon em regimes não de pré-treinamento. Ao substituir o branqueamento uniforme por um filtro de alta passagem controlável, o Pion estabiliza o treinamento em ambientes ruidosos (RLVR) e de baixo posto (VLA) onde o Muon falha.

Os autores enfatizam que as vantagens do Pion não são meramente arquitetônicas, mas derivam de uma correção fundamental do viés indutivo espectral do otimizador. O trabalho sugere que, enquanto a exploração espectral uniforme é benéfica para o pré-treinamento, a seletividade espectral (preservar sinal, suprimir ruído) é essencial para adaptação e raciocínio eficazes a jusante. Os resultados em hardware de robôs reais validam ainda mais que essas melhorias algorítmicas se traduzem em ganhos tangíveis no desempenho de tarefas físicas.