Riepilogo Tecnico: Principio Compatibile con la Simmetria per la Progettazione di Ottimizzatori

1. Enunciato del Problema

Esiste una disparità geometrica fondamentale nella pratica moderna dell'apprendimento profondo. Mentre le architetture delle reti neurali (embedding, testine LM, MLP SwiGLU, router MoE) possiedono ricche proprietà intrinseche di simmetria ed equivarianza, i metodi di ottimizzazione dominanti come Adam, AdamW e RMSprop operano coordinata per coordinata. Questi metodi trattano i parametri matriciali come vettori appiattiti di coordinate indipendenti, ignorando di conseguenza la struttura matriciale e rompendo l'equivarianza naturale dello spazio dei parametri.

Questo approccio "cieco alla geometria" non riesce a distinguere tra le geometrie distinte di diversi tipi di livelli. Di conseguenza, gli ottimizzatori popolari possono scartare la struttura spettrale, indurre geometrie non corrispondenti e non rispettare le strutture di invarianza ed equivarianza intrinseche a specifici blocchi di parametri (ad esempio, l'invarianza per permutazione nei vocabolari o l'intercambiabilità degli esperti nei MoE). La comprensione teorica attuale e i confronti empirici tra metodi adattivi coordinata per coordinata e ottimizzatori spettrali rimangono limitati, basandosi spesso su benchmark su larga scala senza un'analisi approfondita del comportamento geometrico.

2. Metodologia: Il Principio Compatibile con la Simmetria

Il documento propone un principio compatibile con la simmetria a livello di livello per la progettazione degli ottimizzatori. La tesi centrale è che la regola di aggiornamento $\mathcal{U}$ dell'ottimizzatore per un blocco di parametri $W$ dovrebbe essere equivariante sotto il gruppo di simmetria $G$ che agisce su quel specifico blocco.

2.1 Quadro Teorico

Sia $W \in \mathbb{R}^{m \times n}$ una matrice di parametri. Se le coordinate degli spazi di input e output vengono trasformate da matrici invertibili $P$ e $Q$, la mappa lineare si trasforma come $\tilde{W} = P W Q^{-1}$. Il gradiente si trasforma in modo controvariante/covariante di conseguenza. La mappa di aggiornamento dell'ottimizzatore $\mathcal{U}$ deve soddisfare:
$$\mathcal{U}(g \cdot D) = g \cdot \mathcal{U}(D)$$
dove $D$ è una direzione di aggiornamento (ad esempio, gradiente o momento) e $g$ è un elemento del gruppo di simmetria.

Il documento deriva classi specifiche di ottimizzatori basate sui gruppi di simmetria di diversi tipi di livelli:

1. Livelli Matriciali Ordinari (Lineari/Attention):

   • Simmetria: Bi-ortogonale ($O_m \times O_n$).
   • Classe di Ottimizzatore: Ottimizzatori Spettrali. L'aggiornamento deve essere un operatore spettrale che agisce sui valori singolari preservando i vettori singolari. Questo unifica la Discesa Spettrale Stocastica (SSD), Muon, Scion e i metodi del Gradiente Polare. L'aggiornamento è tipicamente il fattore polare ortogonale della direzione del momento.

2. Embedding e Testine LM:

   • Simmetria: Permutazione sinistra ($P_v$) e Ortogonale destra ($O_d$). Le righe indicano gli elementi del vocabolario (permutazione discreta), le colonne indicano le caratteristiche nascoste (rotazione continua).
   • Classe di Ottimizzatore: Aggiornamenti Equivarianti Sinistra-Permutazione Destra-Ortogonale (LPRO).
     • Spettrale Destra: Agisce globalmente tramite la matrice Gram destra ($D^\top D$).
     • Norma di Riga: Agisce localmente sulle norme delle righe (elementi del vocabolario).
     • Ibrido: Composizioni di aggiornamenti di norma di riga e spettrale destra.

3. Proiezioni MLP SwiGLU:

   • Simmetria: Permutazione dei neuroni intermedi ($P_{d_{ff}}$). Le proiezioni gate e up hanno righe corrispondenti ai neuroni intermedi; la proiezione down ha colonne corrispondenti ad essi.
   • Classe di Ottimizzatore: Varianti Consapevoli delle Righe e delle Colonne. Le proiezioni Gate/Up utilizzano aggiornamenti compatibili con LPRO (norma di riga/spettrale destra); le proiezioni Down utilizzano la geometria trasposta.

4. Router MoE:

   • Simmetria: Permutazione degli esperti ($P_e$) e invarianza allo spostamento dei logit condivisi (aggiungere una costante a tutti i logit non cambia la distribuzione softmax).
   • Classe di Ottimizzatore: Aggiornamenti Compatibili con il Router. Questi operano sulla direzione centrata $D_c = \Pi^\perp D$ (rimuovendo lo spostamento condiviso).
     • Spettrale Sinistra Centrato: Agisce sulla matrice Gram degli esperti centrata.
     • Norma di Riga Centrata: Agisce sulle norme delle righe degli esperti centrati.

2.2 Implementazione Pratica

Gli autori istanziano questi principi in ottimizzatori pratici basati sul momento:

• RightPolarGradM / LeftPolarGradM: Aggiornamenti spettrali unilaterali che utilizzano iterazioni di Newton-Schulz per le radici quadrate inverse di matrici.
• RowNormM: Aggiornamenti di normalizzazione riga per riga.
• HybridPolarGradM: Composizioni di passi di norma di riga e spettrali unilaterali.
• Centratura: Per i router MoE, gli aggiornamenti vengono applicati al momento centrato $M_{k,c} = \Pi^\perp M_k$.

3. Contributi Chiave

1. Principio Compatibile con la Simmetria: Un argomento formale secondo cui gli ottimizzatori coordinata per coordinata sono geometricamente non corrispondenti per parametri a valore matriciale. Il documento stabilisce che le regole di aggiornamento dovrebbero corrispondere al gruppo di simmetria del blocco di parametri.
2. Visione Unificante degli Ottimizzatori Spettrali: Fornisce un'interpretazione basata sulla simmetria degli ottimizzatori spettrali (Muon, SSD, ecc.) come l'unica classe di mappe di aggiornamento bi-ortogonalmente equivarianti.
3. Nuove Classi di Ottimizzatori Equivarianti: Deriva regole di aggiornamento specifiche per livelli con simmetrie non standard:
   • Aggiornamenti equivarianti LPRO per embedding e testine LM.
   • Aggiornamenti consapevoli di righe/colonne per MLP SwiGLU.
   • Aggiornamenti di permutazione degli esperti centrati per router MoE.
4. Evidenza di Pre-addestramento End-to-End: Costruisce una pila completa di ottimizzatori a livello di livello dove ogni principale classe matriciale riceve un aggiornamento corrispondente alla simmetria. Questo viene valutato su modelli linguistici MoE densi e sparsi (Qwen3-0.6B, Gemma 3 1B, OLMoE-1B-7B, gpt-oss ridimensionato).

4. Risultati Sperimentali

Gli autori hanno condotto esperimenti di pre-addestramento su sottoinsiemi di 10 miliardi di token di FineWeb-Edu attraverso quattro architetture.

• Modelli Densi (Qwen3-0.6B, Gemma 3 1B):

  • Sostituire AdamW sulle matrici indicizzate dal vocabolario (embedding/testine LM) con RowNormM o HybridPolarGradM ha migliorato costantemente la perdita di validazione finale.
  • I guadagni sono stati modesti per il modello Qwen3 più piccolo ma più pronunciati per il modello Gemma 3 più grande, suggerendo che gli aggiornamenti consapevoli della geometria diventano più criti man mano che le matrici del vocabolario crescono e i gradienti diventano più anisotropi.
  • L'uso di HybridPolarGradM per le proiezioni MLP SwiGLU ha ulteriormente migliorato le prestazioni, probabilmente a causa della correzione dello squilibrio di scala delle righe nelle matrici alte e sottili.

• Modelli MoE Sparsi (OLMoE-1B-7B, gpt-oss ridimensionato):

  • Le configurazioni che utilizzavano aggiornamenti compatibili con la simmetria per embedding, testine LM e router (ad esempio RowNormM + LeftPolarGradM) hanno superato le configurazioni che utilizzavano AdamW per questi componenti.
  • La baseline completamente coordinata (AdamW ovunque) ha ottenuto i risultati peggiori.
  • Gli aggiornamenti compatibili con la simmetria per i router hanno migliorato rispetto agli aggiornamenti per router coordinata per coordinata e, in alcuni casi, hanno ridotto i picchi di perdita durante l'addestramento, suggerendo una maggiore stabilità.

• Tendenza Generale: In tutti i modelli, la sostituzione di AdamW su grandi matrici indicizzate dal vocabolario con ottimizzatori compatibili con la simmetria ha costantemente abbassato la perdita di validazione finale. I benefici sono persistiti attraverso diverse architetture e livelli di sparsità.

5. Significato e Affermazioni

Il documento afferma di offrire un cambio di paradigma nella progettazione degli ottimizzatori, passando da una visione coordinata "cieca alla geometria" a una visione matriciale "consapevole della simmetria".

• Co-Progettazione Architettura-Ottimizzatore: Il lavoro sostiene che la progettazione dell'ottimizzatore dovrebbe essere un problema consapevole dell'architettura. Diverse classi di parametri (embedding, router, livelli lineari) hanno simmetrie naturali distinte e richiedono geometrie di ottimizzatore distinte.
• Fondamento Teorico: Fornisce una caratterizzazione matematica rigorosa (tramite equivarianza) del motivo per cui gli ottimizzatori spettrali funzionano per i livelli lineari e estende questa logica per derivare nuovi ottimizzatori per altri livelli.
• Utilità Pratica: I risultati dimostrano che le assegnazioni compatibili con la simmetria non sono solo teoricamente valide ma anche praticamente vantaggiose, migliorando la perdita finale e potenzialmente la stabilità dell'addestramento nel pre-addestramento su larga scala.
• Modestia: Gli autori affermano esplicitamente di non sostenere che gli ottimizzatori equivarianti dominino i metodi coordinata per coordinata in tutti i regimi. Invece, presentano prove che il principio compatibile con la simmetria a livello di livello porta ad assegnazioni competitive e spesso vantaggiose in impostazioni di pre-addestramento rappresentative. Considerano i risultati come prove dell'utilità del principio piuttosto che come un benchmark esaustivo.

Il documento conclude che l'introduzione di simmetria ed equivarianza nella progettazione degli ottimizzatori apre la strada a un pre-addestramento di modelli su larga scala più stabile, coerente con la parametrizzazione e fondato teoricamente.