Attention Residuals spiegati: ripensare la profondità dei Transformer

I moderni LLM sono diventati più profondi, più ampi e più esigenti in termini di calcolo, ma impilare più layer Transformer non porta sempre a guadagni proporzionali. Una ragione è che le connessioni residuali standard aggregano le uscite dei layer con pesi unitari fissi, quindi ogni layer eredita una somma uniforme di tutto ciò che lo ha preceduto. Questo può diluire le rappresentazioni iniziali, amplificare la magnitudine degli stati nascosti e rendere più difficile per la rete il riuso selettivo delle caratteristiche intermedie più utili.

Invece di trattare la profondità come una ricorrenza additiva fissa, il paper del team Kimi sugli Attention Residuals consente a ciascun layer di effettuare attenzione sulle uscite dei layer precedenti usando pesi softmax appresi. In questo articolo capiremo perché l’aggregazione residua standard diventa un collo di bottiglia, come funzionano gli Attention Residuals, perché la variante a blocchi è importante e cosa suggeriscono davvero i risultati sullo scaling di modelli linguistici più profondi.

Se vuoi approfondire alcune idee dietro le architetture Transformer moderne, ti consiglio il corso DataCamp Tutorial sui modelli Transformer in PyTorch.

Il problema: il segnale che affoga

Le connessioni residuali sono fondamentali per il deep learning da anni. Nei Transformer, l’aggiornamento residuo standard è:

hl=hl−1​+fl−1​(hl−1​)

Questo aiuta il flusso dei gradienti attraverso reti molto profonde. Ma i residuali non sono solo un trucco per i gradienti. Definiscono anche come l’informazione viene aggregata in profondità. Se srotoliamo la ricorrenza, otteniamo:

hl​=h1​+∑i=1l−1​fi​(hi​)

Questo significa che lo stato nascosto al layer l è solo l’embedding ed è una somma pesata uniformemente di tutte le uscite dei layer precedenti. Quindi, a ogni contributo viene assegnato di fatto lo stesso peso.

Questo diventa un problema su larga scala. Il paper mostra che nelle architetture PreNorm, l’accumulo residuo non pesato fa crescere la magnitudine dello stato nascosto con la profondità, approssimativamente come O(L). 

Con la crescita del flusso residuo, le uscite dei layer precedenti diventano sempre più diluite dentro una somma progressiva più ampia. Una volta che un segnale di un layer iniziale è mescolato in questo accumulo, i layer più profondi non possono recuperarlo selettivamente. Di conseguenza, operano solo sullo stato aggregato.

Questo porta a quello che il paper chiama effetto del “segnale che affoga”. Una forte evidenza empirica di questa inefficienza viene dagli studi di pruning dei layer, cioè una frazione sostanziale di layer nei modelli addestrati può spesso essere rimossa con un impatto minimo sulle prestazioni. Questo suggerisce che, pur continuando ad aumentare la profondità, ai modelli manca un meccanismo efficace per sfruttarla appieno. Invece di formare una catena di ragionamento gerarchica, la rete si comporta più come un relè ridondante in cui i segnali iniziali vengono progressivamente diluiti.

C’è anche un compromesso nel forward pass. Con la crescita del flusso residuo, i layer successivi possono dover produrre uscite di magnitudine maggiore per influenzare in modo significativo lo stato accumulato. Il paper collega questo al comportamento PreNorm, in cui le magnitudini degli stati nascosti aumentano monotonamente con la profondità sotto l’aggregazione residua standard.

Le connessioni residuali aiutano il flusso dei gradienti, e il gradiente rispetto a uno stato nascosto intermedio può essere scritto come:

Il termine identità qui preserva un percorso di gradiente diretto. Ma i residuali costringono comunque il percorso di aggregazione in forward a trattare ogni layer precedente con un peso fisso di 1,0. Questa è la limitazione strutturale che AttnRes cerca di risolvere.

Dai residuali additivi all’attenzione lungo la profondità

Il passaggio concettuale più interessante del paper è la dualità tempo–profondità. Le connessioni residuali comprimono l’informazione lungo la profondità, come le reti ricorrenti la comprimono nel tempo. Nel modeling di sequenze, l’attenzione ha sostituito la ricorrenza permettendo a ogni posizione di accedere selettivamente alle posizioni precedenti. AttnRes applica la stessa transizione alla profondità di rete. Invece di definire il prossimo stato nascosto come una somma fissa sui layer precedenti, AttnRes consente a ogni layer di effettuare attenzione sulle uscite dei layer passati:

dove i pesi αi→l sono pesi di attenzione softmax lungo la profondità e sommano a 1. Questi pesi sono calcolati come:

con il paper che utilizza:

Ogni layer ottiene un vettore pseudo-query appreso wl, e quella query fa attenzione su chiavi e valori costruiti dalle uscite dei layer precedenti. RMSNorm è applicata alle chiavi affinché i layer con uscite naturalmente di magnitudine maggiore non dominino la softmax solo per la loro scala.

Il cambiamento chiave nel paper è che i residuali standard sono trattati come una sorta di attenzione lineare lungo la profondità, mentre AttnRes la eleva a un’attenzione softmax lungo la profondità. Invece di un accumulo uniforme, otteniamo un recupero selettivo lungo la profondità.

Un piccolo ma importante dettaglio di implementazione è l’inizializzazione. Gli autori suggeriscono che tutti i vettori pseudo-query dovrebbero essere inizializzati a zero, il che assicura che i pesi di attenzione iniziali siano uniformi, così il modello inizia l’addestramento come una media a pesi uguali invece di un meccanismo di attenzione casualmente sbilanciato.

Figura 1: Panoramica degli Attention Residuals (paper sugli Attention Residuals)

Full Attention Residuals

Nei Full Attention Residuals, ogni layer fa attenzione su tutte le uscite dei layer precedenti, consentendogli di aggregare selettivamente informazioni attraverso la profondità. Questo dà al modello la massima flessibilità, così un layer più profondo può enfatizzare il suo immediato predecessore, l’embedding originale o qualsiasi layer precedente se è lì che risiede il segnale più utile.

Nel training standard, l’overhead di memoria dei Full AttnRes è minore di quanto possa sembrare, perché molte uscite dei layer sono già mantenute per la backpropagation. Ma l’addestramento su larga scala cambia il quadro. Una volta introdotte la ri-computazione delle attivazioni e il pipeline parallelism, quelle uscite precedenti devono essere esplicitamente preservate e comunicate affinché i layer successivi possano farvi attenzione, e questo diventa costoso.

Quindi i Full AttnRes sono il modo migliore per comprendere l’idea di base, ma non è la versione che la maggior parte dei team vorrebbe distribuire su scala.

Block Attention Residuals

Il paper introduce l’idea dei Block Attention Residuals per ridurre l’overhead di memoria e comunicazione. I layer del modello sono partizionati in blocchi tali che all’interno di ciascun blocco le uscite siano combinate usando l’accumulo residuo additivo standard, mentre tra i blocchi il modello esegue attenzione su riassunti a livello di blocco invece che su ogni singolo layer precedente.

Se Bn è l’insieme dei layer nel blocco n, allora la rappresentazione del blocco è data da:

Il modello quindi fa attenzione sull’embedding b0=h1 riassunti dei blocchi precedenti, e sulla somma parziale del blocco corrente man mano che il calcolo procede. Questo riduce memoria e comunicazione da O(Ld) a O(Nd), dove N è il numero di blocchi.

Il paper riporta che usando circa otto blocchi si recupera la maggior parte del beneficio della versione completa, e il divario di prestazioni tra Full AttnRes e Block AttnRes si riduce con l’aumentare della scala. Dimensioni di blocco come S=2,4,8 restano tutte vicine alla versione completa, mentre raggruppamenti molto più grossolani tendono a tornare verso il comportamento di base.

L’insegnamento più utile del paper è che non serve un’attenzione completa lungo la profondità su ogni layer per ottenere gran parte del guadagno. Mentre Block AttnRes rende l’attenzione lungo la profondità fattibile dal punto di vista computazionale, distribuirla in modo efficiente su scala richiede comunque un’attenta progettazione dei sistemi. Il paper introduce anche diverse ottimizzazioni che rendono AttnRes pratico per training e inferenza nel mondo reale.

Figura 2: Esempio di comunicazione in pipeline basata su cache con 4 rank fisici e 2 stage virtuali per rank, dove i riquadri tratteggiati indicano la fine dei blocchi AttnRes (paper sugli Attention Residuals)

Ecco tre ottimizzazioni di sistema chiave che rendono AttnRes un sostituto plug-and-play dei residuali standard:

• Computazione a due fasi: disaccoppiando la pseudo-query dallo stato nascosto, possiamo batchare le query per tutti i layer all’interno di un blocco. La Fase 1 calcola in parallelo l’attenzione inter-blocco, mentre la Fase 2 gestisce le dipendenze sequenziali intra-blocco usando un Online Softmax Merge. Questo riduce significativamente l’I/O per layer rispetto a metodi multi-stream precedenti come mHC.
• Caching tra stage: nel pipeline parallelism, eliminiamo trasferimenti di dati ridondanti memorizzando in cache localmente i blocchi ricevuti in precedenza. Questo fornisce un miglioramento di V× (dove V è il numero di stage virtuali come mostrato in Figura 2) nel costo massimo per transizione, consentendo alla comunicazione di sovrapporsi completamente al calcolo.
• Prefilling efficiente in memoria: suddividendo in shard le rappresentazioni dei blocchi tra dispositivi in tensor parallel, l’impronta di memoria per sequenze a lungo contesto (128K+) è ridotta di ordini di grandezza (ad es., da 15GB a 0,3GB per dispositivo).

Risultati e analisi delle prestazioni

Il paper convalida AttnRes a più livelli, come leggi di scala, ablation, dinamiche di training e benchmark a valle.

Figura 3: curve delle leggi di scala per gli Attention Residuals. (paper sugli Attention Residuals)

Tuttavia, il risultato principale emerge dagli esperimenti sulle leggi di scala. Su cinque dimensioni di modello, sia Full AttnRes sia Block AttnRes ottengono costantemente una validation loss inferiore rispetto al baseline. In base alle curve adattate, Block AttnRes raggiunge la stessa loss di un baseline addestrato con circa 1,25× più compute. Questo suggerisce che AttnRes non è solo teoricamente migliore, ma è anche più efficiente dal punto di vista computazionale.

Nei task a valle, AttnRes migliora rispetto al baseline su tutti i benchmark valutati. Alcuni dei guadagni maggiori sono sui task ad alta intensità di ragionamento, come mostrato nella Tabella 1:

Tabella 1: confronto delle prestazioni di AttnRes con il baseline, entrambi dopo la stessa ricetta di pre-training (paper sugli Attention Residuals)

I guadagni sono particolarmente forti nel ragionamento multi-step, nella matematica e nel codice, in linea con l’ipotesi centrale del paper secondo cui, se i layer successivi possono recuperare selettivamente rappresentazioni precedenti invece di ereditare un aggregato sfocato, il ragionamento composizionale dovrebbe migliorare.

Rispetto al baseline, AttnRes mostra una validation loss inferiore durante tutto il training, magnitudini di output più contenute lungo la profondità e magnitudini di gradiente più uniformi tra i layer. Il baseline soffre del solito pattern di diluizione PreNorm, in cui le magnitudini degli stati nascosti crescono monotonamente con la profondità, mentre Block AttnRes produce un pattern più controllato e limitato grazie all’aggregazione selettiva ai confini dei blocchi.

C’è anche un risultato più ampio a livello architetturale: in uno sweep a compute e parametri fissi, la configurazione ottimale per AttnRes si sposta verso un modello più profondo e più stretto rispetto al baseline. Questo suggerisce che AttnRes può rendere ulteriore profondità più utile di quanto non lo sia sotto l’aggregazione residua standard.

Conclusione

I residuali sono spesso trattati come un trucco necessario per il training, cioè una scorciatoia identitaria che mantiene vivi i gradienti. AttnRes suggerisce che questa visione è troppo ristretta. I percorsi residuali sono anche il meccanismo con cui l’informazione viene instradata attraverso la profondità, e un accumulo additivo fisso può semplicemente essere troppo primitivo per modelli molto profondi.

Gli Attention Residuals non servono tanto a mettere una pezza a un problema noto quanto ad aggiornare una scelta progettuale trascurata. Il sequence modeling è passato dalla ricorrenza all’attenzione perché la ricorrenza fissa era troppo restrittiva. AttnRes sostiene che anche la profondità potrebbe essere pronta per la stessa transizione.

Il modello mostra comunque una forte località, con i layer che spesso pongono più attenzione ai predecessori vicini, ma apprende anche pattern di salto non banali, preserva un peso persistente sull’embedding e mantiene un comportamento distinto tra i layer pre-attention e pre-MLP. 

La considerazione pratica non è che tutti i modelli debbano sostituire i residuali con un’attenzione completa lungo la profondità. Piuttosto, riformula l’aggregazione residua come uno spazio di progettazione architetturale centrale.