Aurora di Tilde Research Risolve un Problema Nascosto di Morte dei Neuroni in Muon

La ricerca di Tilde Research presenta Aurora, un ottimizzatore per l’addestramento di reti neurali che corregge un problema strutturale nel noto ottimizzatore Muon. Il difetto, finora nascosto, comporta la morte permanente di neuroni in reti di tipo MLP durante il training e rende impossibile il recupero. Aurora propone una soluzione matematica diversa, con un esperimento di pre-training su 1,1 miliardi di parametri, nuovi risultati su benchmark come modded-nanoGPT e codice aperto.

Che cos'è Muon?

Per comprendere Aurora, è utile introdurre prima Muon, che ha visto il ritorno di interesse nella comunità ML dopo aver superato AdamW nel tempo di convergenza nel concorso modded-nanoGPT — una misura del velocità di addestramento di GPT-style modelli su loss obiettivo specifico. Da allora, Muon è stato adottato in sperimentazioni su modelli di grandi dimensioni da gruppi di ricerca avanzati.

Il passo chiave di Muon è calcolare il fattore polare della matrice dei gradienti. Data una matrice di gradiente G con SVD sottile G = UΣVᵀ, Muon calcola polar(G) = UVᵀ, che è la matrice semiortogonale più vicina a G nella norma di Frobenius. Questo gradiente ortogonalizzato viene poi utilizzato per aggiornare i pesi: W ← W − η UVᵀ dove η è il tasso di apprendimento. L’utilizzo di algoritmi soltanto matmul per il calcolo del fattore polare rende Muon realizzabile a largo raggio.

L’enigma di NorMuon: La Normalizzazione di Riga aiuta, ma perché?

Prima dell’arrivo di Aurora, NorMuon dominava il concorso modded-nanoGPT. Introduceva un passo di normalizzazione delle righe, simile alla scalatura per parametro nel modello Adam. Anche se questo sovrasposta inizialmente i gradienti ortogonali, NorMuon mostra ancora buoni risultati. Il team Tilde si è posto domande precise: Che cosa NorMuon risolve che Muon non riesce a gestire? Che problema sta veramente colmando?

Ill problemi centrali: Anisotropia della Norma Riga e la Morte dei Neuroni in Matrici Alte

Gli algoritmi di Muon comportano un problema strutturale nascosto: in matrici di tipo alta (molte righe rispetto alle colonne), la normalizzazione forzata non riesce a mantenere un aggiornamento uniforme a tutti i neuroni. Questo genera un problema di anisotropia della norma, che produce aggiorni enormi a certi neuroni e virtualmente nessuno ad altri. Il risultato è un "effetto effetto spirale" dove neuroni inattivi ricevono sempre meno segnale, diventando definitivamente inutili nel modello.

I test iniziali evidenziano che entro il 500° passo di training, più di un quarto dei neuroni finisce inattivo. Questo problema non è isolato; influisce sull’intero modello, indebolendo l’apprendimento in strati successivi. Aurora affronta il problema introducendo un approccio matematico nuovo che mantiene aggiorni uniformi, senza compromettere la precisione dell’ortogonalizzazione.

L’intermediato U-NorMuon

Prima dell’arrivo ufficiale di Aurora, la squadra ha introdotto un’altra ottimizzazione temporanea detta U-NorMuon. Il punto chiave di U-NorMuon è che NorMuon normalizza ciascuna riga a norma 1, ma in matrici alte, la norma media corretta per riga dovrebbe essere √n/m, non 1. U-NorMuon corregge questo, normalizzando in modo appropriato le righe a norma √n/m, anziché 1.

Nel contesto di modelli su larga scala (su 340 miliardi di parametri), U-NorMuon batte entrambi Muon e NorMuon standard, eliminando completamente la morte dei neuroni. Le leve di gradi (scores) diventano isotropiche durante l’addestramento. Il beneficio di U-NorMuon si estende anche alle matrici che non modifica direttamente, stabilizzando matrici in giù, senza intervento diretto.

Tuttavia, U-NorMuon ha un limite: sovrascrive in modo forzato il fattore polare impostando norme uguali. Questo compromette la precisione del fattore polare, riducendo l’efficacia complessiva e non risolvendo completamente il problema. È questo l’aspetto che ha spinto i ricercatori a sviluppare Aurora.

Aurora: Discesa più veloce da due vincoli congiunti

Aurora ricalcola l’approccio in modo radicale. Invece di prima ortogonalizzare e poi correggere con normalizzazione riga per riga, Aurora si chiede: Che cosa è l’aggiornamento ideale con vincoli congiunti di semi-ortogonalità a sinistra e norme riga uniformi?

Formalmente, per matrici alte, Aurora risolve il problema:

U∗ = arg⁡max_U Tr(G^T U), con le condizioni:

• U^T U = In per semi-ortogonalità
• ∥Ui:∥ = √n/m ∀i per uniformità

I ricercatori dimostrano che combinando questi due vincoli, tutti i valori singolari di U diventano esattamente 1; dunque, l’aggiornamento risultante mantiene semi-ortogonalità perfetta. Questo è il punto chiave che distingue Aurora da U-NorMuon e NorMuon: raggiunge la distribuzione uniforme e l’ortogonalità insieme, non facendoci a vicenda.

Implementazioni e Algoritmi

Aurora introduce due soluzioni algoritmiche:

• Riemannian Aurora utilizza un approccio di proiezione gradiente limitata sulla varietà di Stiefel/equal-row-leverage.
• Vanilla Aurora è un’implementazione più semplice e pratica.

Entrambi sono open source. Nelle matrici non alte (quadrate o larghe), l’uniformità delle norme è già garantita dall’ortogonalità, così Aurora non modifica tali parametri.

Risultati e Confronti

Gli esperimenti con Aurora hanno portato al training di un modello di 1,1 miliardi di parametri che mostra efficienza del dato a 100x rispetto al dataset internet open source, battere modelli più grandi su valutazioni generali come HellaSwag e ottenere record mondiali su modded-nanoGPT.

A questa dimensione (1 miliardo), Aurora supera