Riepilogo Tecnico: AMUSE (Anytime Muon con Valutazione Stabile del Gradiente)

Enunciato del Problema

L'ottimizzazione moderna nell'apprendimento profondo si basa tipicamente sull'ottimizzatore AdamW combinato con schedule di tasso di apprendimento prescritti (ad esempio, decadimento coseno). Recenti avanzamenti hanno messo in discussione questi componenti individualmente: l'ottimizzazione Schedule-Free (SF) rimuove gli schedule espliciti valutando i gradienti in punti interpolati tra le iterazioni correnti e quelle mediate, mentre Muon migliora la geometria dell'aggiornamento per i parametri a valore matriciale ortogonalizzando il momento.

Nonostante il successo empirico di Muon nel preaddestramento di grandi modelli linguistici (LLM), il suo meccanismo sottostante non è completamente compreso. Il documento identifica una specifica instabilità in Muon: mentre la sua ortogonalizzazione accelera il progresso lungo le direzioni "piatte" (il "fiume" del paesaggio della funzione di perdita), amplifica involontariamente il rumore nelle direzioni "ripide" ad alta curvatura (le "pareti della valle"), portando a traiettorie oscillatorie. Al contrario, mentre i metodi Schedule-Free stabilizzano le traiettorie mantenendole vicino al fiume, potrebbero mancare del rapido progresso di massa offerto dall'ortogonalizzazione di Muon. Il problema fondamentale è progettare un ottimizzatore che combini la rapida convergenza di Muon con la stabilità della media Schedule-Free senza richiedere schedule manuali di tasso di apprendimento.

Metodologia: AMUSE

Gli autori propongono AMUSE (Anytime Muon con Valutazione Stabile del Gradiente), che integra gli aggiornamenti matriciali ortogonalizzati di Muon con il framework Schedule-Free utilizzando un coefficiente di interpolazione variabile nel tempo.

Fondamento Teorico: Paesaggio della Funzione di Perdita Fiume-Valle

L'analisi si basa sulla prospettiva del paesaggio della funzione di perdita fiume-valle:

• Il Fiume: Uno spazio vettoriale bulk ad alta dimensionalità e bassa curvatura dove avviene il progresso utile dell'addestramento.
• Le Pareti della Valle: Uno spazio vettoriale dominante a bassa dimensionalità e alta curvatura che induce oscillazioni.
• Meccanismo di Muon: L'ortogonalizzazione della matrice del momento aumenta la componente bulk relativa degli aggiornamenti, accelerando il progresso lungo il fiume. Tuttavia, impedisce anche alle componenti dominanti di decadere, causando l'amplificazione di piccoli rumori nelle direzioni dominanti in oscillazioni attraverso le pareti della valle.
• Meccanismo Schedule-Free: Valutare i gradienti su una traiettoria media (più vicina al fiume) filtra le fluttuazioni ad alta curvatura prima del passo di aggiornamento.

Progettazione dell'Algoritmo

AMUSE impiega un coefficiente di interpolazione variabile nel tempo $\beta_t$ per determinare il punto di valutazione del gradiente $Y_t$ tra la sequenza base veloce $Z_t$ (aggiornamenti Muon) e la sequenza media stabile $X_t$:
$$Y_t = (1 - \beta_t)Z_t + \beta_t X_t$$

Lo schedule per $\beta_t$ è progettato per bilanciare l'adattamento iniziale e la stabilità nella fase finale:

1. Fase di Warmup ($t \le T_0$): $\beta_t$ è fissato a un valore piccolo $\beta_1$. Ciò permette all'ottimizzatore di valutare i gradienti vicino alla sequenza veloce di Muon ($Z_t$) per un adattamento iniziale rapido.
2. Post-Warmup ($t > T_0$): $\beta_t$ aumenta gradualmente verso 1 secondo la formula:
   $$\beta_t = 1 - \left( \frac{T_0 - 1}{t - 1} \right)^\rho (1 - \beta_1)$$
   Qui, $\rho \in [0, 1]$ controlla il tasso di transizione. All'aumentare di $\beta_t$, il punto di valutazione del gradiente si sposta verso la sequenza media $X_t$. Ciò sopprime le componenti ad alta curvatura (dominanti) del gradiente prima che subiscano l'ortogonalizzazione di Muon, stabilizzando così la traiettoria.

Crucialmente, questo schedule è indipendente dall'orizzonte totale di addestramento, abilitando l'addestramento anytime (l'addestramento può essere interrotto in qualsiasi momento senza degradazione delle prestazioni). AMUSE non richiede alcuno schedule di decadimento del tasso di apprendimento.

Contributi Chiave

1. Analisi Geometrica di Muon: Il documento fornisce una spiegazione geometrica dell'efficienza e dell'instabilità di Muon. Dimostra che l'ortogonalizzazione promuove aggiornamenti orientati al bulk (accelerando il progresso del fiume) ma amplifica il rumore nelle direzioni dominanti (causando oscillazioni sulle pareti della valle).
2. Stabilizzazione tramite Interpolazione: Gli autori mostrano che valutare i gradienti più vicini alla traiettoria media (tramite Schedule-Free) riduce le componenti dominanti prima dell'ortogonalizzazione, mitigando le oscillazioni.
3. Algoritmo AMUSE: Un nuovo ottimizzatore che sposta dinamicamente il punto di valutazione del gradiente dalla traiettoria veloce di Muon alla traiettoria media stabile. Ciò combina la rapida convergenza di Muon con la stabilità dei metodi Schedule-Free.
4. Supporto all'Addestramento Anytime: Utilizzando uno schedule variabile nel tempo indipendente dal numero totale di passi ed evitando il decadimento del tasso di apprendimento, AMUSE supporta orizzonti di addestramento flessibili.

Risultati Sperimentali

AMUSE è stato valutato su diversi benchmark, inclusa la classificazione di immagini (CIFAR-10/100, SVHN, ImageNet-1k), la segmentazione di immagini (ISIC 2018) e il preaddestramento di LLM (modelli in stile Llama con 124M, 720M e 1B di parametri su FineWeb).

• Frontiera di Pareto Prestazioni-Iterazioni: AMUSE supera costantemente (Schedule-Free) AdamW e Muon in tutti i compiti.
• Guadagni di Efficienza:
  • Sul preaddestramento di LLM da 720M, AMUSE raggiunge le prestazioni finali di Muon con 1,51$\times$ meno passi.
  • Su ImageNet con ResNet-50, ottiene le stesse prestazioni con 1,12$\times$ meno passi.
  • Su ImageNet con fine-tuning di ViT, richiede 3,08$\times$ meno passi.
• Stabilità: A differenza di Muon vanilla, che beneficia significativamente della media esponenziale dei pesi (EWA) a posteriori o del decadimento del tasso di apprendimento, AMUSE mantiene una traiettoria stabile durante tutto l'addestramento, mostrando un miglioramento minimo quando queste tecniche a posteriori vengono applicate.
• Scalabilità: AMUSE scala efficacemente a grandi dimensioni di batch (fino a 2048) e grandi scale di modelli (1B di parametri), dove altre varianti Schedule-Free a volte faticano.

Significato e Affermazioni

Il documento afferma che AMUSE rappresenta un approccio principiato alla progettazione di ottimizzatori derivato da una comprensione sistematica delle dinamiche di addestramento e della geometria del paesaggio della funzione di perdita. Affrontando esplicitamente il compromesso tra rapido progresso di massa e stabilità delle pareti della valle, AMUSE raggiunge prestazioni superiori senza la necessità di una pianificazione manuale del tasso di apprendimento.

Gli autori notano modestamente che, sebbene AMUSE superi gli ottimizzatori esistenti, comporta un leggero sovraccarico di memoria (una copia aggiuntiva dello stato per la sequenza base $Z_t$) rispetto a Muon vanilla, sebbene rimanga paragonabile ad AdamW e Schedule-Free AdamW. Suggeriscono che il lavoro futuro potrebbe concentrarsi sul raggiungimento di una stabilità e prestazioni simili senza questo costo aggiuntivo di memoria. Il documento sottolinea che la prospettiva "fiume-valle" offre un quadro robusto per progettare ottimizzatori che siano sia efficienti che stabili.