🤖 machine learning

Approximate Muon with low-rank adapters

Questo articolo introduce sMuon, un metodo efficiente che approssima l'ottimizzatore Muon per il fine-tuning parametrico-efficiente a basso rango linearizzando l'obiettivo e risolvendolo tramite minimi quadrati, abilitando così i benefici prestazionali di Muon in contesti come SFT e pretraining ReLoRA dove precedentemente era incompatibile con LoRA.

Autori originali: Ben Anson, Conor Houghton, Edward Milsom

Pubblicato 2026-08-17

📖 4 min di lettura☕ Lettura da pausa caffè

CC0 1.0

Autori originali: Ben Anson, Conor Houghton, Edward Milsom

Articolo originale dedicato al pubblico dominio sotto CC0 1.0 (http://creativecommons.org/publicdomain/zero/1.0/). ✨ Questa è una spiegazione generata dall'IA dell'articolo qui sotto. Non è stata scritta né approvata dagli autori. Per precisione tecnica, consulta l'articolo originale. Leggi il disclaimer completo

Immagina di cercare di insegnare a un robot gigante e super intelligente una nuova abilità, come scrivere poesie o risolvere problemi di matematica. Non vuoi riaddestrare l'intero cervello del robot da zero perché richiede un tempo infinito e costa una fortuna. Invece, vuoi attaccare un piccolo "modulo di addestramento" leggero al suo cervello esistente. Questo è chiamato Parameter-Efficient Fine-Tuning (PEFT). Pensa a questo come al mettere una cuffia specializzata a un robot versatile: la cuffia gli insegna il nuovo trucco senza cambiare l'hardware principale del robot.

Un modo popolare per costruire queste cuffie è chiamato LoRA (Low-Rank Adaptation). È come costruire la cuffia usando due sottili fogli di carta piatti che, quando sovrapposti, creano la forma della nuova abilità. Per insegnare al robot, abbiamo bisogno di un ottimizzatore — un coach matematico che dice alla cuffia come regolare la sua forma per migliorare. Il coach più comune è chiamato AdamW, ma esiste un coach più nuovo e potente chiamato Muon. Muon è famoso per essere incredibilmente efficiente durante il pretraining (insegnare al robot le sue primissime lezioni da zero). Funziona assicurando che gli aggiornamenti del cervello del robot avvengano in un modo perfettamente bilanciato, "ortogonale", come un ballerino che ruota perfettamente su un piede senza barcollare.

Tuttavia, c'è un problema. Quando provi a usare Muon per addestrare la piccola cuffia LoRA, la matematica si rompe. Muon vuole bilanciare l'intero strato del cervello del robot, ma la cuffia LoRA è solo una piccola fetta a basso rango di quello strato. È come cercare di usare la bacchetta di un direttore d'orchestra per dirigere un singolo violinista in un modo che abbia senso per l'intera orchestra; la geometria semplicemente non si adatta. Il workaround standard è stato quello di ignorare l'orchestra e dire semplicemente al violinista di ruotare da solo, ma questo perde la magia di Muon. Questo articolo si chiede: possiamo trovare un modo per far funzionare Muon per queste piccole cuffie senza rompere la matematica o rallentare il processo?

Gli autori, Ben Anson, Conor Houghton ed Edward Milsom, dicono di sì. Introducono un nuovo metodo chiamato sMuon (small Muon). Invece di forzare l'impossibile, utilizzano un astuto trucco matematico: "linearizzano" il problema. Immagina di cercare di camminare su una fune tenendo un ombrello gigante e traballante. Invece di cercare di bilanciare tutto l'ombrello in una volta sola, fingi che l'ombrello sia solo un bastone dritto per un breve istante, calcola il passo perfetto e poi ti adegui. Facendo questo, trovano il modo migliore per aggiornare la cuffia LoRA che approssima la danza perfetta di Muon.

L'articolo trova che questo nuovo metodo sMuon funzioni sorprendentemente bene. Quando lo hanno testato su vari modelli linguistici, sMuon ha spesso performato meglio del coach standard AdamW e di altri tentativi con Muon. È stato particolarmente impressionante quando usato su un robot che era già stato pre-addestrato con Muon (chiamato Moonlight), dove sMuono ha ottenuto i punteggi più alti in sei di undici diversi compiti. In un esperimento di pre-training, sMuon si è pareggiato per i migliori risultati nel ridurre il tasso di errore, ma lo ha fatto circa il 30% più velocemente di un suo rivale più stretto, un metodo chiamato Riemannian.

Fondamentalmente, gli autori dimostrano che il loro metodo non è solo un'idea teorica; è anche pratico. Mentre altri metodi avanzati richiedono calcoli pesanti e lenti (come smontare e rimontare matrici complesse), sMuon si basa quasi interamente su semplici moltiplicazioni di matrici — il tipo di matematica per cui i moderni chip informatici sono costruiti per lavorare a velocità fulminea. Ciò significa che sMuon è abbastanza veloce da poter essere usato nell'addestramento reale senza rallentare il processo. L'articolo suggerisce che, rispettando la geometria dell'intero strato invece di trattare le parti LoRA come pezzi separati e non correlati, possiamo ottenere un significativo aumento delle prestazioni con pochissimo costo aggiuntivo. È un promemoria del fatto che, a volte, il modo migliore per insegnare un nuovo trucco a un robot gigante è dargli una piccola, intelligente ed elegante spinta matematica.

=== RIASSUNTO TECNICO ===

Riassunto Tecnico: Approssimazione di Muon con Low-Rank Adapters (sMuon)

Definizione del Problema

L'ottimizzatore Muon ha dimostrato prestazioni superiori nel pretraining di grandi modelli linguistici (LLM) ortogonalizzando gli aggiornamenti dei pesi ad ogni livello, equiparando efficacemente i valori singolari. Questo comportamento corrisponde alla discesa più ripida sotto la norma spettrale, controllando il cambiamento peggiore di un livello come un operatore lineare piuttosto che come un vettore appiattito. Tuttavia, Muon è raramente utilizzato per il Parameter-Efficient Fine-Tuning (PEFT), in particolare con il Low-Rank Adaptation (LoRA).

L'ostruzione matematica fondamentale è che Muon richiede l'ortogonalizzazione del passo di discesa per l'intera matrice dei pesi del livello. Nel contesto LoRA, dove l'aggiornamento del peso è vincolato a una parametrizzazione a basso rango $\Delta W = BA$ (con $B \in \mathbb{R}^{d_1 \times r}$ e $A \in \mathbb{R}^{r \times d_2}$ ), è matematicamente impossibile trovare generalmente aggiornamenti $\delta B$ e $\delta A$ tali che il risultante $\delta \Delta W$ corrisponda perfettamente al gradiente ortogonalizzato a pieno rango richiesto da Muon.

Le soluzioni di contorno esistenti includono:

Per-factor Muon: Ortogonalizzazione separata di $B$ e $A$ , che ignora la geometria del livello completo.
LoRA-Muon: Rilassamento dell'obiettivo utilizzando la disuguaglianza triangolare.
Riemannion: Formulazione del problema come ottimizzazione Riemanniana su un manifold a basso rango, che richiede costose decomposizioni in Valori Singolari (SVD) troncate durante il passaggio di ritrazione, portando a un significativo overhead computazionale.

Metodologia

Gli autori propongono sMuon (small Muon), un metodo che approssima la soluzione di un obiettivo Muon rilassato nel contesto a basso rango. L'approccio si basa sulla linearizzazione e sull'ottimizzazione dei minimi quadrati per trovare il passo di discesa ottimale.

1. Linearizzazione e Formulazione dei Minimi Quadrati

L'obiettivo è trovare gli aggiornamenti $\delta B, \delta A$ tali che l'aggiornamento implicito $\delta \Delta W \approx \delta B A + B \delta A$ approssimi l'ideale passo Muon $\delta \Delta W^* = -\eta \cdot \text{msign}(G_{\Delta W})$ , dove $\text{msign}(X) = UV^\top$ è il fattore polare della SVD.

Poiché l'ideale passo è generalmente a pieno rango e non può essere realizzato esattamente dai fattori a basso rango, gli autori linearizzano il problema scartando i termini del secondo ordine ( $\delta B \delta A$ ) e risolvendo un problema di minimi quadrati:
$(\delta B^*, \delta A^*) = \arg \min_{\delta B, \delta A} \| \delta \Delta W^* - (\delta B A + B \delta A) \|_F^2$
Ciò produce una soluzione in forma chiusa che coinvolge le pseudoinverse di Moore-Penrose ( $B^\dagger, A^\dagger$ ). Fondamentalmente, gli autori dimostrano che l'aggiornamento ottimale non è unico in termini di fattori $(\delta B, \delta A)$ , ma il risultante cambiamento del primo ordine al livello completo $\delta \Delta W$ è unico.

2. Computazione Efficiente Basata Solo su Matmul

Un'implementazione diretta della soluzione in forma chiusa richiederebbe l'ortogonalizzazione di una grande matrice $d_1 \times d_2$ , il che è computazionalmente proibitivo. Per affrontare questo problema, gli autori derivano una formulazione efficiente (Algoritmo 2) che evita l'esplicita ortogonalizzazione di grandi matrici:

Gradiente a Basso Rango: Inveve di utilizzare il gradiente completo $G_{\Delta W}$ , utilizzano una proiezione a basso rango $H = P_B G_{\Delta W} + (I-P_B)G_{\Delta W} P_A$ , che è equivalente al primo ordine per l'obiettivo di ottimizzazione.
Riduzione della Dimensionalità: Sfruttano l'identità $\text{msign}(H) = U \text{msign}(U^\top H V) V^\top$ , dove $U$ e $V$ sono basi ortonormali che spaziano negli spazi colonna e riga di $H$ . Poiché $H$ ha un rango al massimo di $2r$ , ciò riduce l'costosa ortogonalizzazione $d_1 \times d_2$ a una ortogonalizzazione $2r \times 2r$ .
Implementazione: L'intero algoritmo si basa esclusivamente su moltiplicazioni tra matrici, radici quadrate di matrici inverse e decomposizioni polari di piccole matrici $r \times r$ o $2r \times 2r$ . Ciò evita l'overhead di decomposizioni QR e SVD richiesti da Riemannion.

3. Dettagli Algoritmici

L'ottimizzatore sMuon incorpora:

Split Weight Decay: Per corrispondere alla dinamica del training a parametri completi, il weight decay viene applicato tramite un fattore di scala $s = \sqrt{1-\lambda\eta}$ ai fattori $A$ e $B$ prima dell'aggiornamento del gradiente.
Trasporto del Momentum: Per garantire l'invarianza della parametrizzazione (ovvero, l'ottimizzatore si comporta in modo coerente indipendentemente dalla specifica fattorizzazione $B, A$ ), i buffer del momentum vengono proiettati sugli spazi dei fattori correnti ad ogni step.

Contributi Chiave

Derivazione di sMuon: Un metodo in forma chiusa per l'addestramento di adapter LoRA con Muon, prendendo un'approssimazione del primo ordine dell'obiettivo di discesa con norma spettrale regolarizzata.
Implementazione Efficiente: Un algoritmo "matmul-only" che è computazionalmente efficiente per l'hardware moderno, evitando decomposizioni costose. Offre una velocità comparabile a Per-factor Muon e LoRA-Muon, pur essendo significativamente più veloce di Riemannion, specialmente per rank LoRA più elevati.
Valutazione Empirica: Esperimenti estesi su quattro modelli base (incluso uno pre-addestrato con Muon) e undici benchmark (common sense e coding).

Risultati

Il documento riporta le seguenti scoperte:

Performance: sMuon si comporta favorevolmente attraverso i benchmark SFT. Sebbene i risultati dipendano dal modello, sMuo ha raggiunto la migliore accuratezza in sei degli undici task durante il fine-tuning del modello Moonlight-16B-A3B (che è stato pre-addestrato con Muon). In generale, supera AdamW e Per-factor Muon.
Pretraining ReLoRA: In un contesto di pretraining utilizzando ReLoRA (fusione e aggiornamento periodico degli adapter), sMuon ha raggiunto la perdita di validazione finale più bassa, pareggiando con Riemannion ma essendo circa il 30% più veloce.
Proprietà Spettrali: L'analisi empirica degli spettri di aggiornamento dei pesi mostra che sMuon mantiene uno spettro più uniforme (più vicino al comportamento ideale di Muon) rispetto a Per-factor Muon e LoRA-Muom.
Efficienza: La Tabella 4 dimostra che, sebbene sMuon sia leggermente più lento di AdamW, è significativamente più veloce di Riemannion. Per un rank di $r=64$ , Riemannion aggiunge circa il 33% di overhead al tempo di step, mentre sMuon aggiunge solo lo ~0.8%.

Significato e Rivendicazioni

Il paper sostiene che sMuon fornisce un modo principato, efficiente in termini di memoria e computazionalmente trattabile per applicare l'ottimizzazione Muon agli adapter a basso rango. Rispettando la geometria del livello completo invece di trattare i fattori LoRA come matrici indipendenti, sMuon ottiene miglioramenti moderati delle prestazioni, particolarmente quando il modello base è stato anch'esso pre-addestrato con Muon. Gli autori sottolineano che il loro metodo evita gli "onerosi vincoli del manifold" dell'ottimizzazione Riemanniana senza sacrificare le prestazioni, rendendo l'ottimizzazione basata su Muon più accessibile per applicazioni pratiche.

Gli autori rimangono modesti, osservando che i risultati sono "dipendenti dal modello e dalla valutazione" e che il contributo primario è un'approssimazione pulita ed efficiente che colma il divario tra i benefici teorici di Muon e i vincoli pratici di PEFT.