💻 computer science

Multi-Head Attention Residuals

Questo articolo introduce i Multi-Head Attention Residuals (MHAR), una modifica architettonica priva di parametri che consente a diversi sottospazi di caratteristiche di prestare attenzione indipendentemente alla cronologia della profondità attraverso molteplici teste di routing, risolvendo così il disaccordo tra sottospazi e migliorando significativamente la perdita di validazione e le prestazioni nei compiti a valle rispetto ai Transformer standard in varie scale di modelli.

Autori originali: Cheng Luo, Zefan Cai, Junjie Hu

Pubblicato 2026-07-31

📖 4 min di lettura☕ Lettura da pausa caffè

CC BY 4.0

Autori originali: Cheng Luo, Zefan Cai, Junjie Hu

Articolo originale sotto licenza CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Questa è una spiegazione generata dall'IA dell'articolo qui sotto. Non è stata scritta né approvata dagli autori. Per precisione tecnica, consulta l'articolo originale. Leggi il disclaimer completo

Immagina di stare costruendo una biblioteca massiccia e multipiano, dove ogni nuovo piano aggiunge una nuova stanza di conoscenza. Nel mondo dell'Intelligenza Artificiale, queste "biblioteche" sono chiamate Transformer, e sono i cervelli dietro molti degli strumenti intelligenti che usiamo oggi. Per dare un senso a una storia o a un problema di matematica, la biblioteca deve ricordare cosa è successo al primo piano mentre sta lavorando al decimo. Tradizionalmente, queste biblioteche hanno un unico, stretto corridoio che collega ogni piano. Quando una stanza al decimo piano ha bisogno di guardare indietro, può solo sbirciare nella stanza immediatamente sottostante. Deve fidarsi del fatto che la stanza sotto abbia riassunto perfettamente tutto dal primo piano in su fino a quel punto. È come cercare di ricordare un intero film chiedendo solo alla persona seduta proprio davanti a te cosa è successo; se quella persona ha perso un dettaglio, lo perdi anche tu.

Per risolvere questo problema, i ricercatori hanno recentemente provato una nuova idea chiamata "Attention Residuals" (Residui di Attenzione). Invece di guardare solo la stanza sottostante, la biblioteca ha dotato ogni stanza di un telescopio magico che poteva guardare qualsiasi piano precedente. Tuttavia, c'era un intoppo: il telescopio aveva un solo obiettivo. Ogni singolo dettaglio nella stanza — ogni colore, ogni suono, ogni numero — doveva guardare attraverso quell'unico obiettivo condiviso per decidere su quale piano concentrarsi. Era come costringere un gruppo di amici con gusti totalmente diversi (uno ama i film d'azione, un altro i documentari, un terzo i cartoni animati) a mettersi tutti d'accordo su esattamente un unico canale da guardare insieme. Dovevano scendere a compromessi, e spesso, quel compromesso significava che nessuno riusciva a vedere ciò di cui aveva realmente bisogno.

Questo articolo introduce un aggiornamento intelligente chiamato Multi-Head Attention Residuals (MHAR). I ricercatori si sono resi conto che, proprio come un cinema ha più schermi per mostrare generi diversi contemporaneamente, il telescopio della biblioteca dovrebbe avere più lenti. Hanno diviso quell'unico obiettivo condiviso in diverse lenti più piccole e indipendenti. Ora, la parte della stanza che ama i "film d'azione" può guardare i piani pieni di azione, mentre la parte dei "documentari" guarda i piani storici, il tutto senza litigare o scendere a compromessi.

Il team ha testato questa idea costruendo biblioteche di diverse dimensioni, da modelli piccoli da 100 milioni di parametri fino a uno massiccio da 1 miliardo di parametri. Hanno scoperto che, per le biblioteche piccole, il vecchio sistema a obiettivo singolo funzionava bene. Ma man mano che le biblioche diventavano più grandi e complesse, l'obiettivo singolo diventava un collo di bottiglia, rendendo in realtà l'IA peggiore rispetto al sistema del corridoio standard. Il nuovo sistema a più lenti, invece, migliorava sempre di più man mano che le biblioteche crescevano. Alla scala più grande, il nuovo metodo ha ridotto la confusione dell'IA (misurata come perdita di validazione) di circa 0,063 rispetto al modello standard, mentre il vecchio metodo a obiettivo singolo ha effettivamente peggiorato le cose di 0,140.

I ricercatori hanno anche dimostrato che questo aggiornamento non richiedeva la costruzione di un telescopio più grande o l'aggiunta di ulteriore peso alla biblioteca; era solo un modo più intelligente di disporre il vetro esistente. Hanno persino costruito strumenti speciali, super veloci, per garantire che la biblioteca non venisse rallentata dall'extra attività di osservazione, mantenendo la velocità quasi identica a quella del design originale. Infine, hanno dimostrato che questo trucco funziona anche su una gigantesca biblioteca pre-costruita da 8 miliardi di parametri, potenziando la sua capacità di risolvere problemi matematici e rispondere a difficili domande scientifiche senza dover essere ricostruita da zero. La scoperta fondamentale è che man mano che i modelli di IA diventano più ampi e complessi, smettono di concordare su cosa ricordare, e dare loro molteplici modi indipendenti di guardare indietro è la chiave per sbloccare tutto il loro potenziale.

Riepilogo Tecnico: Multi-Head Attention Residuals (MHAR)

1. Definizione del Problema

I Transformer standard propagano l'informazione attraverso la profondità tramite un singolo flusso residuo additivo, dove ogni sottostrato legge solo lo stato immediatamente precedente ( $h_\ell = h_{\ell-1} + f_\ell(h_{\ell-1})$ ). Sebbene le caratteristiche utili calcolate negli strati precedenti rimangano nel totale corrente, esse diventano individualmente non indirizzabili.

Gli Attention Residuals (Kimi, 2025) rilassano questo vincolo trattando tutti i prodotti dei sottostrati precedenti come una memoria, permettendo a ogni sottostrato di recuperare una combinazione appresa e pesata tramite softmax dell'intera cronologia della profondità. Tuttavia, questo meccanismo soffre di un collo di bottiglia strutturale: utilizza un singolo query di routing condiviso ( $q \in \mathbb{R}^d$ ) per valutare la cronologia della profondità. Di conseguenza, tutte le $d$ coordinate delle feature sono costrette a leggere dagli stessi strati nelle stesse proporzioni attraverso una singola distribuzione $\alpha$ .

Il paper sostiene che ciò crei un "compromesso forzato". All'aumentare della larghezza del modello, diverse sottospazi di feature probabilmente svilupperanno preferenze divergenti riguardo a quali strati storici siano più rilevanti. Un singolo query non può soddisfare simultaneamente queste diverse preferenze, facendo collassare distinte preferenze di profondità in una singola lettura. Il costo di questo compromesso cresce con la larghezza del modello, rendendo potenzialmente dannoso il routing a singolo query su larga scala.

2. Metodologia: Multi-Head Attention Residuals (MHAR)

Gli autori propongono i Multi-Head Attention Residuals (MHAR), una generalizzazione rigorosa degli attention residuals che elimina il compromesso forzato senza aggiungere parametri o un calcolo significativo.

Meccanismo Core

Invece di un singolo query $q$ , MHAR rimodella il query di routing in $H$ teste per sottospazio, ciascuna con la propria indipendente softmax sulla cronologia della profondità.

Rimodellamento (Reshaping): Il singolo vettore query $q \in \mathbb{R}^d$ viene rimodellato in un tensore di forma $(H, d/H)$ .
Routing Indipendente: Per ogni testa $h$ , il modello calcola una distinta distribuzione di profondità $\alpha^{(h)}$ basata sulla porzione del vettore sorgente corrispondente a quella testa.
Lettura a Diagonale di Blocco: La miscela recuperata è a diagonale di blocco; ogni sottospazio legge la cronologia della profondità in modo indipendente.
Recupero: Impostando $H=1$ , si recuperano esattamente gli originali attention residuals a singola testa.

Dettagli di Implementazione

Efficienza dei Parametri: La trasformazione è un reshape, il che aggiunge zero parametri rispetto alla baseline a singola testa.
Efficienza Computazionale: L'overhead è trascurabile, comportando solo $H$ piccole softmax sull'asse della profondità (lunghezza $\le 2L+1$ ) invece di una grande softmax.
Ottimizzazione del Sistema: Riconoscendo che il routing della profondità è limitato dalla memoria (richiede traffico $O(L^2 B T d)$ ), gli autori forniscono kernel Triton fusi. Questi kernel eseguono la softmax online con RMSNorm fusa, eliminando la materializzazione intermedia di chiavi e sorgenti impilate. Ciò riduce il tempo di routing di $2\times$ rispetto a torch.compile e di $6.4\times$ rispetto alle implementazioni di riferimento eager. Questo porta il throughput di addestramento a $0.55\text{--}0.88\times$ della baseline.

Impostazione degli Iperparametri

Il numero di teste di routing $H$ è impostato uguale al numero di teste KV nel modello. Gli autori ritengono che questo sia un default quasi ottimale e privo di iperparametri, che allinea la granularità del routing con la granularità di consumo del modello.

3. Contributi Chiave

Architettura MHAR: Una modifica priva di parametri che suddivide il query di routing in teste indipendenti, eliminando il compromesso forzato del routing di profondità a singolo query.
Kernel Efficienti: Kernel Triton fusi che rendono pratico il routing della profondità su larga scala, riducendo significativamente l'overhead di memoria e la latenza.
Conversione Preservante l'Identità: Un metodo per innestare MHAR su modelli pre-addestrati usando "delta attention residuals" (aggiungendo il termine routato allo stream esistente invece di sostituirlo), consentendo il mid-training senza picchi di perdita.
Insight Meccanicistico: Il probing diretto dei query addestrati conferma che il disaccordo tra sottospazi cresce con la larghezza e che è il driver primario del degrado del routing a singola testa.

4. Risultati Sperimentali

Pretraining (Da zero)

Addestrato su FineWeb-Edu su tre scale (100M, 350M, 1B):

Loss di Validazione: MHAR migliora costantemente la loss di validazione rispetto ai Transformer standard e agli attention residuals a singola testa.
- 100M: $-0.049$ nats rispetto alla baseline.
- 350M: $-0.080$ nats rispetto alla baseline.
- 1B: $-0.063$ nats rispetto alla baseline.
Comportamento di Scaling: Il vantaggio di MHAR rispetto al routing a singola testa si amplia con la scala. Mentre il routing a singola testa aiuta a 100M ($-0.039$), diventa dannoso a 350M ( $+0.055$ ) e 1B ( $+0.140$ ).
Equivalenza Computazionale: MHAR raggiunge un guadagno computazionalmente equivalente di $1.3\text{--}1.5\times$ rispetto alla baseline.

Mid-Training (8B)

Utilizzando un controllo con schema di annealing su un modello da 8B (Marin-8B) su un corpus orientato alle capacità:

Preservazione dell'Identità: La conversione non introduce shock di ottimizzazione; le curve di loss coincidono con la baseline per i primi 100 step.
Guadagni Downstream: MHAR produce miglioramenti statisticamente significativi nei compiti di ragionamento:
- GSM8K: $+3.2$ punti (McNemar accoppiato $p=0.004$ ).
- GPQA: $+3.1$ punti (McNemar accoppiato $p=0.038$ ).
- I guadagni in MMLU, MATH e compiti di coding non sono statisticamente significativi in questa specifica configurazione.

Studi di Ablazione

Conteggio delle Teste: Il numero ottimale di teste $H$ segue il numero di teste KV. Nei regimi sotto-addestrati, più teste aiutano sempre; nei regimi convergenti, $H=KV$ è il "ginocchio" della curva.
Larghezza vs Profondità: Il probing dei query addestrati rivela che il "disaccordo di larghezza" (preferenze di routing divergenti tra sottospazi) cresce monotonicamente con la larghezza del modello, validando l'ipotesi del compromesso forzato.
Allineamento delle Teste: Allineare le teste di routing ai gruppi KV dell'attenzione non produce alcun beneficio misurabile rispetto al routing di sottospazio arbitrario, suggerendo che il beneficio derivi dalla molteplicità delle teste piuttosto che dal loro specifico allineamento alle strutture di attenzione.

5. Significato e Rivendicazioni

Il paper sostiene che Multi-Head Attention Residuals risolva un limite architettonico fondamentale nel modo in cui i Transformer accedono alla cronologia della profondità. Permettendo a diversi sottospazi di feature di prestare attenzione a diversi strati in modo indipendente, MHAR scala efficacemente dove il routing a singolo query fallisce.

Gli autori sottolineano che:

Il miglioramento non è dovuto all'aggiunta di parametri o calcolo, ma piuttosto alla rimozione di un collo di bottiglia strutturale.
Il metodo è pratico su larga scala grazie ai kernel fusi che mitigano la natura della memoria limitata del routing della profondità.
L'approccio è versatile, funzionando sia da zero che tramite mid-training preservante l'identità.
La necessità del routing multi-testa è una diretta conseguenza della larghezza del modello; man mano che i modelli crescono, il "compromesso forzato" di un singolo query diventa un limite, rendendo MHAR un componente critico per scalare le architetture consapevoli della profondità.