CoWindow Attention: la copertura causale completa è una proprietà collettiva

Abstract degli autori

FullAttn rende ripetutamente accessibile a ogni testa di attenzione l’intera storia causale, generando un notevole carico di calcolo ridondante e traffico di memoria anche con kernel densi efficienti nelle operazioni di input/output. Presentiamo CoWA, un’architettura di attenzione strutturata che distribuisce tra le teste KV l’accesso alla storia causale. Tutte le teste condividono finestre vicine alla diagonale e finestre prefix-sink, mentre finestre complementari a lungo raggio si suddividono la storia restante. La loro unione garantisce una copertura causale completa, sebbene ogni testa presti attenzione solo a una parte dei token distanti. Questo schema di attenzione, definito in base alla posizione, non richiede router o indicizzatori appresi, viene usato in modo coerente durante l’addestramento e l’inferenza ed è compatibile con il parallelismo tensoriale delle teste KV. Uno studio di ablazione con finestre di pari dimensioni a 8K isola l’effetto dell’assegnazione di finestre complementari a lungo raggio: CoWA, con una copertura collettiva del 100%, raggiunge un’accuratezza dell’89,73%, contro l’89,97% di FullAttn, mentre le finestre a lungo raggio duplicate ottengono risultati nettamente peggiori. In un confronto controllato più ampio sul richiamo associativo, con budget di token equivalenti, CoWA mantiene risultati vicini a quelli di FullAttn all’aumentare del contesto, mentre altri schemi di attenzione sparsa perdono una quota sostanziale delle associazioni. In un benchmark dell’operatore di attenzione a 128K token con parallelismo tensoriale, CoWA riduce rispetto a FullAttn la latenza dei passaggi forward e backward durante l’addestramento di un fattore rispettivamente di 7,4 e 8,6, e quella della decodifica durante l’inferenza di un fattore 3,0. Il picco di memoria dell’operatore per rank è pari a quello di FullAttn durante l’addestramento e inferiore di un fattore 7,6 durante la decodifica. Nell’addestramento secondo le leggi di scala, da 0,6B a 14B parametri, CoWA ottiene una perplexity molto vicina a quella di FullAttn, riducendo al contempo i FLOPs totali dell’addestramento. I modelli 14B così ottenuti e i modelli 32B derivanti da un addestramento continuato separato raggiungono punteggi comparabili a quelli di FullAttn nelle prove di conoscenza, ragionamento e recupero di informazioni in contesti lunghi. Questi risultati mostrano che la copertura causale completa può essere una proprietà collettiva dell’insieme delle teste, anziché una proprietà duplicata in ogni testa.