Riepilogo Tecnico: Ripensare Muon Oltre il Pre-addestramento

Enunciato del Problema

Il documento identifica limitazioni fondamentali di Muon (MomentUm Orthogonalized by Newton–Schulz), un ottimizzatore consapevole delle matrici che ha ottenuto successo nel pre-addestramento di Modelli Linguistici su Grande Scala (LLM), quando applicato a due regimi di addestramento emergenti: l'addestramento Vision-Language-Action (VLA) e l'Apprendimento per Rinforzo con Ricompense Verificabili (RLVR).

Muon opera guidando tutti i valori singolari della matrice di momento verso 1 tramite iterazioni Newton–Schulz (NS), eseguendo efficacemente un whitening spettrale uniforme. Sebbene ciò favorisca l'esplorazione nel pre-addestramento, gli autori sostengono che ciò induca un "mismatch spettrale" nelle attività a valle:

1. Nell'addestramento VLA: I gradienti del modulo di azione sono intrinsecamente a basso rango (basso rango effettivo, o erank). Il whitening uniforme di Muon amplifica le direzioni rumorose della coda (che dominano lo spettro a basso rango) alla stessa magnitudine dei valori singolari informativi principali, corrompendo l'aggiornamento.
2. Nel post-addestramento RLVR: I gradienti della politica presentano un basso rapporto Segnale-Rumore (SNR) a causa delle ricompense a livello di traiettoria e dei meccanismi di stabilizzazione (es. clipping, campionamento per importanza). Il whitening uniforme di Muon amplifica questo rumore stocastico, portando al collasso della politica. Inoltre, Muon tratta le proiezioni di attenzione come singoli blocchi, ignorando l'eterogeneità per testa (scale di aggiornamento specializzate) ereditata dal pre-addestramento.

Varianti esistenti come Low-Rank Muon (LRMuon) affrontano i problemi di rango ma si basano su SVD o schizzi computazionalmente costosi, non riuscendo a eguagliare la scalabilità di Muon.

Metodologia: Pion

Gli autori propongono Pion (sPectral hIgh-pass Optimization on momeNtum), un sostituto plug-in per Muon che mantiene la sua efficienza computazionale (costo per passo identico a Muon) ma sostituisce il whitening uniforme con un filtro passa-alto spettrale.

Meccanismo Principale: Iterazione NS Passa-Alto

Pion reimmagina l'iterazione NS come una mappa polinomiale scalare $f(\sigma)$ che agisce sui valori singolari normalizzati $\sigma \in [0, 1]$. Invece di un singolo polinomio che guida tutti i $\sigma \to 1$, Pion suddivide i $k=5$ passi NS in una sequenza a due stadi di Promozione + Soppressione:

1. Stadio di Promozione ($f_p$): Applicato per $k_p$ passi. Questo polinomio è progettato per amplificare monotonicamente i valori singolari, sollevando le direzioni informative al di sopra di una soglia di soppressione mentre ne preserva l'ordinamento relativo. I coefficienti sono derivati per massimare la pendenza all'origine ($a_p = 1.875$) ancorando $\sigma=1$.
2. Stadio di Soppressione ($f_s$): Applicato per $k_s = k - k_p$ passi. Questo polinomio ancora i valori singolari principali a 1 mentre guida i piccoli valori singolari rumorosi verso 0. Impone una derivata nulla all'origine ($f'_s(0)=0$) per eliminare l'amplificazione lineare del rumore.

La mappa composita risultante crea un filtro passa-alto netto: i valori singolari dominanti sono preservati (ancorati vicino a 1), mentre la coda rumorosa è soppressa (contratta verso 0). La forza del taglio è controllata dall'iperparametro $k_p$ (numero di passi di promozione).

Modalità Per-Testa

Per affrontare la limitazione relativa alla specializzazione delle teste di attenzione, Pion introduce una modalità per testa. Invece di applicare l'iterazione NS all'intera matrice dei pesi, ridisegna le proiezioni di attenzione lungo la dimensione della testa e applica l'NS passa-alto indipendentemente a ciascuna testa. Ciò preserva le scale di aggiornamento eterogenee richieste tra le teste senza incorrere in costi computazionali aggiuntivi oltre a un semplice ridimensionamento.

Contributi Chiave

1. Identificazione dei Fallimenti Spettrali: Il documento è il primo a dimostrare sistematicamente che il whitening spettrale uniforme di Muon è dannoso nel VLA (a causa di gradienti di azione a basso rango) e nel RLVR (a causa di gradienti a basso SNR), portando a prestazioni subottimali o al collasso del modello.
2. Ottimizzatore Pion: La proposta di Pion, che ridisegna le iterazioni NS in un meccanismo a due stadi Promozione+Soppressione. Raggiunge un effetto passa-alto spettrale allo stesso costo computazionale di Muon, evitando operazioni SVD costose.
3. Adattamento Per-Testa: L'introduzione di una modalità di applicazione per testa che rispetta l'eterogeneità delle teste di attenzione pre-addestrate, cruciale per un post-addestramento RLVR stabile.
4. Validazione Empirica Completa: Estesi esperimenti che mostrano come Pion superi costantemente sia AdamW che Muon in contesti diversificati, inclusi simulazioni e dispiegamento su robot reali.

Risultati Sperimentali

Addestramento VLA (LIBERO & LIBERO-Plus)

• Configurazione: Addestramento di VLA-Adapter (regressione $\ell_1$) e VLANeXt (flow-matching) sui suite LIBERO.
• Risultati:
  • Pion supera costantemente Muon e AdamW. Nel compito LIBERO Object, Pion raggiunge un tasso di successo del 100% dopo 1.500 passi, rispetto al 97,0% di Muon e al 32,2% di AdamW.
  • Pion dimostra una convergenza più rapida e una maggiore robustezza sotto spostamenti di distribuzione (es. rumore, perturbazioni del robot) nel benchmark LIBERO-Plus.
  • Valutazione su Robot Reale: Su un robot Franka Research 3 utilizzando un backbone $\pi0.5$, Pion ha raggiunto un tasso di successo medio del 85,6% nei compiti di afferrare e posizionare, superando significativamente Muon (38,9%) e AdamW (31,1%).

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

• Configurazione: Fine-tuning dei modelli Qwen su MATH e GSM8K utilizzando gli algoritmi GRPO e GMPO.
• Risultati:
  • Collasso di Muon: Muon fallisce completamente negli ambienti RLVR, con l'accuratezza che scende a quasi zero a causa dell'amplificazione del rumore.
  • Successo di Pion: Pion recupera un segnale di addestramento significativo, superando AdamW in tutte le dimensioni del modello e nei benchmark. Ad esempio, su Qwen3-1.7B con GRPO su MATH, Pion ottiene una crescita stabile dell'accuratezza mentre Muon collassa.
  • Ablazione: Una variante "Low-pass Muon" (LPMuon), che filtra nella direzione opposta (mantenendo i piccoli valori, scartando quelli grandi), fallisce nell'addestramento, confermando che la natura passa-alto di Pion è il fattore critico per il successo.

Significato e Affermazioni

Il documento afferma che Pion offre una soluzione plug-in ed efficientemente computazionale ai problemi di mismatch spettrale intrinseci nell'applicazione di Muon a regimi non di pre-addestramento. Sostituendo il whitening uniforme con un filtro passa-alto controllabile, Pion stabilizza l'addestramento in ambienti rumorosi (RLVR) e a basso rango (VLA) dove Muon fallisce.

Gli autori sottolineano che i vantaggi di Pion non sono meramente architetturali, ma derivano da una correzione fondamentale del bias induttivo spettrale dell'ottimizzatore. Il lavoro suggerisce che, sebbene l'esplorazione spettrale uniforme sia benefica per il pre-addestramento, la selettività spettrale (preservare il segnale, sopprimere il rumore) è essenziale per un'adattazione e un ragionamento efficaci a valle. I risultati sull'hardware di robot reali validano ulteriormente che questi miglioramenti algoritmici si traducono in guadagni tangibili nelle prestazioni dei compiti fisici.