Attenzione con testa virtuale

Abstract degli autori

La progettazione dei meccanismi di attenzione nei grandi modelli linguistici a contesto lungo si scontra con un «triangolo impossibile» tra dimensione della cache KV, FLOPs di addestramento e prefill e qualità del modello. GQA offre una qualità elevata, ma richiede una cache KV di grandi dimensioni; approcci alternativi come MLA e MFA possono comprimerla, ma MLA comporta una perdita di qualità, mentre MFA aumenta notevolmente il calcolo necessario per l’addestramento e il prefill. Proponiamo Virtual Head Attention (VHA), che migliora simultaneamente tutti e tre i vertici grazie a due operazioni lineari leggere: Q Premix applica alle query di ciascun gruppo KV una trasformazione prossima all’identità per recuperare la diversità delle teste virtuali pur avendo dimezzato il numero di teste di query fisiche; Linear Postmix combina le caratteristiche tra teste mediante una struttura residua a basso rango $I + AB^\top$. Poiché VHA non introduce operazioni non lineari, durante l’inferenza Premix può essere incorporato nella proiezione delle query e Postmix nella proiezione dell’output, riducendo il modello a un’architettura GQA-2 standard che riutilizza direttamente tutte le ottimizzazioni esistenti per l’inferenza di GQA. Con due sole teste KV, nei modelli di grandi dimensioni la dimensione dei gruppi raggiunge facilmente 64, sfruttando appieno il parallelismo dei Tensor Core durante la decodifica. Su quattro scale, con un preaddestramento su 50B token e un contesto di 4K, VHA supera costantemente la baseline GQA standard a ogni scala (da $+$0,05\% a $+$0,81\%), riducendo al contempo la cache KV di $4\times$ vs.\ GQA-8 alle scale dei modelli densi e di $2\times$ vs.\ GQA-4 alla scala del modello MoE, e abbassando i FLOPs di addestramento e prefill del 7,9\% (1,7B, $S{=}4096$). Su un modello MoE industriale 30B-A3B addestrato su 500B token con un contesto di 8K, VHA supera sia MLA sia GQA-4 di $+$0,70\,/\,$+$2,60\,\% in una suite di valutazione del modello base composta da 12 benchmark; in un confronto separato, a parità di condizioni, tra modelli MoE 10B-A1B dopo il mid-training su contesti lunghi, VHA ottiene un punteggio RULER medio di 84,4 nell’intervallo 8K–128K, contro 73,0 di GQA, pur usando due teste KV anziché quattro nei livelli di attenzione completa.