Attenzione casuale: ripensare l’espulsione dalla cache KV per un ragionamento efficiente

Abstract degli autori

I modelli linguistici di grandi dimensioni ottengono prestazioni superiori nei compiti che richiedono un ragionamento prolungato, ma le lunghe catene di pensiero rendono la cache KV un grave collo di bottiglia per la memoria. I metodi esistenti di compressione della cache KV seguono tutti lo stesso approccio: attribuire a ogni token nella cache un punteggio che stimi quanto sarà importante in seguito e conservare quelli con i punteggi più alti. Mostriamo che il segnale di selezione incide pochissimo. Random Attention conserva il prompt ed espelle i token con una selezione casuale uniforme all’interno di ciascuna testa di attenzione, senza calcolare alcun punteggio; su quattro modelli e sei compiti di ragionamento, ottiene risultati pari a quelli del migliore metodo precedente di espulsione, offrendo al contempo un throughput superiore del 32-43% rispetto a quest’ultimo in un deployment con vLLM. Esperimenti controllati spiegano perché, mostrando che 1) il prompt è la parte più vulnerabile della cache e gran parte della differenza tra i metodi di selezione dipende semplicemente dal fatto che il loro segnale abbia conservato o meno il prompt; 2) la traccia del ragionamento si protegge dall’espulsione grazie alla ridondanza su due livelli: nel testo (il modello ribadisce ciò di cui ha ancora bisogno mentre procede) e tra le teste di attenzione (ciascuna conserva una propria copia della traccia). Perciò, una volta messo al sicuro il prompt, una selezione casuale conserva un numero sufficiente di copie di ciò di cui il modello ha ancora bisogno, senza che serva un punteggio per sceglierle. Il nostro codice è disponibile pubblicamente all’indirizzo https://github.com/SalesforceAIResearch/Random-Attention.