SubGen: Token-Generierung in sublinearer Zeit und mit sublinearem Speicher

SubGen: Token Generation in Sublinear Time and Memory

February 8, 2024

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

cs.AI

Zusammenfassung

Trotz der beachtlichen Erfolge großer Sprachmodelle (LLMs) stellen deren umfangreiche Speicheranforderungen Herausforderungen bei der Bereitstellung für die Erzeugung von Tokens mit langem Kontext dar. Der erhebliche Speicherbedarf von LLM-Decodern ergibt sich aus der Notwendigkeit, alle vorherigen Tokens im Attention-Modul zu speichern, eine Anforderung, die durch das Key-Value (KV)-Caching auferlegt wird. In dieser Arbeit konzentrieren wir uns auf die Entwicklung einer effizienten Kompressionstechnik für den KV-Cache. Empirische Belege deuten auf eine signifikante Clusterbildung innerhalb der Key-Embeddings im Attention-Modul hin. Aufbauend auf dieser zentralen Erkenntnis haben wir eine neuartige Caching-Methode mit sublinearer Komplexität entwickelt, die Online-Clustering für Key-Tokens und Online-ell_2-Sampling für Werte einsetzt. Das Ergebnis ist ein nachweislich genauer und effizienter Attention-Decodierungsalgorithmus, genannt SubGen. Dieser Algorithmus gewährleistet nicht nur einen sublinearen Speicherbedarf und eine sublineare Zeitkomplexität, sondern wir stellen auch eine enge Fehlerschranke für unseren Ansatz auf. Empirische Auswertungen bei Aufgaben zur Beantwortung von Fragen mit langem Kontext zeigen, dass SubGen bestehende und state-of-the-art KV-Cache-Kompressionsmethoden in Bezug auf Leistung und Effizienz deutlich übertrifft.

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.