Sintesi Tecnica: Miglioramento dell'addestramento delle reti neurali tramite il disaccoppiamento della magnitudo e della direzione dei vettori di peso

1. Definizione del problema

L'addestramento delle moderne reti neurali si affida ad ottimizzatori come Adam e Muon che trattano le matrici di peso come oggetti singoli e monolitici. Tuttavia, ogni matrice di peso $W$ trasporta intrinsecamente due quantità distinte: una magnitudo ($\|W\|$) e una direzione ($\hat{W} = W/\|W\|$). Nell'ottimizzazione standard, queste due quantità sono accoppiate, portando a dinamiche di addestramento instabili e imprevedibili:

• Dipendenza dal cambiamento direzionale: Il cambiamento angolare indotto da un aggiornamento dipende inversamente dalla magnitudo corrente. Di conseguenza, il tasso di apprendimento (LR) non controlla direttamente il tasso di cambiamento direzionale, che può variare tra i livelli e nel tempo.
• Deriva della magnitudo non intenzionale: Gli aggiornamenti tendono a essere perpendicolari ai pesi attuali (a causa dell'invarianza di scala o del rumore), causando una crescita della magnitudo anche in assenza di un gradiente radiale. Questa "crescita incontrollata" costringe i professionisti a fare affidamento su meccanismi indiretti come il weight decay per stabilizzare l'addestramento e il warmup per prevenire l'instabilità precoce.
• Limitazioni di scalabilità: Gli attuali metodi per scalare l'addestramento (ad es. $\mu$P) spesso richiedono una ritaratura degli iperparametri quando si cambia la larghezza o la profondità del modello. Inoltre, approcci recenti che vincolano i pesi a una sfera a norma fissa mancano di parametri di magnitudo apprendibili, lasciando il controllo della scala esclusivamente ai livelli di normalizzazione.

2. Metodologia: Disaccoppiamento Magnitudo–Direzione (MD)

Gli autori propongono il Disaccoppiamento Magnitudo–Direzione (MD), una modifica dell'ottimizzatore che fattorizza ogni matrice di peso in una direzione a norma fissa e guadagni di magnitudo apprendibili, aggiornandoli con tassi di apprendimento separati.

Fattorizzazione Core

Ogni matrice di peso $W$ è riparametrizzata come:
$$W = \text{diag}(\gamma_{\text{row}}) \, \hat{W} \, \text{diag}(\gamma_{\text{col}})$$
Dove:

• $\hat{W}$ è la direzione, vincolata a risiedere su un'ipersfera fissa (ad es. norma di Frobenius).
• $\gamma_{\text{row}}$ e $\gamma_{\text{col}}$ sono guadagni di magnitudo apprendibili (vettori che agiscono rispettivamente su righe e colonne).

Procedura di Ottimizzazione

Il metodo opera interamente all'interno dello step dell'ottimizzatore, mantenendo invariato il forward pass del modello (il modello vede un singolo tensore fuso $W$):

1. Recupero: Recuperare $\hat{W}$ e i guadagni $\gamma$ dal peso fuso $W$.
2. Suddivisione del Gradiente: Dividere il gradiente $G = \partial L / \partial W$ in un gradiente di direzione ($\partial L / \partial \hat{W}$) e gradienti di guadagno (sommati sugli assi rispettivi).
3. Aggiornamento della Direzione: Aggiornare $\hat{W}$ utilizzando l'ottimizzatore base (ad es. Adam, Muon) con un tasso di apprendimento $\eta_W$, quindi proiettarlo nuovamente sulla sfera per mantenere la norma fissa.
4. Aggiornamento dei Guadagni: Aggiornare i guadagni $\gamma$ (opzionalmente tramite una mappa fluida come softplus per garantire la positività) utilizzando un tasso di apprendimento separato $\eta_\gamma$.
5. Ricostruzione: Ricostruire $W$ per il successivo forward pass.

Proprietà Chiave

• Indipendente dall'Ottimizzatore: Il metodo agisce come un wrapper attorno a qualsiasi ottimizzatore base (Adam, Muon, Shampoo, ecc.).
• Controllo Diretto del LR: Fissando la direzione su una sfera, il tasso di apprendimento imposta direttamente il relativo aggiornamento angolare, rimuovendo la dipendenza dalla magnitudo attuale del peso.
• Eliminazione degli Euristiche: Il metodo elimina la necessità di weight decay (poiché i pesi sono già vincolati nella norma) e warmup (poiché gli aggiornamenti iniziali grandi e destabilizzanti sono regolati fin dal primo passo).
• Controllo della Scala Fine: A differenza dei metodi che si affidano esclusivamente ai guadagni dei livelli di normalizzazione, l'MD Decoupling fornisce un controllo esplicito e appreso sulla scala di singole righe e colonne, permettendo al modello di amplificare o attenuare specifiche caratteristiche.

3. Contributi Principali e Risultati

Risultati Empirici

Gli autori validano l'MD Decoupling attraverso modelli densi (da 181M a 1.29B di parametri) e grandi modelli Mixture-of-Experts (MoE) (fino a 6.7B di parametri totali).

• Guadagni di Performance: L'MD Decoupling migliora rispetto alle baseline ben tarate (AdamW e Muon). Sui grandi modelli MoE, MuonMD (Muon con MD Decoupling) raggiunge la perdita di AdamW con circa 2$\times$ meno calcolo.
• Trasferimento del Tasso di Apprendimento: Il metodo consente il trasferimento zero-shot del tasso di apprendimento attraverso la larghezza del modello. Il LR ottimale per matrice trovato su un modello piccolo rimane ottimale per modelli significativamente più grandi senza ritaratura. Questo vale sia per la scalabilità della larghezza che della profondità.
• Addestramento Senza Warmup: Gli esperimenti dimostrano che eliminare il warmup non solo mantiene la stabilità ma migliora la perdita finale, poiché i primi passi di addestramento non sono più sprecati in tassi di apprendimento ridotti.
• Insight dalle Ablazioni:
  • Asse di Normalizzazione: Si è scoperto che vincolare la direzione tramite norma di Frobenius (intera matrice) è il metodo più flessibile e performa comparabilmente ai vincoli riga/colonna.
  • Parametrizzazione del Guadagno: Un guadagno combinato per riga e colonna ($\gamma_{\text{row}} + \gamma_{\text{col}}$) supera i guadagni scalari o a singolo asse. La parametrizzazione softplus per i guadagni fornisce un vantaggio lieve ma costante rispetto agli aggiornamenti diretti o esponenziali.
  • Embedding: Vincolare i vettori di embedding alla norma L2 unitaria performa leggermente meglio rispetto al lasciarli non vincolati.

Confronto con il Lavoro Correlato

Il paper colloca l'MD Decoupling all'interno di un corpo crescente di studi sul training sferico e sulla riparametrizzazione:

• A differenza di nGPT, che unisce i vincoli sferici a specifiche modifiche architettoniche (ad es. normalizzazione L2, residui interpolati), l'MD Decoupling è una modifica lato ottimizzatore che può essere applicata ad architetture standard. Applicato all'architettura di nGPT, l'MD Decoupling migliora ulteriormente le prestazioni.
• A differenza di Muown o HyperP, che possono fissare le norme senza magnitudo apprendibili o legare il metodo a ottimizzatori specifici, l'MD Decoupling apprende esplicitamente le magnitudo ed è indipendente dall'ottimizzatore.
• Affronta l'interferenza "magnitudo-direzione" identificata in lavori precedenti (ad es. Kosson et al., 2024a) separando la dinamica delle due quantità.

4. Significato e Rivendicazioni

Il paper afferma che trattare la magnitudo e la direzione come quantità controllate separatamente produce dinamiche di addestramento più prevedibili e un miglioramento semplice e ampiamente applicabile ai moderni ottimizzatori.

• Semplicità: Il metodo non richiede cambiamenti architettonici e aggiunge un sovraccarico computazionale minimo (operazioni elemento per elemento durante lo step dell'ottimizzatore).
• Robustezza: Semplifica la ricetta di addestramento eliminando la necessità di weight decay e warmup, riducendo lo spazio di ricerca degli iperparametri.
• Scalabilità: Fornisce un meccanismo principato per il trasferimento del tasso di apprendimento attraverso la larghezza del modello, affrontando un collo di bottiglia chiave nella scalabilità dei grandi modelli linguistici.

Gli autori concludono che, sebbene la logica aggiuntiva nello step dell'ottimizzatore introduca un modesto sovraccarico, i benefici nella stabilità dell'addestramento, nell'efficienza computazionale e nel trasferimento degli iperparametri rendono l'MD Decoupling un avanzamento significativo per l'addestramento di reti neurali su larga scala.