Переосмысление Muon за пределами предварительного обучения: спектральные сбои и высокочастотные коррективы для VLA и RLVR

Аннотация

Мюон — это матрично-осознанный оптимизатор, использующий итерации Ньютона–Шульца (NS) для принудительной ортогонализации спектрального градиента путем подтягивания всех сингулярных значений матрицы импульса к 1. Хотя такое равномерное спектральное отбеливание улучшает исследование и превосходит AdamW при предобучении больших языковых моделей, мы показываем, что оно может приводить к фундаментальным ограничениям за пределами предобучения в двух режимах: (i) кросс-модальное обучение «зрение–язык–действие» (VLA), где градиенты модуля действий, имеющие по своей природе низкий ранг, вызывают усиление шумных хвостовых направлений, и (ii) обучение с подкреплением на основе проверяемых наград (RLVR), где низкое отношение сигнал/шум градиентов и необходимость сохранения специализации по головам из предыдущего обучения делают отбеливание нестабильным. Для решения этих проблем мы предлагаем Pion — замену Muon, работающую без изменения кода, которая сохраняет его вычислительную эффективность, заменяя равномерное спектральное отбеливание двухэтапным механизмом Promotion+Suppression, называемым нами высокочастотной итерацией NS. Эта конструкция создает резкий спектральный высокочастотный эффект, фиксируя доминирующие сингулярные значения на 1, подавляя шумные хвостовые компоненты до 0 с контролируемой силой фильтра. Для сохранения предобученной гетерогенности по головам Pion также поддерживает режим работы по головам, который применяет обновления независимо для каждой головы внимания посредством простого изменения формы, без дополнительных затрат. В обучении VLA на наборах LIBERO и LIBERO-Plus Pion стабильно превосходит оба базовых метода как в архитектурах с l₁-регрессией (VLA-Adapter), так и с согласованием потоков (VLANeXt), например, достигая 100% успешности на LIBERO Object после 1500 шагов обучения с VLA-Adapter против 97,0% у Muon и всего 32,2% у AdamW. Преимущество Pion распространяется и на реального робота Franka Research 3 с основой pi₀.₅ в конфигурации DROID для трех задач «схватить и разместить». При дообучении с подкреплением на Qwen3-1,7B/4B с использованием GRPO и GMPO Pion также превосходит AdamW по MATH и GSM8K, в то время как Muon коллапсирует до нуля.

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.