Feed-Forward Steering in Transformer Residual Dynamics

Sintesi Tecnica: Guida Feed-Forward nella Dinamica Residua dei Transformer

Problema
Le attuali teorie continue della dinamica dei Transformer modellano principalmente lo stream residuo come un sistema di particelle che si aggregano su una sfera, guidate esclusivamente dal meccanismo di attenzione. Sebbene questo framework "solo-attenzione" spieghi fenomeni come il collasso del rango e il clustering, non riesce a rendere conto del ruolo della rete Feed-Forward (FFN). Parallelamente, la ricerca sull'interpretabilità identifica le FFN come siti di archiviazione per memorie chiave-valore e target efficaci per l'editing, eppure queste visioni rimangono scollegate dalla prospettiva dei sistemi dinamici. Il problema centrale affrontato è come integrare la FFN in un modello dinamico unificato dell'evoluzione dello stream residuo per spiegare il suo impatto geometrico sulle direzioni dei token e per determinare la fattibilità di approssimare blocchi sequenziali con blocchi paralleli.

Metodologia
Gli autori estendono il framework continuo della dinamica dell'attenzione incorporando la FFN come un campo di guida locale che agisce sugli stati dei token. Il modello tratta lo stream residuo $x_i = r_i u_i$ (dove $u_i$ è la direzione sulla sfera unitaria e $r_i$ è l'ampiezza) come un'evoluzione secondo:
$$\dot{u}_i = \frac{1}{s_i(t)} P_{u_i} \left[ A_i(U) + \Phi_i(u_i) \right]$$
Qui, $A_i(U)$ rappresenta il campo di aggregazione non locale proveniente dall'attenzione, e $\Phi_i(u_i)$ è il campo di guida locale fornito dalla FFN. $P_{u_i}$ è la proiezione ortogonale sullo spazio tangente, garantendo che solo la componente tangenziale del campo vettoriale alteri la direzione, mentre la componente radiale influenzi l'ampiezza (e quindi la velocità angolare).

Lo studio impiega quattro protocolli sperimentali primari su modelli GPT-2, Pythia, Mistral e Llama:

1. Diagnostica del Deficit Angolare: Misurazione dell'allineamento tra l'aggiornamento basato solo sull'attenzione e l'aggiornamento del blocco realizzato per quantificare il contributo della FFN al moto angolare.
2. Ablazione Radiale/Tangenziale: Decomposizione dell'output della FFN nelle sue componenti radiale ($\langle u, \Phi(u) \rangle u$) e tangenziale ($P_u \Phi(u)$) per isolare quale guidi i cambiamenti di direzione.
3. Diversità sotto Aggregazione: Analisi delle statistiche di diversità sferica prima e dopo i layer di attenzione e FFN per osservare come la FFN contrasti la concentrazione indotta dall'attenzione.
4. Difetto Sequenziale-Parallelo: Stima di un punteggio di sensibilità per layer ($C_\ell$) per determinare se i blocchi sequenziali ($x + A(x) + \Phi(x+A(x))$) possano essere approssimati da surrogati paralleli ($x + A(x) + \Phi(x)$) senza perdite significative di performance.

Ulteriori diagnostiche includono la calibrazione del gate della FFN per layer, l'analisi della geometria dei punti critici (identificazione di attrattori, selle e repeller) e test di sensibilità al prefisso.

Contributi Chiave

1. Framework di Aggregazione-Guida: Il documento formula la dinamica dello stream residuo dei Transformer come un flusso di aggregazione-guida su una sfera, dove l'attenzione fornisce l'aggregazione non locale e la FFN fornisce la guida locale.
2. Identificazione della Guida Tangenziale: Gli autori dimostrano matematicamente ed empiricamente che la componente tangenziale del campo FFN è necessaria per il moto nello spazio della direzione residua. La componente radiale riscala principalmente l'ampiezza, mentre la componente tangenziale detta la direzione.
3. Geometria delle Direzioni Critiche: Lo studio identifica che le direzioni residue critiche corrispondono a equilibri proiettivi non lineari. Contrariamente all'intuizione che le FFN creino attrattori stabili, l'analisi rivela che la guida della FFN riorganizza attivamente la geometria locale verso modalità di trasporto rotazionali e di tipo sella, preservando la diversità.
4. Metrica del Difetto Sequenziale-Parallelo: Viene derivata una metrica di difetto diretta per caratterizzare quando un blocco sequenziale attenzione-FFN può essere accuratamente approssimato da una mappa parallela e additiva. Questo difetto è governato dal termine della derivata direzionale $D_\Phi A$, distinto dal solo bracket di Lie.

Risultati

• Disallineamento Angolare: Il ramo dell'attenzione da solo presenta un sostanziale deficit di allineamento angolare rispetto all'aggiornamento del blocco realizzato. Questo deficit aumenta con la scala del modello, passando da 0,41 in GPT-2 a 0,63 in Llama-3-8B, indicando che la FFF non è una mera correzione secondaria ma un driver primario degli aggiornamenti angolari.
• Risultati dell'Ablazione: Mantenere solo la componente tangenziale della FFN preserva la maggior parte della qualità del modello (incrementi di perdita tra 0,07 e 1,72), mentre mantenere solo la componente radiale causa il collasso delle performance (incrementi di perdita tra 7,19 e 12,96), paragonabile alla rimozione totale della FFN.
• Preservazione della Diversità: L'attenzione tipicamente contrae le direzioni residue (riducendo la diversità), mentre la FFN spesso ripristina o rimodella questa diversità. La componente tangenziale della FFN inietta modalità di trasporto rotazionali e di tipo sella che prevengono l'eccessiva concentrazione sotto la pressione di aggregazione.
• Fattibilità della Parallelizzazione: I layer con piccoli difetti sequenziale-parallelo (bassi punteggi di sensibilità) possono essere approssimativamente parallelizzati con modesti incrementi di perdita. Ad esempio, parallelizzare 18 layer a basso difetto in GPT-2-large produce un'accelerazione stimata di 1,24× con solo +0,02 di perdita, mentre parallelizzare i layer ad alto difetto degrada rapidamente le prestazioni.
• Geometria Critica: L'ablazione delle componenti del campo delle direzioni critiche mostra che i candidati FFN-only sono prevalentemente di tipo sella (spettri a segno misto). La FFN non crea semplicemente attrattori stabili, ma ristruttura il ritratto di fase per mantenere una geometria di rappresentazione ad alta dimensionalità.

Significatività e Affermazioni
Il documento sostiene l'interpretazione delle FFN come campi di guida direzionale che modellano la geometria residua dei Transformer. Argomenta che la FFF è essenziale per governare la fattibilità delle interazioni a livello di blocco e per mantenere la diversità delle rappresentazioni contro la pressione di collasso dell'attenzione.

Gli autori pongono i loro risultati come un ponte tra la teoria dei sistemi dinamici e l'interpretabilità, suggerendo che la visione di "guida" spiega perché le FFN siano siti di editing efficaci e perché siano necessarie per la stabilità geometrica del modello. La significatività pratica risiede nell'identificazione dei layer che possono essere parallelizzati in sicurezza per ridurre la latenza di inferenza senza degradare significativamente le prestazioni.

Il lavoro rimane modesto riguardo al proprio ambito, notando che i risultati di media di campo descrivono l'attenzione completa, mentre i modelli decoder-only causali richiedono dinamiche indicizzate al prefisso. Le analisi dei punti critici sono presentate come diagnostiche locali piuttosto che descrizioni globali dell'intero spazio residuo. Inoltre, i risultati pratici sulla parallelizzazione sono inquadrati come accelerazioni stimate del tempo dei sottomoduli piuttosto che come affermazioni di ottimizzazione di kernel di produzione. Gli autori dichiarano esplicitamente che le loro scoperte sono diagnostiche e non rivendicano garanzie universali di anti-collasso o di parallelizzazione perfetta in tutte le architetture.