SubGen: Token-generatie in sublineaire tijd en geheugen

SubGen: Token Generation in Sublinear Time and Memory

February 8, 2024

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

cs.AI

Samenvatting

Ondanks het aanzienlijke succes van grote taalmmodellen (LLM's), vormen hun uitgebreide geheugeneisen uitdagingen bij het inzetten ervan voor het genereren van tokens in lange contexten. De aanzienlijke geheugenvoetafdruk van LLM-decoders ontstaat door de noodzaak om alle voorgaande tokens op te slaan in de aandachtmodule, een vereiste die wordt opgelegd door key-value (KV)-caching. In dit werk richten we ons op het ontwikkelen van een efficiënte compressietechniek voor de KV-cache. Empirisch bewijs wijst op een significante clusteringstendens binnen key-embeddings in de aandachtmodule. Gebaseerd op dit inzicht hebben we een nieuwe cachingmethode ontwikkeld met sublineaire complexiteit, waarbij online clustering wordt toegepast op key-tokens en online ell_2-steekproeven op waarden. Het resultaat is een bewezen nauwkeurig en efficiënt aandacht-decodeeralgoritme, genaamd SubGen. Dit algoritme zorgt niet alleen voor een sublineaire geheugenvoetafdruk en sublineaire tijdcomplexiteit, maar we leggen ook een strakke foutgrens vast voor onze aanpak. Empirische evaluaties op taken voor het beantwoorden van vragen in lange contexten tonen aan dat SubGen aanzienlijk beter presteert dan bestaande en state-of-the-art KV-cachecompressiemethoden wat betreft prestaties en efficiëntie.

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.