Sintesi Tecnica: Stabilità dei Transformer sotto la Normalizzazione di Livello (Layer Normalization)

Definizione del Problema

Nonostante il loro dominio nel deep learning moderno, l'addestramento di Transformer profondi rimane incline all'instabilità, in particolare all'aumentare della profondità del modello. Queste instabilità si manifestano in due aree critiche: instabilità forward (crescita incontrollata degli stati nascosti che porta a problemi numerici) e instabilità backward (esplosione dei gradienti durante la backpropagation).

Una componente architettonica primaria utilizzata per mitigare questi problemi è la Normalizzazione di Livello (Layer Normalization - LN). Tuttavia, il posizionamento della LN all'interno dei blocchi Transformer si è evoluto principalmente attraverso euristiche empiriche piuttosto che attraverso una rigorosa giustificazione teorica. Il documento identifica tre configurazioni principali:

• Post-LN: LN applicata dopo le connessioni residue. Spesso richiede una pianificazione delicata dell'ottimizzazione ed esibisce prestazioni subottimali.
• Pre-LN: LN applicata prima dei moduli di attenzione e feedforward. Sebbene sia lo standard per la sua migliore stabilità di addestramento rispetto a Post-LN, è empiricamente osservato che produce stati nascosti eccessivamente grandi, portando a potenziali instabilità numeriche e all'esplosione dei gradienti in reti profonde.
• Peri-LN: Una configurazione più recente che applica la LN sia all'input che all'output dei moduli. Gli studi empirici suggeriscono che offre una stabilità superiore e stati nascosti più regolari rispetto a Pre-LN, ma le sue proprietà teoriche rimangono scarsamente comprese.

Il problema centrale affrontato è la mancanza di un quadro teorico fondato per spiegare perché questi diversi posizionamenti portino a dinamiche di addestramento e profili di stabilità distinti, e se l'addestramento spinga i modelli verso soluzioni regolari o comportamenti patologici.

Metodologia

Gli autori propongono un nuovo quadro teorico basato sulla teoria del controllo ottimo a tempo discreto per analizzare le dinamiche di addestramento dei Transformer.

1. Formulazione come Controllo Ottimo: Il problema di addestramento del Transformer è formulato come un problema di controllo mean-field a tempo discreto. L'indice dello strato $i$ rappresenta il tempo discreto, e lo stato nascosto $X_i$ rappresenta lo stato del sistema. Le connessioni residue agiscono come la velocità delle dinamiche.
2. Analisi Hamilton-Jacobi-Bellman (HJB): Gli autori analizzano le condizioni di ottimalità del problema di addestramento utilizzando l'equazione differenziale alle derivate parziali (PDE) HJB accoppiata con equazioni di trasporto di massa.
   • Esaminano l'Hamiltoniano, che funge da meccanismo di controllo a feedback definendo il campo di velocità ottimale per gli stati nascosti.
   • Per Pre-LN, dimostrano che l'Hamiltoniano è illimitato, portando a condizioni di ottimalità mal poste in cui gli stati nascosti possono crescere arbitrariamente grandi.
   • Per Peri-LN, la normalizzazione in uscita vincola il campo di velocità a un ellissoide (basato sul Lemma 1), garantendo che l'Hamiltoniano sia limitato e il problema sia ben posto.
3. Analisi del Gradiente: Il documento analizza la stabilità inversa esaminando la propagazione dei gradienti attraverso il prodotto di sensibilità locali (matrici Jacobiane) attraverso gli strati.
4. Scaling del Residuo: Il framework viene esteso includendo un parametro di incremento temporale $\Delta t < 1$, che scala efficacemente l'entità dei passi residui per controllare la propagazione del segnale.

Contributi Chiave

1. Framework Teorico per Modelli Addestrati

A differenza del lavoro precedente che spesso si concentra sull'inizializzazione del modello o si affida a prove empiriche, questo framework analizza i modelli addestrati alla convergenza. Fornisce una valutazione sistematica del fatto che le dinamiche di addestramento conducano a soluzioni regolari o patologie.

2. Analisi Rigorosa dei Posizionamenti della LN

• Instabilità Pre-LN: Il documento dimostra che, per Pre-LN, la soluzione ottimale del problema di addestramento è illimitata in magnitudo. Anche con il weight decay, gli stati nascosti possono crescere esponenzialmente ($O(e^D)$), e i gradienti crescono proporzionalmente alle attivazioni, portando all'instabilità dell'addestramento.
• Stabilità Peri-LN: Gli autori derivano limiti espliciti che dimostrano come Peri-LN garantisca una crescita lineare per il valore assoluto medio degli stati nascosti ($O(D)$) e una crescita quadratica per la varianza ($O(D^2)$). Fondamentalmente, la sensibilità locale di Peri-LN è mostrata essere invariante rispetto alla magnitudo dell'attivazione, prevenendo l'esplosione del gradiente anche quando le attivazioni sono grandi.

3. Quantificazione dell'Incertezza

L'analisi stabilisce che Peri-LN fornisce una propagazione controllata dell'incertezza (misurata tramite la distanza di Wasserstein) dall'input agli stati nascosti terminali, mentre Pre-LN può amplificare l'incertezza dell'input fino a livelli illimitati.

4. Scaling del Passo Residuo

Guidati dall'analisi della stabilità, gli autori introducono un meccanismo di scaling del passo residuo ($\Delta t < 1$). Dimostrano teoricamente che questo scaling:

• Affina i limiti sulla crescita degli stati nascosti e sull'incertezza in uscita.
• Mitiga l'esplosione del gradiente scalando la sensibilità locale.
• Migliora sia la stabilità forward che quella backward senza costi computazionali aggiuntivi.

Risultati Sperimentali

Gli autori validano i loro risultati teorici utilizzando varianti di GPT-2 (da 124M a 1.5B parametri) addestrate sul dataset OpenWebText.

• Stabilità dell'Addestramento: In esperimenti con 5 prove, i modelli Pre-LN sono divergi in 1 run quando il weight decay era abilitato e in 3 run quando il weight decay era disabilitato. Al contrario, i modelli Peri-LN sono rimasti stabili in tutte le 5 prove sotto entrambe le impostazioni di weight decay.
• Performance: Peri-LN ha raggiunto prestazioni comparabili a Pre-LN anche quando gli iperparametri erano ottimizzati per Pre-LN.
• Impatto dello Scaling: Combinare Peri-LN con lo scaling del residuo ($\Delta t = 0.1$) ha costantemente prodotto le migliori prestazioni in tutti i parametri (Val Loss, Perplexity, ROUGE, punteggi BERT) e ha controllato efficacemente la crescita dei momenti degli stati nascosti (valore assoluto medio e varianza) attraverso gli strati, allineandosi alle previsioni teoriche.

Significato e Rivendicazioni

Il documento sostiene di fornire la prima rigorosa giustificazione teorica per le osservazioni empiriche riguardanti i posizionamenti della Normalizzazione di Livello nei Transformer.

• Spiegazione dell'Instabilità: Spiega perché Pre-LN porta a stati nascosti illimitati ed esplosione del gradiente tramite l'illimitata natura della PDE HJB sottostante.
• Giustificazione per Peri-LN: Fornisce una base teorica per la stabilità di Peri-LN, mostrando che mantiene una crescita controllata e una sensibilità invariante.
• Guida al Design: Il framework offre un metodo principato per "verificare la sanità" (sanity-check) di nuove modifiche architettoniche prima di costosi addestramenti empirici.
• Miglioramento Pratico: L'introduzione dello scaling del passo residuo è presentata come una modifica semplice e priva di costi che migliora significativamente stabilità e prestazioni, validata da esperimenti che mostrano come Peri-LN con scaling superi le configurazioni standard.

Gli autori sottolineano che il loro lavoro colma il divario tra euristiche empiriche e comprensione teorica, offrendo una prospettiva continua che genera sia intuizioni teoriche che miglioramenti pratici all'architettura.