Prism Attention: meccanismo di condivisione degli indici con affinamento delle proposte per un’inferenza efficiente degli LLM

Abstract degli autori

Nella decodifica di contesti lunghi, i modelli linguistici di grandi dimensioni (LLM) sono sempre più limitati dalla memoria e dal calcolo necessari per applicare l’attenzione a una cache Key-Value (KV) che cresce a ogni passo autoregressivo. L’attenzione sparsa a livello di token riduce questo carico limitando il calcolo a un piccolo sottoinsieme di token passati rilevanti, ma i benefici pratici sono spesso vanificati dalla necessità di ricalcolare ripetutamente gli indici sparsi per ciascun livello man mano che il contesto cresce. Una soluzione comune consiste nel condividere gli indici sparsi tra i livelli, calcolandoli soltanto in determinati livelli di ancoraggio. Questo riutilizzo statico, però, non tiene conto dell’evoluzione delle configurazioni dell’attenzione da un livello all’altro: gli indici riutilizzati si allontanano progressivamente da quelli ottimali per ciascun livello, causando una perdita di attenzione sostanziale. Riesaminando la sparsità tra livelli dal punto di vista della dinamica dell’attenzione, osserviamo che lo spostamento dei token salienti tra i livelli è fortemente localizzato: sebbene gli indici esatti dei primi $κ$ token differiscano da un livello all’altro, i token centrali che ricevono più attenzione restano stabili e i cambiamenti si concentrano vicino al confine ampliato delle regioni selezionate in precedenza. Su questa base proponiamo Prism Attention, un framework di attenzione sparsa che affina le proposte. Sostituisce il rigido riutilizzo degli indici tra livelli con una strategia dal grossolano al fine, preservando la capacità di adattamento a ciascun livello e mantenendo i vantaggi di efficienza della condivisione tra livelli. Numerosi esperimenti mostrano che Prism Attention ottiene prestazioni all’avanguardia tra i principali metodi di attenzione sparsa, sia nei benchmark di ragionamento sia in quelli per contesti lunghi, e raggiunge una velocità di decodifica end-to-end pari a $3.4\times$ quella di FlashAttention con una lunghezza del contesto di 128K.