SQuat: Subspace-orthogonal KV Cache Quantization

Thanks for your careful reading of our paper and thoughtful comments!

R1 + Q1. There are two new hyperparameters introduced. This paper lacks discussion on how to set them effectively. A more detailed analysis regarding the specific influence of r on SQuat's behavior.

A1. Our recommended configuration for these two hyperparameters is $r=5$ and $\lambda=0.001$, and we use these values consistently across all LLMs and tasks in LM-eval. As shown in our paper, this choice consistently outperforms all other baselines (see our response to your Q2 for results for LongBench results). Below, we provide the intuition behind this selection.

For $r$, we recommend a small value (e.g., $r=5$). As you pointed out, a larger $r$ defines a higher-rank subspace, which makes satisfying the orthogonality constraints more difficult and can degrade model performance. This aligns with our ablation study (Table 3), where for various values of $\lambda$ (=$5 \times 10^{-4}$, $10^{-3}$, and $10^{-2}$), $r=5$ achieves the highest scores among all configurations. One exception occurs when $\lambda$ is very small (=$10^{-4}$); in this case, the orthogonality constraint is relaxed, and a slightly larger $r$ (e.g., $r=20$) yields better performance. As for $\lambda$, we recommend $\lambda=0.001$ (or even smaller), based on the observation that the largest singular value of the query subspace (i.e., $\Sigma[1]$) typically falls between 100 and 300. This choice of $\lambda$ ensures the regularization is effective without being too restrictive.

While we do not claim that this configuration is universally optimal (in fact, the best hyperparameters will vary depending on the specific task), our recommended configuration delivers consistent performance (outperforming all baselines tested) across different LLMs and evaluation tasks.

R2. Despite Llama 3.1's support for a 128k context length and the use of LongBench (a benchmark for long-context evaluation), the authors limited their experiments to a maximum sequence length of 8192 tokens. This is considerably shorter than the contexts explored in recent long-context modeling studies.

A2. To address your question, we increased the context length to 30,000 tokens. On this larger window, our results remain consistent with those reported in the original paper: our method achieves near-lossless performance even when quantizing the KV cache to INT2. Specifically, the Baseline (16-bit) achieves a score of 53.206, while SQuat (2-bit) achieves 52.554.

We chose 30,000 tokens because nearly all sequences from the selected LongBench tasks fall below this threshold. We acknowledge that there are other benchmarks having much longer sequences; however, 30,000 tokens is also the maximum context length we can experiment with due to our limited computational resources.

R3. The analysis lacks a detailed speed profile or computational cost breakdown for SQaut. Detailed and separated profiling of overheads in prefilling and generation is crucial for assessing practical performance and identifying bottlenecks.

A3. Please refer to our global response.

Q2. Why use different parameters values ($r$, $\lambda$) on LongBench?

A4. In this paper, we explore various parameter settings to evaluate the optimal performance of SQuat on LongBench, which consists of relatively simple tasks. Nonetheless, the default parameter values we recommend (i.e., $r = 5$, $\lambda = 0.001$) can also achieve near-lossless accuracy on LongBench and consistently outperform all baseline methods (see below):

Llama-2-7B-hf

Llama-3.1-8B-Instruct

Mistral-7B-Instruct-v0.3

Q3. What is the length of the input sequence used to generate Figure 2? Do the result reported in Sec 3.2, i.e., prompt only is a good enough approximation, hold for substantially longer sequences?

A5. The input sequence is 203 tokens long. We believe the observation that “prompt-only is a good enough approximation" holds even for much longer sequences (we tested 8k tokens and observed consistent results). This is likely due to the low-rank structure of the projection matrix that maps attention inputs to query vectors, which constrains the query subspace. In Figure 2 (left, red curve), even when the query subspace is derived from a different task (document QA instead of math), it still provides a reasonable approximation for the math-500 dataset’s token query tensors.

To be continued...

We thank the reviewer for the feedback and kind comments!

Q1. Insufficient guidance on hyperparameter selection beyond ablation studies.

A1. The main hyperparameters of our method are the subspace dimension $r$ and the regularization coefficient $\lambda$. Our recommended configuration for these two hyperparameters is $r=5$ and $\lambda=0.001$, and we use these values consistently across all LLMs and tasks in LM-eval. As shown in our paper, this choice consistently outperforms all other baselines (see our response to Reviewer hLUL’s Q2 for LongBench results). Below, we provide the intuition behind this recommendation.

For $r$, we recommend a small value (e.g., $r=5$). A larger $r$ defines a higher-rank subspace, which makes satisfying the orthogonality constraints more difficult and can degrade model performance. This aligns with our ablation study (Table 3), where for various values of $\lambda$ (=$5 \times 10^{-4}$, $10^{-3}$, or $10^{-2}$), $r=5$ achieves the highest scores among all configurations. One exception occurs when $\lambda$ is very small ($=10^{-4}$); in this case, the orthogonality constraint is relaxed, and a slightly larger $r$ (e.g., $r=20$) yields better performance. As for $\lambda$, we recommend $\lambda=0.001$ (or even smaller), based on the observation that the largest singular value of the query subspace (i.e., $\Sigma[1]$) typically falls between 100 and 300. This choice of $\lambda$ ensures the regularization is effective without being too restrictive.

While we do not claim that this configuration is universally optimal (in fact, the best hyperparameters will vary depending on the specific task), our recommended configuration delivers consistent performance (outperforming all baselines tested) across different LLMs and evaluation tasks.

In addition to $r$ and $\lambda$, our method includes two other hyperparameters: residual buffer length and group size. Throughout the paper, we set both of them to 32 tokens, following the recommendations from prior work [1, 2].

[1] Liu, Zirui, et al. "Kivi: A tuning-free asymmetric 2bit quantization for kv cache." ICML, 2024.

[2] Sheng, Ying, et al. "Flexgen: High-throughput generative inference of large language models with a single gpu", ICML, 2023.

Q2. Limited analysis of computational overhead during inference, particularly for initial tokens.

A2. Please refer to our global response for a detailed discussion.

We thank the reviewer for the thoughtful review and for appreciating the merits of the work!

Q1. The query subspace is estimated solely from prompt tokens, assuming it generalizes well to future tokens.

A1. To address your concern, we conducted two experiments.

Experiment 1: dynamic query subspace update

In this experiment, we dynamically update the query subspace using the residual buffer, rather than estimating it solely from the prompt tokens (an idea suggested by Reviewer hLUL). This approach is particularly useful in dialog scenarios where user responses may shift semantics during generation. We evaluate this method on the deepseek-distill-llama-8B model using 100 samples from the math-500 dataset. We update the query subspace every time 512 new tokens are generated. The resulting accuracy is 77%, matching the FP16 baseline and indicating no loss from quantization. For comparison, our original implementation (static subspace) has accuracy 72% and FP16 (no quantization) has accuracy 77%.

(Additional results can be found in our response to Reviewer hLUL.)

Experiment 2: robustness of query subspace construction

We test the robustness of our method by constructing the query subspace from easier or even unrelated tasks. Using the llama-3.1-8B-instruct model on 100 samples from the math-500 dataset, we compare the performance under three different subspace:

1. Current setup (subspace constructed from the first prompt in each batch): 47% accuracy.
2. GSM8K subspace (constructed from a prompt in the GSM8K dataset, an easier math dataset): 46% accuracy.
3. QuALITY subspace (constructed from a prompt in the QuALITY dataset, which is a document QA dataset, unrelated to math: 44% accuracy.

For reference, KIVI achieves 44% accuracy on the same set of math-500 samples.

These results are consistent with our observations in Figure 2 (left) from Section 3.2, suggesting that query vectors lie in a low-dimensional subspace. Even when constructed from different tasks, this subspace remains a reasonable approximation which is likely due to the low-rank nature of the projection matrix mapping attention inputs to query vectors.

Q2. The practical cost of SVD, subspace projection, and iterative updates, especially for longer sequences and large-scale deployment, needs clearer profiling in the main paper.

A2. Please refer to our global response. In addition, we conducted an experiment using longer sequences as per your recommendation. Specifically, we used the deepseek-distill-llama-8B model with 100 prompt tokens followed by 8,000 response tokens. This setup is designed to better reflect modern reasoning workloads, which often require extended responses for multi-step reasoning.

SQuat (our method):

FP16:

Q3. Real-world deployment scenarios, where key tensors may be quantized asynchronously or under hardware constraints (e.g., mobile inference), are not addressed.

A3. This is a great suggestion! We believe our method can indeed benefit from performing quantization asynchronously to further reduce latency. Our approach maintains a residual buffer that stores the most recent tokens’ KV cache in FP16. Once this buffer reaches a predefined length, the cached KV tensors are quantized together. This design naturally supports asynchronous quantization, as it avoids the need to quantize the KV cache immediately with each new token. Instead, decoding can continue using the FP16 KV cache in the residual buffer, while quantization of the buffer proceeds in parallel.

To address hardware constraints, one potential solution is to reimplement our algorithm using the MLX framework, which supports apple silicon. For example, adapting our current implementation (which builds on huggingface code) for the llama model to MLX could be done by adapting it using mlx-examples/llms/llama/llama.py. The main operations in our algorithm, such as SVD, are already supported in MLX (e.g., mlx.core.linalg.svd).