PaperlayerView exact revision on arXiv

Source-linked AI summary

SubGen: Token Generation in Sublinear Time and Memory

Amir Zandieh, Insu Han, Vahab Mirrokni, Amin Karbasi

arXiv:2402.06082v1cs.LGcs.AIcs.DS

TL;DR

Long-context autoregressive decoding faces linear memory and runtime growth from retaining prior tokens in the KV cache. SubGen combines key clustering with value-token sampling to approximate attention in sublinear time and memory. Under its stated assumptions, it provides an approximation guarantee and outperforms other KV-cache compression algorithms at the same memory budget on zero-shot line retrieval.

  • Problem

    KV caching stores prior tokens, so memory requirements and per-token runtime grow linearly with context length.

    Source: §1 · p. 2

  • Method

    SubGen approximates attention using online key clustering and sampling of key-value pairs according to squared value-token norms.

    Source: §1.3 · p. 3, §2.3 · p. 7

  • Results

    SubGen outperforms other KV-cache compression algorithms with the same memory budget on zero-shot line retrieval tasks.

    Source: §4 · p. 13

  • Takeaways & Limitations

    Under bounded-query-norm and key-clusterability assumptions, SubGen has sublinear memory and runtime with an attention approximation guarantee.

    Source: §4 · p. 13

    Fast worst-case runtime for producing output embeddings is likely impossible without assumptions about the input tokens.

    Source: §1.1 · p. 2

Abstract

from arXiv · show

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.

1 Introduction

SubGen targets the linear memory and runtime costs of long-context autoregressive decoding with a sublinear-time, sublinear-memory attention approximation. The introduction presents its key-clustering approach, conditional guarantees, and empirical comparisons with KV-cache compression methods.

1 Introduction

  • Autoregressive decoding must produce contextually relevant tokens one at a time, while KV caching makes memory needs and per-token runtime grow linearly with context length.The paper asks whether attention output can be approximated with sublinear space and time in context length.Source: §1 · pp. 1–2

1.1 Related Work

  • Prior KV-cache methods include attention-based eviction, adaptive compression, recent-and-pivotal-token retention, and contextual-sparsity acceleration; the paper says none achieves fully sublinear time and memory.Fast worst-case runtime for arbitrary input tokens is likely impossible without assumptions about those tokens.Source: §1.1 · p. 2

1.2 Streaming Attention Problem

  • Streaming attention processes query, key, and value vectors for each token, but storing all preceding keys and values requires O(nd) memory.The attention output is used to predict the next token.Source: §1.2 · pp. 2–3

1.3 Overview of Contributions

  • SubGen approximates the attention matrix-vector product by subsampling O(ε^-2 d log n) key-value pairs using squared value-token norms.The passage describes implementing this sampling in a streaming setting with reservoir-sampling variants.Source: §1.3 · p. 3
  • To approximate the softmax normalizer, SubGen clusters keys and samples keys from each cluster; its assumption allows any sublinear number of bounded-diameter clusters.The clustering algorithm is inspired by streaming k-center.Source: §1.3 · p. 3
  • When keys form m = n^(1−Ω(1)) bounded-diameter clusters, SubGen uses sublinear memory and runtime while meeting the stated approximation guarantee.The result is conditional on the key-clustering assumption.Source: §1.3 · p. 4
  • SubGen outperforms other KV-cache compression algorithms with the same memory budget on zero-shot line retrieval tasks.The introduction also reports stronger results particularly on question-answering tasks with various sequence lengths.Source: §1.3 · p. 4

2 Sublinear Time and Memory Algorithm

SubGen compresses streaming attention by maintaining separate summaries for the value-weighted matrix product and the softmax partition function. Under clusterability and query-norm conditions, it provides a probabilistically accurate estimate with sublinear memory and runtime.

2.2 Matrix Product Data Structure

  • Row-norm sampling approximates the value-weighted matrix product using O(ε^-2 d log n) samples, with a reservoir-sampling variant for streams.Sampling probability is proportional to the ℓ2 norm of rows in the value matrix.Source: §2.1 · p. 5

2.1 Streaming Attention Data Structure

  • For the softmax denominator, SubGen assumes keys form m=o(n) clusters of diameter at most δ, enabling sublinear-memory storage.The authors cite observed key-embedding clusterability as motivation for using stream clustering.Source: §2.1 · p. 5

2.3 Softmax Normalizer (Partition Function) DS

  • UpdateSoftmaxNormalizer assigns keys to the nearest representative within δ, creates a cluster otherwise, and keeps t sampled keys per cluster.The per-cluster samples are maintained through reservoir sampling.Source: §2.3 · p. 7

2 Sublinear Time and Memory Algorithm

  • SubGen uses its streaming data structures to produce an accurate approximation of the attention output for a given query.The method approximates both the matrix-vector product and the softmax partition function.Source: §2.4 · p. 9

2.4 Streaming Attention: Main Theorem

  • With (m, δ)-clusterable keys and bounded queries satisfying δr=o(log n), SubGen returns an attention estimate with probability at least 0.99 and memory/runtime O(d·(mt+s)).The theorem bounds both total memory and runtime during the n-th iteration.Source: §2.4 · pp. 9, 11
  • Under sublinear cache budgets, the k-center-based method outperforms other methods across LongEval line-retrieval sequence lengths.The reported context-length range is 5k–9k tokens.Source: Table 1 · p. 11

3 Experiments

Experiments find key embeddings more clusterable than values, and SubGen outperforms H2O and Attention Sink on LongEval line retrieval. The comparison uses equal cache budgets across algorithms.

3 Experiments

  • Experiments use a single NVIDIA RTX6000 GPU with 48 GB VRAM.Source: §3 · p. 11

3.1 Ablation Study on Clusterability

  • Key embeddings are more clusterable than value embeddings across the randomly selected layers and heads.The key cluster centers are distributed across the embedding space.Source: §3.1 · p. 12
  • RoPE transformations are offered as the reason keys appear well-separated while values show an unstructured, random distribution.This motivates compressing keys through clustering.Source: §3.1 · p. 12

3.2 End-to-end Text Generation

  • LongEval line retrieval varies targets across 200, 300, and 400, corresponding to sequence lengths n=5,000, 7,000, and 9,000, with 50 questions per dataset.Answers are generated with longchat-7B.Source: §3.2 · p. 12
  • SubGen retains the most recent r embeddings plus k cluster centers, and all algorithms use the same r + k cache budget, which scales sublinearly with context length n.Greedy k-center clustering is applied once to compress the full KV cache.Source: §3.2 · p. 12
  • At 9k tokens, SubGen reaches 44% accuracy using half of the cached KV embeddings, while H2O and Sink are 10% lower.Source: §3.2 · p. 13

4 Conclusion

SubGen compresses the KV cache through stream clustering, using clustered keys and recent-token retention. Under bounded-query-norm and clusterability assumptions, it has a spectral error bound with sublinear time and memory, and outperforms other KV cache compression algorithms on zero-shot line retrieval at the same memory budget.

  • SubGen is an efficient KV cache compression algorithm based on stream clustering.Source: §4 · p. 13
  • Under bounded-query-norm and clusterability assumptions, SubGen guarantees a spectral error bound with sublinear time and memory.Source: §4 · p. 13
  • On zero-shot line retrieval tasks, SubGen outperforms other KV cache compression algorithms at the same memory budget.Source: §4 · p. 13
Loading 2402.06082v1…