Reducing Transformer Key-Value Cache Size with Cross-Layer Attention

Summary

To resolve the costly memory consumption of KV cache in transformer-based LLMs, this paper proposes Cross-Layer Attention (CLA). In contrast to multi-query attention (MQA) and grouped-query attention (GQA) which share KV cache across attention heads, CLA shares KV cache across contiguous layers. Pre-training experiments prove the effectiveness of CLA over MQA and GQA.

Strengths

1. The experiments are comprehensive and the design choices are clearly documented. 2. The problem studied is an important one and the approach is novel.

Weaknesses

1. Experimental results are not presented fully. Figure 3 only shows the validation perplexity for a subset of models, and Table 3 and 4 only presents the accuracy for a subset of models. I couldn't find the results elsewhere in the paper or the appendix. Presenting the results fully contributes to the holistic understanding of the effectiveness of the proposed approach. 2. Suggestions to experimental designs. The authors choose to pre-train 1B and 3B models from scratch to prove the effectiveness of CLA. It would perhaps be more convincing if the authors can choose a fully open source pre-trained model (such as OpenLLaMA) as baseline, and pre-train it with CLA to prove its effectiveness. Pre-training 1B and 3B models using a custom set of hyper-parameters and datasets is less convincing in comparison. 3. Training loss curves are missing. Training loss is important for understanding whether the models have fully converged. 4. Citation formats may be incorrect. Many references on page 10-12 only show author names, title, and year, and omit publication venue. 5. Some typos. Line 16: "which are possible" -> "which are impossible", line 18: "would otherwise be possible" -> "would otherwise be impossible".

Questions

1. Can the authors please provide some insights on why cross-layer sharing of KV cache may produce better models than GQA? The proposed method would be more convincing if some design justifications are provided. 2. Is it possible to adapt pre-trained models to using CLA through fine-tuning?

Rating

6

Confidence

3

Soundness

3

Presentation

2

Contribution

2

Limitations

Yes the authors addressed the limitations. I do not envision additional limitations.

Summary

In this paper, the authors proposed an approach, Cross-Layer Attention (CLA), to accelerate the autoregressive generation procedure of Large Language Models. Going beyond multi-query attention, the main idea of CLA is to share the Key-Value caches among attention layers. Intuitively, this idea is straightforward to further reduce the size of the KV cache that must be stored during generation. The authors conducted several experiments covering different model sizes and combinations between CLA and GQA/MQA to verify the effectiveness of the CLA module.

Strengths

1. The problem this paper aims to tackle has great significance in real-world applications of LLMs. 2. Overall, the proposed CLA approach is simple and easy to implement, which is friendly for practical engineering. The paper is also easy to follow. 3. **Efficiency**: from the experimental results, the CLA module can achieve further acceleration beyond MQA/GQA approaches across different model sizes. 4. **Performance**: it is interesting to see that the cross-layer sharing of KV caches does not hurt the performance of LLMs a lot, which further serves as support evidence that the CLA approach can be applied in practice.

Weaknesses

1. **The comprehensiveness of experiments can be further improved.** - **Model sizes**: the current version of this paper only verifies the effectiveness of CLA on 1B and 3B Transformer language models. Although the experimental results demonstrate that CLA does not bring performance drop on real-world tasks for these two model sizes, it is also questionable whether modern LLMs with >10 and > 100 billion parameters could still use this approach to achieve better accuracy-efficiency. After all, there exist many approaches that are only effective on small-scale architectures in the literature. It is also understandable that the computational resources are restrictions for conducting such verification experiments, but I think it would be a great plus for improving the quality of this work. - **Design choices**: in Table 1 and Appendix B, the authors listed a bunch of Non-CLA baselines and MQA/GQA+CLA2/CLA3/CLA4 models. However, CLA-only models are not compared. It would be questionable whether CLA can only be used with MQA/GQA together.

Questions

1. In lines 139-141, could you explain why using separately-learnable affine layer-norm parameters for the KV projection blocks and Q blocks in attention? 2. Similar to Sec 3.2.2, how is the robustness of CLA against other hyper-parameters like batch size?

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors discussed limitations of this work in Sec 4.

Summary

This paper introduces a new KV cache compression technique by sharing KV cache cross-layers called CLA. It can be integrated with most of the existing KV cache compression technique like MQA, GQA, and quantization. When applying over GQA, CLA can further reduces KV cache size by 2× while maintaining similar accuracy. Experiments with 1B- and 3B-parameter models show that CLA offers better memory/accuracy trade-offs, allowing models to handle longer sequences and larger batch sizes.

Strengths

- Simple and effective idea. It introduces an orthogonal dimension to the existing KV Cache compression methods. It can easily be integrated with existing KV cache techniques, like MQA/GQA, quantization, token eviction. - The writing and presentation is easy to follow - The proposed Cross-Layer Attention (CLA) does not require custom implementation and easily be deployed for any end-devices.

Weaknesses

I personally like this paper and vote for acceptance. However, I think the experimental design could be improved to make it more solid: - Only Wikitext Perplexity is reported, which is a weak indicator of LLM ability. While I understand that LLM evaluation is challenging and time-consuming, I suggest the authors also report the MMLU score, even if only on a subset. - Can CLA work under continuous pretraining scenarios like GQA? Namely, can a pretrained LLM be continuously tuned into a CLA-based model, even though the compression ratio might not be as high in this case? I ask this because pretraining LLMs is extremely expensive, as seen with models like Llama 3.1. It is risky and costly to train another open-sourced CLA model from scratch. If we can adapt existing models into CLA-based models, this paper could have a much greater impact. - One of the main motivations for KV Cache compression is long-context inference. However, long-context tasks like few-shot learning or multi-doc QA are much harder than the "short" tasks. I understand that long-context evaluation is challenging for pretrained models due to lack of alignment. Still, I think the authors could show some passkey retrieval results, as this task does not require much alignments.

Questions

Check Weaknesses point 2

Rating

8

Confidence

5

Soundness

3

Presentation

3

Contribution

4

Limitations

No major limitations.