ClusterGen: Token Generation in Sublinear Time and Memory with Clustering KV Cache

Thank you for your valuable feedback, thoughtful suggestions, and you interest in our result.

Typos and Notational Consistencies: We appreciate you pointing out the typos and inconsistencies in the use of k. We will correct these in the final version to ensure consistency throughout the paper.

Practical Guideline: We will include a practical guide for choosing the hyperparameters in the experiments section. Specifically, our main hyperparameter is the number of clusters, which we fixed at 1024 for the long-context experiments presented in Table 3.

Additional Experiments: We acknowledge the importance of testing with larger n. While computational constraints make this challenging, we will prioritize conducting additional experiments with longer contexts and include the results in the final version.

Thank you again for your constructive feedback!

Thank you for your thoughtful and detailed feedback. Below, we address the specific points raised.

Theorem 2.4 says the algorithm result satisfies equation (1); should it say equation (3) or (4)?

Theorem 2.4: Thanks for pointing out the typo. We will fix it.

Have sampling-based approaches for attention been used previously (in practice or theory)? I'm confused about why operator norm guarantees (eq (2)) are being used here, since they don't seem to have been used much previously for LLMs. Can you give intuition for why operator norm is a reasonable choice?

The primary focus of our paper is not the specific choice of norm guarantees but rather the development of an algorithm that leverages techniques from approximation algorithms—such as online/streaming clustering, random sampling, and approximate matrix multiplication—to address a practical problem in LLM efficiency, that is the KV cache compression. While we use operator norm guarantees (as in Eq. (3)) to establish rigorous theoretical results, the main contribution lies in applying tools from theoretical/algorithmic computer science to improve efficiency in LLMs.

In section 4.1 / table 2, how does the window size l compare between the three caching algorithms? Is the increase in accuracy coming because ClusterGen is able to use a bigger l? How significantly does integrating the window approach improve your algorithm? The use of the sliding window should be mentioned earlier when describing the algorithm, and not just in the experiment section.

To have a fair comparison we set the same values of the recent window size $\ell$ for all methods we compares (AttnSink, H2O, ClusterGen) in order to have equal KV cache size for all competing methods. This part is independent of those quantization methods, i.e., keeping the bottom $\ll$-by-d submatrices separately. We will update the sliding window part in the algorithm description.

Did you verify the clustering assumption (definition 2.1) in the section 4 experiments? Is there a relationship between how clusterable the data is and the accuracy of the model? Does this explain why ClusterGen is better at some tasks than others in table 3?

In Section 3.1 and Appendix B, we provide empirical evidence on our clustering assumptions across different LLMs. Specifically, a better clusterability is generally observed in LLMs showing high evaluation performance based on scores of LLM leaderboard. These observations strongly support the validity and practicality of our clusterability assumption and the proposed algorithm.

How did you choose which of the algorithms described in the related work (section 1.1) to compare with in the experiments (section 4)?

Since our method focuses on KV cache compression through token pruning, we specifically compare it against other token pruning approaches for a fair evaluation. In our experimental section, we benchmark our ClusterGen against prominent token pruning techniques, including H2O and AttentionSink.

How do the running times of the different algorithms compare in experiments? It seems only the memory usage is being considered.

We evaluate the end-to-end token generation times on multi-news dataset from LongBench. Compared to the exact method (without KV cache compression), ClusterGen shows approximately 20% slower runtime in order to compress the KV cache. However, with this marginal delay, the memory space can be saved more than 5 times.

Comparison with "Exact" Attention and FlashAttention implementation:

1. There appears to be a misunderstanding about the context in which our algorithm operates. Our primary focus is on the generation phase of LLMs, where the model generates one token at a time in an online streaming mode. This phase differs significantly from the prompt phase, which involves a batch of n queries and key/value embeddings.

2. FlashAttention, as far as we know, is designed to optimize the prompt phase by leveraging techniques like checkpointing and tiling to reduce memory and computational overhead from quadratic to almost linear in n. However, during the generation phase, the runtime is less critical because attention operations occur for a single query, and the total cost scales linearly with the context size. Instead, memory concerns dominate during generation, as the KV cache size grows proportionally to the context length.

3. Our experiments explicitly target this memory bottleneck during generation. We provide comparisons focused on memory reduction, which is the critical constraint in this phase. Furthermore, in the generation phase, feed-forward layers often dominate runtime costs, while attention operations are comparatively lightweight. This justifies our emphasis on memory optimization in this context.

Comparison to Polynomial based methods:

Polynomial-based methods (e.g., Hedgehog, Porcupine) are an interesting line of research, but they operate under assumptions and constraints that differ significantly from ours:

1. Assumption of Small Embedding entries: These methods require key embeddings to have small entries for polynomial approximations of the exponential function to remain accurate (as shown in Alman & Song). This corresponds to a single-cluster setting in our framework. In contrast, our method supports a more general multi-cluster setting, making it applicable to more realistic scenarios where embeddings exhibit diverse structures.

2. Limitations in polynomial degree: Polynomial kernels correspond to tensor product embeddings whose dimensions grow exponentially with the degree of the polynomial. In practice, these dimensions must remain small (e.g., 3 or 4) to control the dimension of the transformed embedding vectors, seriously reducing the accuracy of the exponential approximation. Consequently, such methods require pre-training or at least fine-tuning to achieve competitive performance.

3. Post-Training Focus: Our work focuses on post-training KV cache compression, where we directly apply our method to pre-trained and fine-tuned models without requiring further training. This distinction is critical, as pretraining or fine-tuning is outside the scope of our study and we believe those problems require other types of techniques.

Comparison to alternative architectures like mamba:

• While alternative architectures such as Mamba provide intriguing pathways to efficient LLMs, they fall outside the scope of our paper. Our focus remains on dot-product attention with softmax activation and optimizing its KV cache usage during inference. Comparisons with attention-free models would require addressing architectural changes and training new models, which is a fundamentally different problem. Including such comparisons would detract from the core contributions of our work.

Accuracy compared to the exact Attention:

We acknowledge the trade-off between accuracy and memory savings in our method. However, it is important to reiterate that:

1. Pretraining/Fine-Tuning is Out of Scope: Unlike works that combine training-phase techniques with hardware optimizations, we focus solely on post-training inference scenarios. Naturally, fine-tuning or pretraining our method would improve accuracy but requires additional methods and resources that are not addressed in this paper.

2. Tailored to the Defined Problem: Our solution is specifically designed for post-training KV cache compression, where memory is the bottleneck during generation. Accuracy improvements via fine-tuning represent a separate problem that can build upon our work but is not its focus.

Thank you for your thoughtful and detailed feedback. Below, we address the specific points raised.

RAM model vs GPU/TPU model: We appreciate your observations regarding the applicability of our algorithm to modern hardware systems, including GPUs and TPUs. We address your concerns below, demonstrating how our algorithm can indeed be efficiently implemented in the GPU/TPU computation model and highlighting its compatibility with operations like GEMM, GEMV, and tensor reductions.

1. Line 11: Distance Computation. To compute all distances $\| x_i - k \|$ for $i \in [m]$ we can store $x_i$’s in rows of a matrix (or tensor) and compute these norms using a broadcast followed by a tensor norm reduction, which is highly optimized in current GPU libraries. Finding the index with the minimum norm is similarly efficient and can be implemented using GPU-supported argmin operations

2. Lines 15 - 16: The batched random sampling step can be done by generating a vector (tensor) of uniform random variables in the range [0,1) and comparing them against a threshold. Given that the vector dimension $t$ is a small constant in practice (e.g. 32), this operation is extremely fast. Rows passing the threshold are copied to the matrix (tensor) S, which is also a straightforward operation on GPUs.

3. Lines 24-25: Single Random Sample and Update. We sample a single random variable in the range [0,1) and compare it to a threshold. If it passes the threshold we copy the new key/value embedding vectors to the corresponding rows of the key-value cache matrix (tensor).

4. Lines 28-29: Matrix-Vector Operations. These steps can be both implemented by GEMV operations as we need to first do a matmul between the tensor of sampled key embeddings and another vector, followed by an elementwise exponential operation. Next we do an inner product with weights vector (tensor) which are inverse sampling probabilities. Then the attention scores need to be multiplied by the sampled value embeddings which can be implemented with GEMV.

5. Next note that in attention layers, computations across batch and head dimensions are entirely independent. Our algorithm naturally supports parallelism across these dimensions, ensuring it scales effectively with hardware-accelerated operations.

Our algorithm in fact exploits tensor operations (e.g., GEMM/GEMV), ensuring compatibility with block-level computations on GPUs. The random sampling mechanism in our algorithm does not violate the block-level memory access paradigm, as we operate within tensorized blocks and optimize access patterns. We welcome further discussion and feedback on these implementation details.

We address your concerns regarding experimental results in the official comment below.

OK, this makes sense but one thing that is not clear. If one were to count the number of block ops and GEMM/GEMV is the improvement over plain vanilla Attention would the improvement be the same? Having bounds in terms of number of block operations would be useful to see.

Thank you for your thoughtful and insightful feedback. Below, we address the specific points raised.

How does ClusterGen perform on tasks other than question-answering?

We appreciate your question about the results of other tasks. Following the reviewer’s question, we additionally experiment on a summarization task, i.e., ``multi_news’’ dataset from LongBench, comparing our method against AttentionSink. In particular, we use a fine-tuned Llama-3-8B-Instruct model [1] with context lengths ranging from 8K to over 1 million. Under the same memory constraints, our method achieves a ROUGE score (higher is better) of 27.2, outperforming AttentionSink (26.99). The score of the exact method is 27.43. This result further demonstrates our effectiveness across diverse long-context tasks. We will include a more comprehensive analysis with diverse tasks in our final draft.

[1] https://huggingface.co/gradientai/Llama-3-8B-Instruct-Gradient-1048k.

Have the authors considered methods to handle cases with low clusterability?

We appreciate your question about low clusterability. Although we haven’t considered specific methods with low clusterability, several approaches that can be applied to the coreset selection problem (which equals to KV cache compression) would be promising future research directions. However, as we studied in Section 3.1, such clusterable properties are commonly observed in LLMs with higher evaluation scores, suggesting a correlation between the clusterability of the key cache and the evaluation performance of a given LLM on a specific dataset. We think that approaches that handle clusterability are much more important than those with low clusterability.