Emergent Mind

SubGen: Token Generation in Sublinear Time and Memory

(2402.06082)

Published Feb 8, 2024 in cs.LG , cs.AI , and cs.DS

Abstract

Despite the significant success of 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.

A t-SNE plot showing cached keys and values embeddings across timesteps from Llama2-7B on MT Bench.

Overview

SubGen presents an approach to compress the KV cache in LLMs using online clustering, aiming to tackle the challenges in generating long sequences by reducing memory and runtime complexities.
It innovatively achieves sublinear-time memory space for KV cache compression, outperforming previous methods and effectively addressing the token generation efficiency.
SubGen's core lies in its clustering method for KV cache compression and online sampling, ensuring sublinear complexity without significantly compromising accuracy.
Empirical evaluations highlight SubGen's superior performance in handling sequences of varying lengths, efficiently balancing between efficiency and performance.

SubGen: An Efficient Clustering-Based Caching Method for LLMs

Introduction

The emergence of LLMs such as Transformer architectures has significantly advanced the capabilities of natural language processing applications. However, deploying these models, especially for tasks that require generating long sequences of tokens, presents substantial challenges due to the linearly scaling memory and computational constraints. A critical aspect of this challenge is the extensive memory footprint required by the attention mechanism's key-value (KV) caching during token generation. Addressing this, the paper introduces SubGen, an innovative approach designed to compress the KV cache utilizing online clustering, significantly reducing both memory and runtime complexities.

Related Work

Efficient token generation has become more critical with the rise of applications necessitating the handling of long-range context datasets. Previous works have proposed various strategies to compress the KV cache efficiently. These methods range from dynamic eviction algorithms, based on accumulated attention scores, to adaptive and deterministic policies that selectively retain crucial tokens. Despite their efforts, none have achieved a fully sublinear-time memory space method for KV cache compression, an accomplishment that SubGen successfully achieves.

Stream Attention Problem

LLMs often employ attention mechanisms where, in streaming settings, the sequence of tokens is decoded autoregressively. This process demands a mechanism to efficiently store and retrieve key and value pairs from all preceding tokens, known as KV caching. However, the linear scaling of memory requirements with the context size remains a primary bottleneck for efficiency. SubGen addresses this by proposing a method that ensures sublinear memory footprint and time complexity without significant compromise on accuracy.

Overview of Contributions

The core contribution of SubGen lies in its novel clustering method for KV cache compression, achieving sublinear complexity. By exploiting the significant clustering tendency observed within key embeddings in the attention module, SubGen compresses the KV cache through online clustering on key tokens and online sampling on values, establishing a tight error bound on its approximate attention decoding. Empirical evaluations confirm that SubGen significantly surpasses existing KV cache compression methods, striking an optimal balance between performance and efficiency.

Sublinear Time and Memory Algorithm

SubGen algorithmically enforces a sublinear space complexity in context length during the streaming attention process. This is achieved through a streaming attention data structure capable of approximating the attention mechanism's output with ample accuracy while conserving memory space efficiently. A notable facet of SubGen is the streaming attention data structure and the two pivotal procedures, UpdateSoftmaxNormalizer and UpdateMatrixProduct, each tailored for specific functionalities within the overall algorithm—in ensuring the space complexity remains sublinear, intrinsic to streamed token triplets.

Correctness and Efficiency Analysis

SubGen demonstrates a provable efficiency and correctness under certain conditions related to clusterability and bounding of query norms, verified through mathematical analysis and empirical evidence. The paper's analysis illustrates that under assumptions of bounded query norms and clusterability of keys, SubGen guarantees acceptable spectral error bounds, concurrently achieving sublinear time and memory. Furthermore, empirical evaluations on question-answering tasks showcase that SubGen holds superior performance, particularly in handling sequences of varying lengths efficiently.

Conclusions

SubGen introduces an innovative caching method leveraging clustering for KV cache compression in LLMs, presenting a practical solution to the challenge of efficiently generating long sequences. By maintaining a sublinear memory footprint and computational complexity without significantly sacrificing performance, SubGen paves the way for more scalable and efficient deployment of LLMs, especially in applications facing lengthy sequences. Its empirical superiority over existing approaches underpins the potential of incorporating clustering dynamics into the realm of LLM optimizations, opening avenues for future explorations in more adaptive, efficient caching methodologies.