Efficient Low Rank Attention for Long-Context Inference in Large Language Models

We thank the reviewer for the constructive feedback. Below are our detailed responses to each concern raised.

W 1: The observations regarding singular values and key-value decomposition are not novel; similar findings have already been reported in prior work such as Loki [1].

[1] Loki: Low-rank Keys for Efficient Sparse Attention

Thank you for highlighting this important context. You are correct that low-rank properties of attention matrices have been explored in prior work, such as Loki [1]. However, the key distinction lies in how we operationalize this property:

• Loki decomposes the key matrix ($\mathbf{K}$) individually to exploit its low-rank structure.
• Our method performs simultaneous low-rank decomposition of both $\mathbf{Q}$ and $\mathbf{K}$, leveraging the property of matrix multiplication:

This enables a tighter approximation:

where $\mathbf{A}\_Q$ and $\mathbf{A}\_K$ are low-rank proxies for $\mathbf{Q}$ and $\mathbf{K}$.

By jointly approximating $\mathbf{Q}$ and $\mathbf{K}$, we avoid the limitations of isolated decomposition. This design, combined with a proxy approximation strategy (rather than direct filling), leads to superior performance in both efficiency and accuracy compared to Loki.

W 2: The proposed method introduces a two-stage approximation involving proxy attention scores and hierarchical cache management, which adds notable complexity to the inference pipeline. However, the paper does not adequately address the implementation or integration overhead.

We acknowledge this important concern and will provide comprehensive implementation details in the revision. The implementation and integration overhead analysis are as follows:

Implementation overhead:

• Prefill stage: Need to add a cost to compute low-rank projections on top of standard attention (increase computation). After projection, we can pick the most relevant tokens and compute attention based on the selected tokens (reduce computation).
• Generation stage: Additional computation for low-rank matrices updates (increase computation). Since the number of selected tokens are limited, the computation cost of attention is reduced (reduce computation).
• Cache management: The KV caches are offloaded to CPU, which requires additional lookup in memory and CPU-GPU transfer (increase computation). Additional computation is required to manage the cache (increase computation).

Integration overheads: There is a brief version of the code in supplementary material, which shows the basic idea of our implementation. A brief overview of the implementation is as follows:

• The implementation is based on the standard huggingface transformers library.
• The original attention mechanism is replaced by our proxy attention mechanism.
• The cache management is following transformers.cache_utils.Cache, which makes it easier to integrate into existing transformer architectures.

W3: Although the authors acknowledge in the Limitations section that the method requires careful hyperparameter tuning, no concrete strategies or guidelines are provided—this significantly limits its practical usability.

Thank you for this valuable comment. We will address this concern by providing guidelines for hyperparameter selection.

Rank $(r)$: The low-rank approximation of matrices $\mathbf{Q}$ and $\mathbf{K}$ is calculated for each attention head separately. Typically, the dimension of each head is $128$, allowing for a range of $r$ values between $1$ and $128$. For practical purposes, common choices for the rank $r$ are $\{8, 16, 32, 48\}$, offering a balance between approximation quality and computational efficiency.

Top-$k$ Tokens: Select and retain the $k$ tokens with the highest attention scores in the GPU cache. The choice of $k$ is pivotal in balancing the sequence length and memory usage. Short sequences typically require a smaller $k$, while longer sequences necessitate a larger $k$ to manage memory efficiently. In the experiment, $k=256$ is used for a 4K token context, and the maximum is $k=2048$, applying for context longer than 32K.

Number of Lite Tokens: This hyperparameter determines the number of most recent tokens that are kept in the GPU. It has some flexibility in its choice; common values are 16, 32, or 64. This choice does not significantly increase the computation.

Number of Iterations and Tolerance: These two parameters are introduced in Algorithm 1 and 2. The number of iterations refers to the number of loops for updating all related low-rank matrices. Typically, a few iterations, such as 2 or 3, is sufficient for the task. The tolerance is the threshold for the convergence of the low-rank matrices, which is defined as the mean square error between the current and previous matrices. The tolerance is usually set to a small number, for example $10^{-2}$ or $10^{-3}$.

Tips: If no prior knowledge, it is recommended to start with the default values, $r=32, k=2048$, 64 lite tokens and 2 iterations with tolerance $10^{-2}$. Then, we can adjust the parameters based on the performance and memory usage.

We appreciate the reviewer's thoughtful evaluation and recognition of the proposed method. Below, we address your specific concerns:

W1. Hyperparameter Sensitivity: The method requires careful tuning of rank and top-k token counts, which vary across model architectures and tasks. Optimal values demand extensive grid searches, increasing implementation overhead.

Thank you for pointing out the sensitivity to hyperparameters. To achive better performance, careful tuning is indeed required. Nevertheless, the search space is relatively small, and there is a recommend configuration: $r=32, k=2048$, 64 lite tokens and 2 iterations with tolerance $10^{-2}$. We can further turn based on this. The detailed guideline for hyperparameter are as follows:

Rank $r$: The matrices $\mathbf{Q}$ and $\mathbf{K}$ are approximated using low-rank decomposition for each attention head independently. Given that each head usually has a dimension of $128$, the rank $r$ can vary from $1$ to $128$. In practice, common choices for $r$ are $\{8, 16, 32, 48\}$, which balance between the quality of approximation and computational efficiency.

Top-$k$ Tokens: This hyperparameter involves selecting and retaining the $k$ tokens with the highest attention scores in the GPU cache. The selection of $k$ is crucial for balancing the sequence length and memory usage. Shorter sequences generally require a smaller $k$, while longer sequences may need a larger $k$ for efficient memory management. Experimentally, $k=256$ is suitable for a 4K token context, with $k=2048$ used for contexts exceeding 32K tokens.

Number of Lite Tokens: This parameter specifies the count of the most recent tokens maintained in the GPU. Common values are 16, 32, or 64, and this choice does not substantially impact computational overhead.

Iterations and Tolerance: These parameters are integral to Algorithms 1 and 2. The iteration count denotes the number of cycles required to update the low-rank matrices. Generally, 2 or 3 iterations suffice. Tolerance is the convergence threshold for the low-rank matrices, calculated as the mean squared error between successive matrices, typically set to a small value like $10^{-2}$ or $10^{-3}$.

W2. The paper lacks performance analysis. It would be better if the authors could compare their method with baselines in terms of throughput and TTFT.

We provide comprehensive timing results on single NVIDIA A100 40GB with Llama-3-8B-1M.

Main Results: The GPU only method can provide the fastest TTFT (Time To First Token; Prefill time) and TPOT (Time To First Token; Prefill time), but will reach out of memory (OOM) when context is too long (64K). CPU offload method can handle longer context but will suffer from limited GPU bandwidth and result in slow decoding speed, expecially for long context. The default proposed method can support longer context and provide a stable throughput. The proposed method without hit/miss mecanism is slower than the default in decoding. Further details can be found in the supplemental material. The following table shows the results.

Dear Reviewer 7jqo,

Thank you for your valuable feedback. We appreciate your interest in our work and are committed to addressing your concerns. A brief summary of your comments and our responses is provided below:

• Hyperparameter Sensitivity: We have carefully considered your concern regarding the sensitivity of the proposed method to hyperparameters. We have provided detailed guidelines for tuning the rank, top-k tokens, number of lite tokens, and iterations and tolerance. These guidelines should help users achieve optimal performance without extensive grid searches.
• System Performance Metrics: We have included detailed timing results on a single NVIDIA A100 40GB with Llama-3-8B-1M to demonstrate the performance of our method. The results show that our method can handle longer contexts while maintaining a stable throughput.

If you have any further questions or require additional information, please do not hesitate to ask. We are here to assist you in any way we can.

We thank the reviewer for the detailed and insightful comments.

W1. The empirical comparisons primarily involve sparse/dynamic attention and offloading methods, but the selection of baselines does not cover some of the most recent quantization methods (e.g., PALU [9], which is only referenced).

Thank you for this kindness suggestion. We acknowledge that direct comparison with PALU requires more careful consideration due to the differences in methodology. PALU decomposes $\mathbf{K} \in \mathbb{R}^{l \times d}$ in dimension $d$, while our method and the comparison methods reduce both dimension $d$ and $l$. Following the your valuable suggestion, KVQuant is applied to both the vanilla model (meta-llama/Llama-3.1-8B-Instruct, Qwen/Qwen2.5-7B-Instruct) and the proposed method with default parameters. The results are shown in the table:

The suggested reference paper [2] is a great paper on applying low rank features on LLM training, which inspires us to explore the training side of the LLM. We will consider this in the future work. Thank you for this suggestion.

W2. While sensitivity to rank and top-k is shown (Tables 3, 4), there is limited discussion of the overhead and tuning effort required, particularly when considering different LLM architectures

We appreciate this concern and provide comprehensive overhead analysis:

Overhead and Tuning Efforts: Different model and dataset should have their own optimal hyperparameters. However, based on our experiments, there is an default choisece: $r=32$, top-$k$ = 2048, and 64 lite tokens. We can further tune the hyperparameters based on this default choice.

W3. Results are only shown for LLaMA-3-8B and Qwen2.5-7B (Tables 1-2), lacking validation on other popular architectures.

Thank you for this valuable feedback. We have extended our evaluation to include two additional popular models: mistralai/Mistral-7B-Instruct-v0.3 and microsoft/Phi-3-mini-128k-instruct. For these experiments, we used the default hyperparameters of LRQK. Due to memory constraints on the A6000 (48GB), the Phi-3-mini model is evaluated in 16K. The results are summarized in below.

W4. Please elaborate further on the differences with other methods focusing on efficient attention for long-context inference.

Thank you for this important suggestions. A comparison table highlighting key methodological differences is provided below:

Distinctions: 1. Joint decomposition, 2. Hierarchical caching, 3. Exact attention preservation.

W5. While the paper claims to balance trade-offs between memory efficiency, numerical precision, and computational latency, it lacks experimental validation of these specific trade-offs.

Thank you for the feedback. We evaluated the trade-offs on a text summarization task using 20 samples from LongBench (context: 32K, output: up to 128 tokens), with meta-llama/Llama-3.1-8B-Instruct. The time and resource usage are measured (Nvidia 3090 24G) and shown as below:

L1. Although presented algorithms are clear, actual running times, cache sizes (absolute MB/GB), and detailed hardware utilization are not given beyond high-level statements about GPU/CPU. Code is not referenced as available, and implementation details are not provided.

We acknowledge this limitation and will provide detailed implementation metrics:

• Cache in GPU: 1. The most recently accessed KV cache, 2. Low rank matrices $\mathbf{A}\_{K}, \mathbf{B}\_{K}, \mathbf{B}\_{Q}$, 3. The weights of LLM.
• Cache in CPU: 1. The full KV cache.
• Actual running times: Please refer to response W5.
• Implementation: A lightweight version of our code is provided in the supplementary material. Key aspects include:
  • Built on the standard Hugging Face transformers library.
  • The original attention mechanism is replaced with our proxy attention mechanism.
  • A self-defined buff session is introduced to manage the cache.
  • Two core algorithms are further compiled with torch.compile.

Thank the reviewer for insightful comments.

W1. Experiments list the hardware (A100) but omit tokens/s throughput, wall-clock latency, and peak/resident GPU memory, making it hard to judge real deployment benefits

Thank you for this valuable feedback. We provide comprehensive timing results on single NVIDIA A100 40GB with Llama-3-8B-1M. The following table shows the results. Further details can be found in the supplemental material.

W2. All results use LLaMA-3-8B-1M and Qwen-2.5-7B; the method’s scalability to 30 B + or MoE LLMs, where KV caches are a bigger bottleneck, remains unverified. The authors themselves flag that rank r and active-set size k require dataset/model-specific tuning, and optimal values can vary across tasks. I am concerned about the ease of adoption of this technique for more general use cases.

Thank you for this insightful question. We provide further validation on larger models: Qwen/Qwen2.5-14B-Instruct (64K context) and Qwen/Qwen2.5-32B-Instruct (16K context). YaRN is applied on 64K context. LRQK use default settings. Results are shown below.

The table shows that the proposed method can be applied to larger models.

W3. Equation 3 formulation is internally inconsistent. The hard constraints $Q=A\_Q B\_Q, K = A\_K B\_K$ in Eq. 3 cannot hold for low rank $r<d$ and are immediately relaxed in the Lagrangian; also, the interaction term omits $B\_Q B\_K^T$. This is mostly cosmetic but weakens the mathematical presentation.

We acknowledge the reviewer's concern and provide the following clarification: You are correct that the hard constraints $Q = A\_Q B\_Q$ and $K = A\_K B\_K$ cannot hold exactly when $r < d$. The formulation should indeed reflect that these are approximation constraints. We will revise the formulation to be more mathematically precise:

Regarding the interaction term omitting $B\_Q B\_K^T$: The key insight is that we are not approximating $Q K^T$ as $(A\_Q B\_Q)(A\_K B\_K)^T = A\_Q B\_Q B\_K^T A\_K^T$.

Instead, our approach recognizes that:

1. Both $Q$ and $K$ have low-rank structure individually
2. Their product $Q K^T$ can be approximated more efficiently by directly learning compact representations $A\_Q$ and $A\_K$ such that $Q K\^T \approx A\_Q A\_K^T$

We will revise the manuscript to clarify these mathematical details and ensure the formulation is both precise and well-motivated.

Q1. Section 5 reports “≈ 60 % reduction in CPU <-> GPU transfers” but gives no tokens-per-second or latency figures on the A100 node you mention (lines 210-216). Could you provide end-to-end throughput (tokens/s) and peak GPU memory for LRQK vs. ShadowKV, InfiniGen and vanilla attention at 128 K tokens?

Thank you for this important question. We provide a timing results on NVIDIA 3090 24GB under text summarization task (20 samples from LongBench, context 32K, output up to 128 tokens) with meta-llama/Llama-3.1-8B-Instruct as the base model. Due to time limit, we evaluate the vanilla attention (2 GPUs) and ShadowKV (1 GPU), and the proposed method (1 GPU). Average results are shown below:

Q2. All experiments use 7-8 B parameters. Are there bandwidth or PCIe limits that prevent LRQK from scaling to 30 B+ or mixture-of-experts LLMs where KV memory dominates? Even a profiling statement would help assess significance.

Bottle neck: PCIe limits and GPU bandwidth is not the main limitation, since the data need to be transfered between CPU and GPU is relatively small. Based on our observation, the main time consumption is the indexing operation in CPU. It takes time to fetch and collect hitted tokens in CPU. We are looking for a better method to reduce this time consumption.

Q3. Since LRQK already uses low-rank factors, can the retained full-precision KV rows be post-quantised to 8-bit or 4-bit (à la KVQuant) without harming accuracy? Discussing complementarity with quantisation would broaden impact.

Thank you so much for this sugessions, the quantization methods are valuable to explore. Following the your suggestion, KVQuant is evaluated to both the vanilla model (meta-llama/Llama-3.1-8B-Instruct and Qwen/Qwen2.5-7B-Instruct) and the model with proposed method. The results are shown in the following table. The hypterparameters of LRQK are $r=32$, top-$k$ = 2048, and 64 lite tokens.

From this table, we can find that in general applying KVQuant with the proposed method does not siginificantly harm the performance.

L1. The authors candidly notes that rank r and active-set size k require task-specific tuning and that optimal values differ across architectures. However, the societal-impact section is missing. Please add a brief discussion on potential misuse (e.g., enabling cheaper deployment of disinformation bots) and mitigation strategies.

Thank you so much for this valuable sugessions.

Hyperparameter Tuning: Different models and datasets may require unique set of optimal hyperparameters. However, there is a default choice that can help reduce the tuning effort. Based on our experiments, we can further turn from a default choice: $r=32$, top-$k$ = 2048, and 64 lite tokens. For instance, the rank $r$ can be chosen from $\{8, 16, 32, 48\}$, and the top-$k$ value can be adjusted to accommodate different context lengths.

Societal impact: LRQK's reduced inference costs could enable broader access to long-context language models, benefiting research and education. However, this efficiency may also lower barriers to deploying harmful applications such as large-scale disinformation campaigns. We believe democratized access to advanced AI capabilities provides net benefits when coupled with appropriate safeguards and responsible deployment practices.