Sintesi Tecnica: GeoNorm: Unificare Pre-Norm e Post-Norm con l'Ottimizzazione Geodetica

Definizione del Problema

La collocazione dei livelli di normalizzazione all'interno delle architetture Transformer — specificamente la scelta tra Pre-Norm e Post-Norm — rimane una decisione progettuale critica che influenza la stabilità dell'addestramento, le prestazioni del modello e il flusso del gradiente. Sebbene la Pre-Norm sia diventata la scelta predominante nei moderni Large Language Models (LLM) grazie alla sua stabilità, essa introduce gradienti sbilanciati in cui gli strati inferiori ricevono aggiornamenti più ampi rispetto a quelli superiori. Al contrario, la Post-Norm, originariamente progettata per facilitare l'addestramento di reti profonde, è soggetta a instabilità dell'addestramento e picchi di perdita (loss spikes). Le soluzioni esistenti come DeepNorm e SandwichNorm migliorano la stabilità attraverso l'inizializzazione dei parametri o l'aggiunta di ulteriori livelli di normalizzazione, ma mancano di un quadro teorico unificato per analizzare tali meccanismi.

Metodologia

Gli autori propongono GeoNorm, un nuovo metodo di normalizzazione basato sull'ottimizzazione su varietà (manifold optimization). L'intuizione teorica centrale è quella di reinterpretare le operazioni standard di aggiornamento del Transformer ($x + \text{FFN}(x)$ e $x + \text{Attention}(x)$) come passi di ottimizzazione su una sfera. In questa visione:

1. Gli embedding dei token agiscono come lo stato iniziale su una varietà.
2. I moduli Attention e FFN generano direzioni di aggiornamento (gradienti stimati).
3. La normalizzazione è vista come una proiezione su un insieme di vincoli (la sfera).

La normalizzazione standard (Post-Norm) esegue un aggiornamento non vincolato seguito da una proiezione euclidea sulla sfera. Gli autori sostengono che questa proiezione scarti informazioni geometriche e possa distorcere le direzioni di aggiornamento. Invece, GeoNorm sostituisce il passo di proiezione con un aggiornamento geodetico utilizzando la mappa esponenziale.

Per un punto $x$ su una sfera e un vettore tangente $v$ (la direzione di aggiornamento proiettata), l'aggiornamento viene eseguito intrinsecamente lungo la geodetica:
$$\text{exp}_x(v) = \cos\left(\frac{\|v\|}{\|x\|}\right)x + \|x\|\sin\left(\frac{\|v\|}{\|x\|}\right)\frac{v}{\|v\|}$$

Questa formulazione assicura che la rappresentazione si muova fluidamente lungo la varietà, preservando la geometria intrinseca dello spazio delle feature. Inoltre, gli autori introducono un meccanismo di decadimento dell'aggiornamento per livello (layer-wise update decay) per la dimensione del passo ($\alpha_k$), analogamente agli schedule del learning rate. Valutano tre strategie di decadimento:

• Armonica: $\alpha_k = \alpha / k$
• Radice Quadrata: $\alpha_k = \alpha / \sqrt{k}$
• Lineare: $\alpha_k = \alpha (T - k) / T$

Il documento dimostra anche teoricamente che la Pre-Norm può essere interpretata come un caso speciale di GeoNorm con angoli non adattivi e dipendenti dalla profondità, unificando così i due approcci sotto un unico quadro geometrico.

Contributi Chiave

1. Quadro Teorico: Gli autori forniscono una prospettiva teorica unificata che interpreta i livelli Transformer come passi di ottimizzazione iterativi su una varietà sferica, colmando il divario tra Pre-Norm e Post-Norm.
2. Algoritmo GeoNorm: Introducono GeoNorm, che sostituisce la normale normalizzazione basata su proiezione con aggiornamenti geodetici tramite la mappa esponenziale. Questo metodo non richiede parametri addestrabili aggiuntivi e aggiunge un sovraccarico computazionale trascurabile.
3. Decadimento per Livello: Una proposta innovativa di decadimento dell'aggimento per livello per i componenti di attenzione e FFN per stabilizzare la dinamica dell'addestramento.
4. Validazione Empirica: Esperimenti estesi su varie dimensioni di modello (da 125M a 1.3B), dataset (Arxiv, Books3, FinWeb-Edu) e lunghezze di addestramento dimostrano l'efficacia del metodo.

Risultati Sperimentali

Il documento riporta che GeoNorm supera costantemente i metodi di normalizzazione esistenti (Pre-Norm, Post-Norm, DeepNorm, SandwichNorm) in tutte le configurazioni testate:

• Prestazioni: Sui dataset Arxiv e Books3 con un modello da 125M, GeoNorm ha ottenuto la perdita di validazione più bassa (1.8792 su Arxiv; 3.4040 su Books3) rispetto a tutti i baseline.
• Scalabilità: Il vantaggio di prestazioni persiste e si amplia all'aumentare della dimensione del modello. Per un modello da 1.3B, GeoNorm ha ottenuto una perdita di 3.1759, superando il miglior baseline (Pre-Norm a 3.2067).
• Stabilità: GeoNorm ha dimostrato una stabilità di addestramento superiore, in particolare in sessioni di addestramento lunghe (fino a una sequenza di 4096). Mentre la Post-Norm ha mostrato picchi di perdita e collassi dell'addestramento a lunghezze maggiori, GeoNorm ha mantenuto una convergenza stabile.
• Meccanismo di Decadimento: Tra le strategie di decadimento, il decadimento armonico ($\alpha/k$) ha prodotto le migliori prestazioni, stabilendo nuovi record sui dataset testati.

Significato e Rivendicazioni

Il documento afferma che GeoNorm offre un'integrazione fluida nelle architetture Transformer standard, richiedendo nessuna modifica ai blocchi core (Attention o FFN) e non introducendo parametri addestrabili aggiuntivi. Gli autori sostengono che, sfruttando l'ottimizzazione geodetica, GeoNorm consente alle rappresentazioni di muoversi in modo più coerente lungo la varietà sferica, portando a una migliore propagazione delle informazioni e stabilità dell'addestramento.

Il lavoro si posiziona come un passo fondamentale per comprendere i meccanismi di normalizzazione attraverso la lente della geometria di Riemann. Suggerisce che il successo empirico della Pre-Norm possa essere spiegato dal suo decadimento implicito della dimensione del passo, un meccanismo che GeoNorm rende esplicito e adattivo. Gli autori concludono che questa visione unificata fornisce una base robusta per future innovazioni nel design delle architetture Transformer, in particolare per scalare i modelli verso profondità e complessità maggiori.