Palu: KV-Cache Compression with Low-Rank Projection

Authors: Chi-Chih Chang, Wei-Cheng Lin, Chien-Yu Lin, Chong-Yan Chen, Yu-Fang Hu, Pei-Shuo Wang, Ning-Chi Huang, Luis Ceze, Mohamed S. Abdelfattah, Kai-Chiang Wu
Primary area: other topics in machine learning (i.e., none of the above)
Keywords: KV-Cache, Low-Rank Compression, Large Language Model

Abstract:
Post-training KV-Cache compression methods typically either sample a subset of effectual tokens or quantize the data into lower numerical bit width. However, these methods cannot exploit redundancy in the hidden dimension of the KV tenors. This paper presents a hidden dimension compression approach called Palu, a KV-Cache compression framework that utilizes low-rank projection to reduce inference-time LLM memory usage. Palu decomposes the linear layers into low-rank matrices, caches compressed intermediate states, and reconstructs the full keys and values on the fly. To improve accuracy, compression rate, and efficiency, Palu further encompasses (1) a medium-grained low-rank decomposition scheme, (2) an efficient rank search algorithm, (3) low-rank-aware quantization compatibility enhancements, and (4) an optimized GPU kernel with matrix fusion. Extensive experiments with popular LLMs show that Palu compresses KV-Cache by 50% while maintaining strong accuracy and delivering up to 1.89× speedup on the RoPE-based attention module. When combined with quantization, Palu’s
inherent quantization-friendly design yields small to negligible extra accuracy degradation while saving additional memory than quantization-only methods and achieving up to 2.91× speedup for the RoPE-based attention. Moreover, it maintains comparable or even better accuracy (up to 1.19 lower perplexity) compared to quantization-only methods. These results demonstrate Palu’s superior capability to effectively address the efficiency and memory challenges of LLM inference posed by KV-Cache. Our code is publicly available at: https://github.com/shadowpa0327/Palu.

Decision: accept (poster)
Reviewer scores: 6, 6, 6, 5

Reviews and discussion:
<thread>
creation date:
1730330651889

thread_type:
official_review

<thread_object>
creator: reviewer_pjhr

<content>
summary:
Palu is a new framework for compressing the key-value cache in large language models, reducing memory usage and improving inference speed. It uses low-rank projection, group-head low-rank decomposition, and automatic rank allocation to achieve up to 50% KV-cache compression with minimal accuracy loss. Palu also integrates quantization compatibility and optimized GPU kernels, resulting in significant speedups, especially for long-context scenarios. The framework outperforms state-of-the-art quantization methods and demonstrates effectiveness across various LLMs and benchmarks.

strengths:
- Novelty: unlike common quantization and pruning approaches, Palu uses low rank decomposition to lower the size of KV caches, which is more mathematically sound and doesn't require hardware changes.
- Significant compression is claimed, and it is to be expected that KV caches can be significantly compressed. However, I do have some concerns about the achieved compression which I cover in the weaknesses.
- Extensive empirical results spanning accuracy studies and speed-up analyses are provided.

weaknesses:
- While caching latent keys and values makes sense from the memory efficiency angle, it introduces significant computational overheads during decoding. Indeed, while the prefill phase can allow for GEMMs to occur at every layer, the decoding phase is all about GEMVs. However, in the palu framework, a full GEMM is required for to rematerialize each key and value matrix to be used in the attention module. This adds a lot of complexity, particularly when decoding a large output sequence length.
- What the paper really is doing is not KV cache compression, but really it is compressing the QKV WxA layer preceding attention by applying SVD on two thirds of the weight matrix. 
- The experimental setup doesn't make a lot of sense. It is mentioned that FlashAttention is not used for apples-to-apples comparison. But FlashAttention anyway is there to speed-up the prefill phase where the KV cache doesn't exist yet. Why does it then matter? Why are you including prefill latency when your solution is supposed to be addressing KV cache compression, which is used in the decoding phase?

questions:
- In Table 1, why does Llama2-7B have better perplexity and accuracy than Llama3-8B-Inst? There is definitely something wrong in the experimental setup.
- Rather than compressing weights, since KV cache compression is desired, can a method such as ESPACE [1] be used to project activations to a low rank directly?
[1] Sakr, Charbel, and Brucek Khailany. "ESPACE: Dimensionality Reduction of Activations for Model Compression." arXiv preprint arXiv:2410.05437 (2024).
</content>
</thread_object>

</thread>
<thread>
creation date:
1730380107545

thread_type:
official_review

<thread_object>
creator: reviewer_2i3g

<content>
summary:
The paper introduces Palu, a KV-Cache compression framework that reduces memory usage in LLM inference by applying low-rank projection to the hidden dimension. Palu decomposes linear layers into low-rank matrices, caching compressed intermediate states and reconstructing KV values on the fly. Key features include medium-grained low-rank decomposition, automated rank search, quantization compatibility enhancements, and a GPU-optimized matrix fusion kernel. Experiments show that Palu compresses KV-Cache by 50% and boosts attention module speed by up to 1.87×.

strengths:
S1: The method creatively combines SVD algorithms with KV-cache compression, optimizing speed and grouping different attention heads with varied compression rates, which is insightful.

S2: The experiments are comprehensive, measuring both the algorithm’s accuracy and speed.

S3: The figures and tables are clear, and the descriptions are precise.

weaknesses:
W1: In Section 3.1, the statement "However, this approach poses significant computational challenges during runtime that make it impractical for deployments" lacks experimental results and theoretical support, even though this claim is foundational to the proposed method.

W2: There is no theoretical proof for the accuracy of the proposed algorithm, such as bounded error analysis.

W3: The time complexity of the algorithm is not analyzed.

W4: In Section 4, the relationship between "compression rate" and "Avg. Bits" is unclear. Additionally, definitions of metrics such as "compression rate" are not specified.

W5: The code is not open-sourced, making reproducibility uncertain.

questions:
Q1: Corresponding to W1, provide theoretical or empirical evidence for the statement in Section 3.1 that " However, this approach poses significant computational challenges during runtime that make it impractical for deployments "

Q2: Corresponding to W2, add a theoretical derivation of the algorithm's error bounds.

Q3: Corresponding to W3, include an analysis of the algorithm's time complexity.

Q4: Corresponding to W4, clarify the relationship between "compression rate" and "Avg. Bits" in Section 4, along with definitions of related metrics.

Q5: It is recommended to open-source the code.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730542784500

thread_type:
official_review

<thread_object>
creator: reviewer_fsn9

<content>
summary:
The paper introduces Palu, a novel framework designed to compress the KV-cache of large language models (LLMs) using low-rank projection. Unlike existing techniques, such as quantization or token eviction, which may overlook redundancies in hidden dimensions, Palu effectively captures these latent efficiencies. Notably, it provides a post-training compression solution, in contrast to low-rank methods like GQA and MLA that require training from scratch. This approach allows Palu to strike a balance between memory efficiency, accuracy, and inference speed.

strengths:
1. **Post-Training Solution:** Unlike MLA, which requires retraining from scratch, Palu can be applied post-training. This versatility allows for the compression of existing models without additional, costly training cycles.
2. **Balanced Accuracy and Efficiency:** The group-head low-rank decomposition (G-LRD) balances accuracy and computational efficiency. It preserves essential information while reducing reconstruction overhead compared to per-head and joint-head decompositions.
3. **Automated Rank Allocation:** By using Fisher information to dynamically assign ranks to matrices based on their importance, Palu ensures compression without significant accuracy loss. This approach optimizes each layer's sensitivity to compression.
4. **Quantization-Friendly Design:** Palu addresses quantization-induced outliers with a Hadamard transformation, effectively supporting low-bit quantization. This further enhances compression without compromising accuracy, yielding competitive perplexity scores compared to quantization-only methods.

weaknesses:
1. **Inefficiency with Positional Encoding**: Palu requires reconstructing the original key cache when calculating positional encodings, which may reduce computational efficiency.
2. **Performance on Complex Tasks**: Palu-50% does not outperform lower-compression 2-bit quantized KV-Cache methods, particularly on complex tasks like Code and Synthetic. This limitation suggests that Palu’s effectiveness may decrease in more challenging scenarios, where high compression ratios can impact accuracy.

questions:
Is it possible to compress position embedding by similar way to remove the key cache reconstruction in attention computing? Or do some better tricks in attention computing rather than reconstructing key states?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730708265025

thread_type:
official_review

<thread_object>
creator: reviewer_12bo

<content>
summary:
This paper addresses the issue of high GPU memory consumption in KV Cache during LLM inference by proposing a novel approach that differs from traditional KV Cache compression methods (such as quantization and token-level compression). Specifically, it employs Singular Value Decomposition (SVD) to reduce the rank of the KV Cache matrix, thereby decreasing its memory usage. Building on this method, the paper further explores several enhancements, including: 1) merging multiple heads into a group for SVD to minimize precision loss; 2) adaptively adjusting the compression ratio based on the redundancy of information at different layers; and 3) integrating the Walsh-Hadamard transform to facilitate quantization of the compressed results.

strengths:
1.	The approach proposed in this paper, which utilizes Fisher information to allocate compression ratios for each linear layer, is both effective and scalable. It offers valuable guidance for other KV Cache and model parameter compression efforts.
2.	This paper implements a CUDA kernel to maximize throughput efficiency.
3.	This paper proposes and experimentally validates the compatibility of its optimization methods with traditional techniques such as quantization and LoRA fine-tuning, ensuring their practical applicability in real-world environments.
4.	This paper conducts comprehensive ablation experiments to validate the effectiveness of each optimization module individually.

weaknesses:
1.	The method proposed in this paper, which utilizes SVD for rank reduction of parameter matrices, is not novel, as it has been employed in various studies (e.g., Sharma et al., "The Truth is in There: Improving Reasoning in Language Models with Layer-Selective Rank Reduction," 2023). The corresponding optimization of KV Cache storage is consequently straightforward given the rank reduction of the parameter matrices.
2.	The paper does not provide open-source code, which raises concerns about the authenticity and reproducibility of their experimental results.
3.	The paper compares a limited number of related works; aside from the baseline, it only includes quantization-based methods such as Atom and KVQuant. However, token-level compression methods are also a key focus in the KV Cache compression field, which this paper completely overlooks. The authors should at least conduct comparative experiments and method combination studies with one of the state-of-the-art approaches in this area, such as H2O (Zhang et al., "H2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models," 2024).

questions:
1.	Suggestion: The authors might consider adding a mathematical derivation section to demonstrate that the SVD method has an upper bound on the accuracy error resulting from the rank reduction of parameter matrices. This would enhance the generalizability and applicability of the method across different models and datasets.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732002737704

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Global Rebuttal

comment:
## **1. Discussion on Palu’s Computation Overhead and Design Rationale**

Given that multiple reviewers raised concerns regarding Palu’s computational cost, particularly for RoPE-based attention, we provide a comprehensive explanation to address these comments and clarify the design rationale behind Palu.

**Memory-Bound Nature of LLM Decoding.** The decoding process in large language models (LLMs) is predominantly memory-bound [4]. A key optimization target in this paper—the attention mechanism—requires frequent retrieval of large key-value matrices and sequential computations. These operations require continuous data movement between global memory and compute units, which places significant pressure on memory bandwidth. This characteristic, combined with the relatively low arithmetic intensity of these operations, results in the underutilization of computational resources. As illustrated by the Roofline Model, performance in memory-bound workloads is **constrained primarily by memory traffic rather than computational throughput.** Also, the powerful **compute units are often underutilized because the data required for computations cannot quickly be supplied from off-chip memory.**

**Trading Compute for Memory Pressure Alleviation** Given memory-bound characteristics, a compression technique that introduces additional computation can be beneficial and effectively accelerate the decoding process by reducing memory bandwidth usage. A notable example is KV-Cache quantization techniques, which introduce additional computation (e.g., dequantization) for memory footprint reduction. In Palu, we build on similar insight from the Roofline model perspective to optimize decoding. By balancing additional computation with memory bandwidth savings, Palu reduces memory requirements while accelerating performance. As demonstrated in **Section 4.5 and Table R4**, **Palu achieves up to a 1.89x speedup over uncompressed baselines in RoPE-based attention.** Furthermore, **Palu’s quantization-friendly design enables up to 2.91x speedups and a 7.59x reduction in memory usage when combined with quantization.**

**Table R4**. Normalized Speedup for RoPE-based attention module. Speedup numbers are extracted from Figure 5 in our paper.
| Method|4K|8K|16K|32K|64K|
|-|-|-|-|-|-|
|KIVI-4-bit|0.80|1.18|1.49|1.68|1.89|
|Palu-50%|1.04|1.24|1.58 |1.76|1.89|
|Palu-50%-4-bit|1.87 | 2.42 | 2.56| 2.75| 2.91|

**Managing Computational Overhead for Reconstruction.** As the Reviewer pjHR notes, it is crucial to ensure the introduced computation is within the manageable level. To achieve this, we designed a novel medium-grained low-rank decomposition (G-LRD) that reduces FLOPs for reconstruction within a manageable level while preserving accuracy. With G-LRD, Palu could successfully deliver acceleration, as seen in Table R5 and [discussion thread with reviewer pJHR](https://openreview.net/forum?id=LWMS4pk2vK&noteId=8Zx8nJnUir).

Together, Palu effectively balances memory-compute trade-offs, achieving significant speedups and memory reductions in memory-bound LLM decoding. 

## **2. Difference between Palu and Prior SVD’s Literature and Design Challenge.**

As pointed out by reviewers, we clarify the fundamental distinction between Palu and weight decomposition works [1, 2, 3]. In the existing literature, SVD is commonly used to decompose parameter matrices into two low-rank matrix pairs. However, these methods focus on reducing parameters and FLOPs and do not explore opportunities to change the caching mechanism during LLM decoding fundamentally.

Palu leverages this overlooked opportunity by designing a compression framework for caching the low-rank latent representations. While this concept might appear conceptually straightforward, making it practically efficient is non-trivial, especially when dealing with RoPE-based attention. One of the critical challenges we address is minimizing the overhead associated with reconstructing the Key-Cache from the low-rank representations using the second decomposed matrix during inference. This led us to develop the group-head low-rank decomposition (G-LRD), which optimizes the decomposition granularity across attention heads. Unlike naive approaches like joint-head low-rank decomposition (J-LRD), which incur higher reconstruction costs and lead to non-negligible slowdowns, our proposed G-LRD achieves a 1.95× speedup at sequence length 16k (see Figure 6.), enabling efficient and practical deployment.

---
[1] The Truth is in There: Improving Reasoning in Language Models with Layer-Selective Rank Reduction. https://arxiv.org/abs/2312.13558

[2] Language model compression with weighted low-rank factorization. https://arxiv.org/abs/2207.00112

[3] Data-aware low-rank compression for large NLP models. https://openreview.net/forum?id=_sSHg203jSu

[4] LLM Inference Unveiled: Survey and Roofline Model Insights https://arxiv.org/abs/2402.16363
</content>
</thread_object>

</thread>
<thread>
creation date:
1732006480768

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response

comment:
We sincerely thank all the reviewers for their thoughtful feedback and for taking the time to read our paper. We have uploaded a revised manuscript with polished text, additional discussions, and new experiments based on your valuable suggestions. Below are the key updates included in the revision:
+ Latency evaluation of performing SVD at runtime (**Appendix I**)
+ Discussion of Token-Level related works (**Appendix J**)
+ Enhanced descriptions of the experimental setup and results (**Sections 4.1, 4.5, and Table 3**)
+ Discussion of Alternative Strategies for Deriving Low-Rank Projection Matrices (**Appendix K**)
+ Theoretical Error Bound Analysis (**Appendix L**)
+ Theoretical Derivation on Compute and Memory Complexity for Computing Attention When Applying Palu (**Appendix M**)

Also, we have uploaded our code implementation in the supplementary materials.
</content>
</thread_object>

</thread>
<thread>
creation date:
1733197108994

thread_type:
official_comment

<thread_object>
creator: reviewer_12bo

<content>
comment:
I would like to provide a summary of my current stance. Initially, I assigned a score of 5 to the paper, citing some unresolved issues that required further clarification. However, after reviewing the authors' detailed supplementary arguments and explanations, I feel that my concerns have been addressed, so I am inclined to raise my score to 6. This is my current position, reflecting a more favorable view of the paper, though I am still open to further considerations by the Area Chair regarding the final decision.
</content>
</thread_object>

</thread>
<thread>
creation date:
1733291950126

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Summary of Rebuttal (Many thanks to all reviewers and AC)

comment:
Dear Reviewers and AC,

We sincerely appreciate the opportunity to clarify and address the reviewers’ feedback on our submission, Palu. Below, we summarize the major concerns raised during the review period and outline how they were addressed, respectively.

**1. Discussion on Additional Compute Requirements:**
+ Reviewers pJHR and Fsn9 expressed concerns about the additional computation required when applying Palu to RoPE-based attention. 
+ In [global rebuttal](https://openreview.net/forum?id=LWMS4pk2vK&noteId=U4BX8kgVtS), we clarified the trade-off of computation for reduced memory footprint, supported by quantitative results in **Tables R3 and R6**. 


**2. Discussion on Difference to Prior SVD Methods for Weight Parameter Reduction:**
+ Reviewers pjHR and 12bo sought clarification on how Palu differs from prior SVD-based methods for weight parameter reduction. 
+ We provided detailed arguments in [global rebuttal](https://openreview.net/forum?id=LWMS4pk2vK&noteId=U4BX8kgVtS), highlighting Palu’s unique focus on KV-Cache compression, its fundamental change in the caching mechanism, and its specific technical contributions (e.g., manage reconstruction cost, automatic rank allocation, quantization compatibility, and efficient GPU kernel design).


**3. Discussion on Related Works**
+ Reviewer 12bo noted the need to discuss related work on token-level sparsity methods. 
+ We include a new section in Appendix covering the discussion of these methods and provided experiments (**Table R1**) demonstrating the orthogonality of Palu to these line of approaches.

**4. Including an additional Mathematical Derivation on Error Bound**
+ Both reviewers 12bo and 2i3G suggested that an error bound analysis on the proposed method would further strengthen the paper.
+ We addressed this by providing the derivation and proof in the **Appendix L**, as requested.

---
**Summary of Reviewer Outcomes**

By the end of the review period, we observed that:
+ Three reviewers (12bo, Fsn9, 2i3G) lean toward acceptance, with booth 12bo and Fsn9 raising their assessments from 5 to 6 and highlighting Palu as valuable work. 
+ Reviewer pJHR raised their score from 3 to 5, acknowledging the value of our work, while maintaining concerns about the trade-off between additional computation and memory footprint reduction we exploited in this work are only beneficial in the cherry-picking scenarios.

To address pJHR's remaining concern, we provided [additional summary](https://openreview.net/forum?id=LWMS4pk2vK&noteId=3QAWNiu7pe) and clarification demonstrating that our evaluations cover a large range of sequence lengths (4k~256k). We also provide benchmarking on end-to-end LLM decoding using popular LLM (e.g., Llama). We believe these results are comprehensive and genuinely show Palu can **practically accelerate LLM decoding in real-world usage**, not just for cherry-picking scenarios. 

**Final Remark** 

It is encouraging to see Palu recognized as a valuable contribution, with reviewers highlighting its **novelty and creativity (pjHR, 2i3G)**, **robust and general framework design (12bo, Fsn9)**, and **comprehensive evaluation (12bo)**. We sincerely express our gratitude to all reviewers and AC for their effort and time dedicated to reviewing the proposed Palu paper.
</content>
</thread_object>

</thread>


Metareview:
metareview:
Dear Authors,

Thank you for your valuable contribution to the ICLR and the ML community. Your submitted paper has undergone a rigorous review process, and I have carefully read and considered the feedback provided by the reviewers.

This paper introduces a new framework for compressing the KV cache in LLMs, which reducing memory usage and improves inference. The approach employs low-rank projections, group-head low-rank decompositions, and automated rank allocation to achieve up to 50% KV cache compression with minimal accuracy loss. Overall, the paper received mostly positive response from the reviewers (6,6,6,5) scores.

Given this positive assessment, I am willing to recommend the acceptance of your paper for publication.

I would like to remind you to carefully review the reviewer feedback and the resulting discussion. While most reviews were positive, the reviewers have offered valuable suggestions that can further strengthen the quality of the paper. Please take another careful look a the 'weaknesses' section of each reviewer comment. I encourage you to use this feedback to make any necessary improvements and refinements before submitting the final version of your paper.

Once again, thank you for submitting your work to ICLR.

Best,
Area Chair