SALS: Sparse Attention in Latent Space for KV Cache Compression

W1: Although the pipeline is robust, the overall methodological novelty is somewhat incremental: SALS combines two well-explored ideas (latent compression and sparse token selection) without introducing fundamentally new theory.

A: Thank you for your comments.

The latent compression is studied by many recent works such as Loki and Palu. The KV cache compression strategy and compression rate are not fully discussed, while the accuracy is hard to maintain. To achieve higher accuracy yet better KV cache compression rate, our work proposes the joint-head compression comparing to the existing works of per-head compression (Loki) and group-head compression (Palu). Joint-head has the advantage of preserving of the common principal components shared among different heads and thereby achieves better compression rate. Loki has not achieved any compression on KV cache whereas Palu achieves high compression rate (50%) but suffers from the accuracy loss (>5%). On the other hand, our work achieves high compression rate (25%）by adopting joint-heads compression. Furthermore, with the sparse attention, the reconstruction overhead is significantly reduced with maintained accuracy (loss gap <1%).

W2: The paper does not deeply investigate how SALS interacts with grouped-query attention (GQA) and multi-query attention (MQA) variants. Although a GQA-based model is included in the experimental suite, the efficiency analysis remains focused on the multi-head-attention (MHA) case and fails to articulate how the latent-space selection and sparse reconstruction behave under GQA/MQA. This omission leaves open the possibility that the reported performance gains may not generalize beyond the standard MHA setting.

A: We are sorry for the confusion caused. As shown in Figure 3, the token is selected based on the joint-head low-rank compressed latent space. For the token selection stage, the latent keys and quires both only have one head. As such, there is only one head to select the token index from latent key cache and these indexes are later shared by all the heads in the sparse attention. Therefore, for GQA/MQA, SALS works fine by sharing the critical tokens among the KV heads.

Our experiment results in Table 5 are also performed based on the GQA setting (32 query heads and 8 KV heads). We have shown clear speed-up (4.27X) comparing to the Flash-attention under the setting of batch size 16 and 16k sequence length.

Q1 It would strengthen the empirical foundation to report SALS’s performance alongside contemporary sparse-attention methods, especially dynamic selection methods like Quest or ArkVale.

W3 While throughput and memory gains are impressive, the experimental evaluation lacks comparisons to several leading sparse-attention frameworks, such as Quest, H2O, and ArkVale that also target long contexts.

A: Thank you for highlighting recent sparse attention baselines. We have included experimental comparisons to Quest and H2O, which are leading dynamic selection methods. The results are summarized in Table R1 below.

Table R1: Accuracy, latency, and memory access comparison with recent sparse attention baselines on LongBench.

We compare SALS with FlashAttention, Quest, and H2O in terms of average accuracy (on LongBench), attention operator latency (ms), and normalized memory access. SALS achieves competitive accuracy with significantly lower memory cost and latency

We also attempted to run ArkVale, but unfortunately its official implementation produced incorrect outputs without modification, and significant internal changes would be required to integrate it into our evaluation pipeline. Given the risk of introducing discrepancies, we chose not to include ArkVale results in this version.

That said, SALS differs from ArkVale in several key aspects (e.g., latent-space selection, fully shared selection heads), and we plan to revisit this comparison in future work with a stabilized version of ArkVale.

Q2 Although you include a GQA-based model, the efficiency analysis may remain tailored to standard multi-head attention. Providing benchmarks on grouped-query attention layers — reporting memory access and throughput — would clarify SALS’s behavior under these widely used architectures.

A: Thank you for the question regarding SALS’s behavior under Grouped-Query Attention (GQA) and Multi-Query Attention (MQA).

Our method is fully compatible with GQA/MQA architectures due to the shared token selection mechanism: token selection is performed in a joint latent space and is shared across all query heads. As a result, the number of memory accesses does not increase with the number of query heads, and the selection cost remains constant regardless of the GQA/MQA configuration.

Furthermore, as discussed in Section 4.5 of our paper, the memory access formula is independent of the number of attention heads, ensuring that the efficiency gains of SALS naturally extend to GQA and MQA architectures.

We provide additional runtime benchmarks in Table R2 below, which reports the attention operator latency and normalized memory access cost (FlashAttention = 1.0) under a GQA configuration (32 query heads, 8 KV heads, batch size = 16, sequence length = 4k). These results are consistent with the end-to-end throughput evaluation presented in Table 6 of our main paper and further demonstrate that SALS achieves strong efficiency improvements under GQA settings.

Table R2: Latency and relative memory‑access ratios (Flash‑Attention = 1.0) at batch = 16, sequence length = 4 k

Q3 Table 3 appears without narrative discussion, leaving its per-task metrics under-interpreted. A brief in-text analysis defining each column, highlighting key trends across Single-QA/Multi-QA/Summarization, and relating memory-access reductions to latency savings could improve clarity.

A: Thank you for the insightful suggestion. We agree that Table 3 would benefit from clearer narrative guidance. In the revised version, we will include a brief in-text analysis defining each column, highlighting trends across key task types (e.g., Single-QA, Multi-QA, Summarization), and explaining how reduced memory access contributes to improved efficiency. Specifically, we will add the following narrative:

Table 3 compares task-level performance across six LongBench categories, alongside normalized memory access costs. The first three columns reflect core reasoning tasks, where SALS achieves consistently strong performance. Notably, SALS-12.5% retains competitive accuracy while reducing memory access to just 6% of the baseline. This indicates that SALS preserves informative tokens more effectively. The memory-access savings translate into practical latency improvements, as fewer KV memory lookups reduce attention computation overhead.

To further support this, we will also include the Table R2 in the final version, which shows a clear correlation between memory-access reduction and latency savings across multiple methods:

W1. The value tensors remain in full rank. Despite the paper adopts quantization to compress value cache, it complicates the overall framework. Experimental results with low rank value caches should be provided.

A: Thank you for your suggestions. We have added the value compression results in the Table R1 below. After the PCA compression, 50% and 25% values are stored. We find that the accuracy drop when performing value compression is larger than the quantization of SALS. The high sensitivity of value cache is also observed in [1].

Table R1: Performance on LongBench under different value compression settings.

[1] Yixin Ji, Yang Xiang, Juntao Li, Qingrong Xia, Zi Ye, Xinyu Duan, Zhefeng Wang, Kehai Chen, and Min Zhang. 2024. Adaptive Feature-based Low-Rank Compression of Large Language Models via Bayesian Optimization. In Findings of the Association for Computational Linguistics: EMNLP 2024, pages 4152–4168, Miami, Florida, USA. Association for Computational Linguistics.

W2. A discussion between SALS and NSA [1] should be provided.

A: Thank you for your suggestions. We will add the following discussion in our revised paper.

Natively sparse attention (NSA) employs a dynamic hierarchical sparse strategy with coarse-grained token compression and fine-grained token selection. The major difference between NSA and SALS is that adopting NSA requires training from scratch, whereas SALS can be calibrated and achieved during the post-training phase. SALS shares the similar idea to perform the top-k token selection. However, NSA performs this token selection based on the trained compressed attention module. SALS performs the token selection based on the low-rank approximation.

W3. Incorporating comparisons of multiple value cache quantization methods (e.g. [2-3]) can further enhance the robustness of the SALS.

A: Thank you for your suggestions. We have added the ZipCache and GEAR quantization methods for comparisons using llama-7b-chat on the longbench dataset. As reported, GEAR, ZipCache and kivi quantization achieve similar accuracy under 4× (25%) compression rate, which illustrates the robustness of our SALS algorithms in incorporating different value cache quantization methods.

Table R2: Performance on LongBench under different value compression methods.

W1 The paper reports memory access as a primary metric, but this does not directly capture the efficiency of the proposed method. It would be more informative to report latency or throughput for SALS and compare it with baseline methods.

A: Thank you for the valuable suggestions. Table R1 reports attention‑operator latency and the corresponding relative memory‑access ratios at batch = 16, sequence length = 4 k.

As you pointed out, the latency of SALS does not reach its theoretical optimum for short sequences. Nevertheless, the measured latency mirrors the trend in the memory‑access ratio, and as the sequence length grows the fixed overhead becomes negligible; the latency ratio therefore converges toward the relative memory‑access ratio, as demonstrated in Table R2.

Table R1: Latency and relative memory‑access ratios (Flash‑Attention = 1.0) at batch = 16, sequence length = 4 k

W2 Most experiments use relatively short context lengths (1K, 2K, 4K as noted on line 301). It would strengthen the evaluation to include results on longer contexts, such as 32K tokens or more, which are more representative of long-context use cases.

A: Thank you for pointing out the importance of long sequences. We have extended our experiments from 8 k up to 128 k tokens; the new measurements are summarized in Table R2 below (batch = 16, A100 80 GB, FP16).

Table R2: Attention operator latency and relative latency ratios of SALS‑25 % / SALS‑12.5 % versus Flash‑Attention at batch = 16, A100 80GB, FP16.

At a 4 k sequence length, SALS delivers a 5.7 × speed-up in attention operator latency. However, when the sequence length increases to 128 k, the fixed-overhead fraction becomes negligible, and the attention operator under SALS-12.5% achieves a 10.44 × acceleration over FlashAttention.

This speed-up aligns with the theoretical latency reduction discussed in Section 4.5 of our paper: since the dominant cost is memory access, the acceleration factor approximates the reciprocal of the memory-access ratio. For SALS-12.5%, the ratio is 0.06, yielding a theoretical upper bound of ~16.6×, which closely matches the trend observed in practice.

W3 The contributions appear limited. The core idea of using low-rank projections to identify top-k tokens resembles existing methods like Loki [1]. The paper would benefit from a clearer explanation of what distinguishes SALS from these prior approaches.

A: Thank you for your comments. Our work significantly differs from Loki in three manifolds:

1) Latent Space with RoPE

Loki obtains its low-rank projection matrix by performing PCA on either the pre-RoPE or post-RoPE keys. Regardless of how the projection is computed, Loki ultimately applies the transformation to post-RoPE keys and queries, and retains full-rank representations. This enables approximate attention score computation for token selection, but does not perform actual compression.

As discussed in Section 3.1 of our paper, performing PCA after RoPE results in a significantly higher intrinsic rank, which leads to poor low-rank approximation accuracy. To address this, SALS applies low-rank projection before RoPE, enabling both accurate token selection and key compression. We also introduce a novel token selection mechanism in the pre-RoPE latent space, achieving better accuracy–latency trade-offs compared to Loki.

2) KV cache compression

As mentioned in the limitation part of Loki paper, Loki has not achieved reduced memory usage of the KV cache. The KV cache is not compressed and stored in the low-rank fashion. Loki only used the low-rank queries and keys to perform the approximate attention. In the appendix E of Loki paper, Loki also proposed PCA-Attn, which directly used the reduced-dimensional scores. In this way, the KV cache is compressed but PCAAttn performs poorly comparing to all the baselines as summarized in Table 5 (Loki paper).

In constrast, our work achieves KV cache 3.57× (1/0.28) compression rate as shown in Table 2 (SALS paper). Moreover, because of the KV cache compression, the memory access is also reduced compared to Loki.

3) Joint-head compression vs. per-head compression

Loki's strategy is to perform per-head low-rank compression. As discussed in Fig. 2 from the Loki paper, the compression rate is around 50% for most models (e.g., for Llama2-7B, head dimension is compressed from 128 to 80 for most layers).

Differently, our strategy is to perform joint-head compression, leading to more compression ratio. As shown in Table R2 of our paper, we can achieve 4× and 8× compression while maintaining the accuracy.

Summary: To achieve KV cache compression using the low-rank strategy brings new challenges to the existing method. Our work tackles this challenge by merging heads for compression and carefully integrating rotary embedding in the SALS architecture. We achieve 6.4× KV cache compression while maintaining the accuracy on 4K sequences compared to FlashAttention 2.

Reference: Singhania, Prajwal, et al. "Loki: Low-rank keys for efficient sparse attention." Advances in Neural Information Processing Systems 37 (2024): 16692-16723. https://arxiv.org/pdf/2406.02542

Q1: You project all multi-head keys into a shared latent space using a joint low-rank projection. Have you considered whether using per-head projections might improve token selection accuracy, despite the increased overhead?

A: Thank you for the insightful question. We explored this alternative and implemented a per‑head variant where each attention head is equipped with its own rank‑r projection matrix. To ensure a fair comparison, we kept the total projection rank constant: for h heads the per‑head rank is set to r / h, so both variants introduce the same number of projected dimensions. To ensure robustness, our results are averaged over 500 input samples and cover all attention layers.

Table R3:  Token‑selection alignment of joint vs per‑head low‑rank projections

To evaluate the token selection module, we report two metrics in Table R3: Overlap Score, which reflects the effectiveness of selection in capturing high-attention tokens (i.e., whether the selected tokens are valuable under full attention), and Overlap Rate, which measures the accuracy of selection by comparing against the top-k tokens in the full attention.

As shown in Table R3, joint-head projections outperform per-head projections. While joint projections may sacrifice per-head specificity, they better preserve the globally important directions in the low-rank space, as discussed in Lemma 1. This results in higher overall token selection accuracy.

In addition, as also observed in prior work (Palu), joint-head projections have the advantage of preserving the common principal components shared across heads, thereby achieving better compression efficiency under a constrained projection budget.

Q2: SALS is conceptually similar to Loki, which also performs low-rank key compression and token selection. Could you clarify the key differences in token scoring or reconstruction strategies between SALS and Loki?

A: Thank you for your comments. Besides the difference in the Latent space with RoPE, KV cache compression, and joint-head compression mentioned in the previous reply, the token scoring and reconstruction strategies are also different:

(1) Since Loki performs the compression on each head, the token scores are for each head and tokens are selected differently for each head. For our case, we perform the joint-head compression, where each head selects the same tokens.

(2) For the reconstruction strategies, Loki performs no KV cache compression, and thereby it can directly select the chosen KV cache. However, for our SALS, the selected token is stored in the compressed fashion and only reconstructed when it is selected.

Dear Reviewer CDzv,

I would like to sincerely thank you for the careful review, the valuable comments and constructive suggestions, which help to improve the quality of our paper. I have addressed your concerns in the response above, including: (1) adding latency evaluations and extending experiments to longer sequences to better demonstrate the efficiency of our method; (2) clarifying our rationale for using joint-head low-rank projections and demonstrating their advantage over per-head projections; and (3) clarifying the key differences between our approach and prior work such as Loki in terms of latent space with RoPE, KV cache compression, and token selection strategy.

Please feel free to let me know if you have any further questions or if anything remains unclear. I’d be happy to provide further clarification if needed.

Best regards,

the authors

Dear Reviewer CDzv, Thank you for reviewing this manuscript. Could you engage in the author-reviewer discussion after reading the author's response?

Dear Reviewer CDzv,

Thank you again for your review and for submitting the Mandatory Acknowledgement. I noticed that there had been no public reply following the acknowledgement, and I just wanted to check whether this might be due to a technical issue on OpenReview or some display limitation on our end. If there was a response that didn’t show up on our side, please kindly let us know.

As the discussion period is coming to a close, we remain available to provide further clarification on any aspects that may still be unclear. We sincerely value your feedback and appreciate the opportunity to engage in this discussion.

Best regards,

the authors

Q1 Can you provide numbers to justify Table 1?

A: Thank you for you suggestions. We have provide the KV data movment, memory size and computation complexity comparision equation. These equations are used to support the claim low/median/high marks.

Table R1: KV data movement, memory cost and complexity comparison of quantization, low rank and token sparse methods with fixed, dynamic and local token selection strategies. We use $b$, $n$, $n_q$, $s$, $d$ to represent the batch size, KV head number, Q head number, sequence length and head dimension, respectively. We further use $s'$ and $s_{fixed}$ for top-k selected token length and fixed pattern length, respectively, for sparse attention. $r$ represents the compressed rank of KV cache.

The accuracy gap is difficult to compared. For the Palu' accuracy, we take the experiments from Table 3 in our paper. StreammingLLM is based on Tabel 3 in Hshare [1] paper on the longbench dataset. For the rest of methods , the accuracy gap is taken from Table 4 in our paper.

[1] Wu, Huaijin, et al. "HShare: Fast LLM decoding by hierarchical key-value sharing." The Thirteenth International Conference on Learning Representations. 2025.

Q2 Additional experiments as stated in the weaknesses.

W1 The evaluation models (Llama2-7b and mistral) are even not long-context models (Llama3.1 is more preferred).

W2 The evaluation dataset is not standard. RULER, InfiniBench, and GSM-Infinite are more convincing.

A: Thank you for your suggestions. We have added benchmark results on the RULER dataset using LLaMA3.1-8B-Instruct. Table R2-a and Table R2-b present our main results on 4k and 32k sequence lengths. Following the LongBench setting, we use an overall sparsity of 1/8. The results demonstrate that SALS retains high performance under long-context conditions, confirming its robustness.

Due to the multimodal and video-based nature of InfiniBench and the lack of standardized implementations for GSM-Infinite, we focused our benchmark extension on RULER, which is fully text-based and compatible with our current evaluation pipeline. We plan to include GSM-Infinite and other long-context datasets in future evaluations as implementations mature.

We selected LLaMA3.1-8B-Instruct to better reflect practical long-context use. We report results at a sequence length of 32k, as even optimized attention implementations (e.g., FlashAttention) face memory constraints beyond this point on single-GPU setups. The 32k setting thus offers a meaningful and representative context length for evaluation.

We will include these extended results in the revised version of the paper.

Table R2-a: Performance comparison of baseline and SALS methods on the RULER dataset with Llama3.1-8B-Instruct (4k sequence length).

Table R2-b: Performance comparison of baseline and SALS methods on the RULER dataset with Llama3.1-8B-Instruct (32k sequence length).

Q3 InfiniGen, ShadowKV also employs low-rank compression for KV cache. Could you provide any insights about the core differences?

A: Thanks for your comments. Although our SALS share some similarities with InfiniGen and ShadowKV in employing low-rank compression for KV cache, there exist several core differences:

(1) The difference from InfiniGen. InfiniGen is built based on two observations: the attention input similarity and the skewed partial weights. Skewed partial weights is the most related to our work, however, skewed weight is designed to identify the large-value channels (outliers) in the queries and keys without considering the RoPE position embedding. As discussed in our work, RoPE leads to increasing rank. In the InfiniGen work, the skewing (low-rank) matrix is fused with weights, and the outliers in the queries and keys are element-wise added to select the top-30% channels for the approximated attention scores computations.

In our work, we directly adopt the low rank keys and queries to compute the approximated attention scores. We do not fuse the orthogonal matrix into weights as discussed in our work, the RoPE operation introduces new variances and increases rank. Another key difference is that we compress the keys and only store the low-rank key cache, whereas InfiniGen stores the full rank keys in the CPU memory and performs prefetching of keys to avoid long communication latency.

(2) The differences fromShadowKV. our method differs in two key aspects: the usage of low-rank compression and the strategy for critical token selection.

a) shadowKV performs SVD-based low-rank decomposition for each token during decoding, producing token-specific low-rank matrices (A, B). While this may offer finer-grained reconstruction, it introduces additional online compression overhead and increases memory usage. In contrast, our method applies offline PCA-based calibration to derive a shared low-rank projection matrix across all tokens, which significantly reduces runtime cost and simplifies the memory layout during inference.

b) ShadowKV selects a chunk of tokens (e.g. 8 tokens) and adopts chunk-level approximation (mean of each trunk) for selecting important tokens. However, such approximation is not accurate enough, therefore, a small set of outlier chunks also has to be stored. In our case, we perform token level approximation by utilizing a smaller rank than stored KV cache. Our selection is much simpler without the need to perform complex outlier selection and also achieve comparable accuracy.

The discussion above will be added in our revised manuscript.

L1 There are many limitations generally for sparse attention. For example, the speed up is limited by how much memory taken by KV cache. I think this paper does not discuss limitations in the work.

A: Thanks for your comments. We will add the following discussion about the limitations in the revised manuscript.

In this paper, we introduce the Sparse Attention in Latent Space (SALS) framework, which projects the KV cache into a compact latent space via low-rank projection and performs sparse token selection in this space for sparse attention. The benefits of adopting sparse attention is greatly depending on the sequence length. As shown in Table 5, when the sequence length is not long enough such as 1k, the sparse attention (latency 0.409ms) performs worse than the standard flash attention (latency 0.23ms). On the other hand, when the sequence length is 4k, sparse attention achieves 3.08X speed-up.

The current paper has not explored the trade-off for the KV cache compression rate and the accuracy of token selection. With consideration of the sequence length, there will be a sweet point to improve the compression rate without losing the accuracy. Moreover, the current paper considers a fixed compression rate and sparse ratio for all the layers. In the future work, we are interested to explore memory-performance-accuracy trade-offs on different layers.