Sintesi Tecnica: Rivedere la Convenzione di Forma dei Modelli Linguistici Transformer

Definizione del Problema

I modelli linguistici Transformer densi hanno storicamente aderito a una rigida convenzione architettonica: la Rete Feed-Forward (FFN) utilizza una forma MLP (Perceptron Multistrato) "stretta–larga–stretta". In questo design, la dimensione del modello ($d_{model}$) viene espansa in una dimensione intermedia più larga ($d_h$), con un rapporto di espansione tipicamente compreso tra 2 e 4, prima di essere proiettata nuovamente. Questa struttura domina l'allocazione dei parametri, spesso rappresentando la maggior parte dei parametri e delle operazioni in virgola mobile del modello.

Gli autori mettono in discussione se questa specifica forma sia essenziale o solo un'istanza conveniente. Recenti lavori teorici suggeriscono che gli MLP residui "larga–stretta–larga" (a clessidra o hourglass) possono fungere da approssimatori di funzioni ottimali, offrendo potenzialmente una capacità espressiva superiore con meno parametri. Inoltre, prove empiriche in modelli di visione generativa e Transformer Mixture-of-Experts (MoE) hanno dimostrato che le strutture a clessidra possono superare i design standard. Il problema centrale affrontato è se la convenzionale FFN "stretta–larga–stretta" sia un vincolo necessario o se una FFN a clessidra più profonda e stretta possa offrire un'allocazione più efficiente della capacità tra i moduli di attenzione e FFN.

Metodologia

Il paper propone una variante Transformer che sostituisce la convenzionale FFN con una FFN a Clessidra (Hourglass FFN). La metodologia prevede:

1. Riprogettazione Architetturale:

   • La FFN a Clessidra sostituisce l'unico strato di espansione larga con una pila di $K$ sub-MLP a forma di clessidra collegati da percorsi residui.
   • Ogni sub-MLP ($MLP_i$) segue una struttura di compressione-espansione: proietta l'input dalla dimensione $d_{model}$ verso una dimensione di collo di bottiglia $d_h$ (dove $d_h < d_{model}$), applica un'attivazione non lineare (SwiGLU) e proietta nuovamente verso $d_{model}$.
   • L'output dello strato è il risultato del passaggio dell'input attraverso $K$ tali sub-blocchi residui.

2. Strategia di Redistribuzione dei Parametri:

   • Riducendo la dimensione intermedia $d_h$ (rendendo la FFN più "stretta"), gli autori risparmiano parametri che tipicamente vengono allocati alla larghezza della FFN.
   • Questi parametri risparmiati vengono riallocati a:
     • Aumentare la dimensione del modello ($d_{model}$) e la capacità di attenzione.
     • Aumentare la profondità dello stack FFN ($K$).
     • Aumentare il numero di strati Transformer ($L$).

3. Configurazione Sperimentale:

   • Baseline: I confronti sono effettuati contro Transformer in stile LLaMA e l'architettura allo stato dell'arte OLMo-2.
   • Scale: Gli esperimenti sono condotti su dimensioni di modello che vanno da 113M a 1B di parametri.
   • Vincoli: Gli esperimenti sono eseguiti sotto budget totali di parametri corrispondenti (entro uno 0,001% di differenza) e dimensioni di attenzione fisse per isolare gli effetti della forma della FFN.
   • Spazio di Ricerca: Gli autori eseguono una ricerca architettonica variando $K$ (profondità dello stack a clessidra), $L$ (numero di strati Transformer), $d_{model}$ e il rapporto $d_h/d_{model}$.

Contributi Principali

Il paper apporta quattro contributi primari:

1. Sfida allo Standard: Mette in discussione la de facto forma FFN "stretta–larga–stretta", proponendo una efficiente FFN a Clessidra "larga–stretta–larga" che mantiene l'espressività attraverso sub-MLP impilati e connessi in modo residuo.
2. Riallocazione dei Parametri: Dimostra una strategia in cui il reindirizzamento dei parametri dalla larghezza della FFN alla larghezza dell'attenzione o alla profondità interna della FFN ($K > 1$) produce prestazioni superiori, in particolare a scale minori (113M).
3. Intuizioni Architetturali: Identifica un compromesso (trade-off) tra larghezza e profondità a forma di U, trovando rapporti ottimali per $d_{model}/L$ e $d_h/d_{model}$, e stabilisce che le strutture FFN a clessidra più profonde sono robuste.
4. Validazione della Scalabilità: Valida la FFN a Clessidra da 113M a 1B di parametri, dimostrando che è un'alternativa competitiva ai baseline standard, con configurazioni specifiche che superano i modelli convenzionali a scale minori e li eguagliano a scale maggiori.

Risultati

La validazione empirica produce le seguenti scoperte:

• Prestazioni a Attenzione Fissa: Alla scala di 113M con parametri di attenzione fissi, le FFN a Clessidra raggiungono una perdita di validazione e una perplessità (PPL) inferiori rispetto alle FFN convenzionali. Ad esempio, con $K=8$, la PPL è scesa da 36,441 (convenzionale) a 36,179.
• Configurazioni Ottimali:
  • Profondità vs Larghezza: Un rapporto $d_h/d_{model} \approx 0,4$ con $K=4$ strati ha prodotto la perdita di validazione più bassa.
  • Spostamento dei Parametri: A 113M, spostare i parametri per rendere dominante il modulo di attenzione (ad esempio, 66M attenzione vs 47M FFN) ha portato alla migliore prestazione (PPL 35,101), suggerendo che l'elaborazione contestuale è più preziosa del recupero fattuale statico per i modelli piccoli.
  • Compromesso Larghezza-Profondità: È stata osservata una curva distinta a forma di U per il rapporto $d_{model}/L$, con un punto ottimale tra 100 e 250.
• Scalabilità:
  • Scale Piccole e Medie (113M – 906M): Le varianti a Clessidra hanno costantemente superato i baseline convenzionali in termini di perdita di validazione, perplessità e compiti di ragionamento downstream (es. Arc Easy, HellaSwag, PIQA).
  • Grande Scala (1B): Alla scala di 1B di parametri, la variante a Clessidra ha raggiunto prestazioni comparabili al forte baseline convenzionale (Val PPL 20,082 vs 20,002) e ha superato la configurazione OLMo-2 (20,466).
• Studi di Ablazione: L'aumento della profondità della clessidra $K$ ha costantemente migliorato le prestazioni, sebbene con rendimenti decrescenti oltre $K=6$. Ridurre $d_h/d_{model}$ al di sotto di 0,4 ha portato a un significativo degrado delle prestazioni, indicando un limite inferiore per la compressione della FFN.

Significatività e Rivendicazioni

Il paper sostiene che la forma MLP "stretta–larga–stretta" a lungo termine nei Transformer non è essenziale e potrebbe essere ridondante. Adottando una FFN a Clessidra, gli autori dimostrano che:

1. Efficienza: I modelli possono raggiungere prestazioni competitive o superiori con una distribuzione dei parametri più efficiente.
2. Riallocazione delle Risorse: La struttura a Clessidra consente un cambiamento fondamentale in cui il meccanismo di attenzione può diventare il consumatore dominante di parametri, sfidando il tradizionale budget dominato dalla FFN.
3. Rivalutazione del Design: Le scoperte spingono a riconsiderare l'equilibrio tra attenzione e FFN, suggerendo che per i modelli linguistici moderni, efficienti ed espressivi, una FFN più profonda e stretta combinata con meccanismi di attenzione più larghi possa essere una strategia di design superiore.

Gli autori rimangono modesti riguardo alla scalabilità verso modelli con decine o centinaia di miliardi di parametri, osservando che le loro scoperte sono state estrapolate dalla scala 113M e che la "maledizione della profondità" (problemi di propagazione del segnale) richiede ulteriori indagini in architetture più profonde. Notano inoltre che l'interazione tra le FFN a Clessidra e i meccanismi di attenzione avanzati (es. Group-Query Attention) rimane inesplorata.