Attenzione sparsa a blocchi con complessità log-lineare

Abstract degli autori

L’estensione dei modelli linguistici a contesti lunghi è limitata dal costo quadratico della self-attention. L’attenzione sparsa a blocchi offre un’alternativa efficiente, ma la scelta dei blocchi da mantenere resta un collo di bottiglia. La selezione convenzionale richiede di assegnare un punteggio a tutte le coppie query-blocco e mantiene quindi una complessità quadratica rispetto alla lunghezza della sequenza. Per affrontare questo problema, proponiamo PISA, un meccanismo di attenzione sparsa a blocchi che impiega una strategia piramidale di selezione Top-$K$. L’idea centrale è restringere gradualmente l’insieme dei candidati attraverso diversi livelli, così da individuare più efficientemente le chiavi più rilevanti. In particolare, costruiamo una gerarchia di chiavi dal livello più grossolano a quello più fine e iniziamo la selezione dal livello più grossolano. A ogni livello, applichiamo il calcolo dei punteggi LogSumExp a un insieme limitato di candidati per selezionare quelli da passare al livello successivo, più fine, fino a raggiungere il livello più fine. Mediante pooling, costruiamo $O(\log N)$ livelli di chiavi, ottenendo una complessità complessiva di $O(N\log N)$, dove $N$ indica la lunghezza della sequenza. Sviluppiamo kernel Triton che tengono conto dell’hardware sia per l’addestramento sia per l’inferenza, integrando l’instradamento gerarchico e il calcolo dei punteggi LogSumExp senza materializzare la matrice dei punteggi query-chiave. Valutiamo inoltre il nostro metodo in compiti di modellazione linguistica. Rispetto alla baseline, il nostro metodo ottiene prestazioni comparabili su benchmark come quelli di ragionamento basato sul senso comune, offrendo al contempo risultati migliori nei compiti di retrieval.