SubGen: Generazione di Token in Tempo e Memoria Sublinerare

SubGen: Token Generation in Sublinear Time and Memory

February 8, 2024

Autori: Amir Zandieh, Insu Han, Vahab Mirrokni, Amin Karbasi

cs.AI

Abstract

Nonostante il notevole successo dei grandi modelli linguistici (LLM), i loro ingenti requisiti di memoria pongono sfide per il loro utilizzo nella generazione di token in contesti lunghi. L'impronta di memoria sostanziale dei decoder LLM deriva dalla necessità di memorizzare tutti i token precedenti nel modulo di attenzione, un requisito imposto dalla memorizzazione in cache delle coppie chiave-valore (KV). In questo lavoro, ci concentriamo sullo sviluppo di una tecnica di compressione efficiente per la cache KV. Le evidenze empiriche indicano una significativa tendenza alla clusterizzazione all'interno degli embedding delle chiavi nel modulo di attenzione. Basandoci su questa intuizione chiave, abbiamo ideato un nuovo metodo di caching con complessità sublineare, utilizzando il clustering online sui token chiave e il campionamento online ell_2 sui valori. Il risultato è un algoritmo di decodifica dell'attenzione provabilmente accurato ed efficiente, denominato SubGen. Questo algoritmo non solo garantisce un'impronta di memoria sublineare e una complessità temporale sublineare, ma stabiliamo anche un limite di errore rigoroso per il nostro approccio. Le valutazioni empiriche su compiti di risposta a domande in contesti lunghi dimostrano che SubGen supera significativamente i metodi esistenti e all'avanguardia di compressione della cache KV in termini di prestazioni ed efficienza.

English

Despite the significant success of large language models (LLMs), their extensive memory requirements pose challenges for deploying them in long-context token generation. The substantial memory footprint of LLM decoders arises from the necessity to store all previous tokens in the attention module, a requirement imposed by key-value (KV) caching. In this work, our focus is on developing an efficient compression technique for the KV cache. Empirical evidence indicates a significant clustering tendency within key embeddings in the attention module. Building on this key insight, we have devised a novel caching method with sublinear complexity, employing online clustering on key tokens and online ell_2 sampling on values. The result is a provably accurate and efficient attention decoding algorithm, termed SubGen. Not only does this algorithm ensure a sublinear memory footprint and sublinear time complexity, but we also establish a tight error bound for our approach. Empirical evaluations on long-context question-answering tasks demonstrate that SubGen significantly outperforms existing and state-of-the-art KV cache compression methods in terms of performance and efficiency.