Boosting Long-Context LLM Inference Efficiency with Intra-Layer Attention Similarity

Authors: Da Ma, Lu Chen, Situo Zhang, Yuxun Miao, Su Zhu, Zhi Chen, Hongshen Xu, Hanqi Li, Shuai Fan, Kai Yu
Primary area: foundation or frontier models, including LLMs
Keywords: Long context LLM, attention similarity, window attention

Abstract:
The increasing context window size in Large Language Models (LLMs), such as the GPT and LLaMA series, has improved their ability to tackle complex, long-text tasks, but at the cost of inference efficiency, particularly regarding memory and computational complexity. Existing methods, including selective token retention and window-based attention, improve efficiency but risk discarding important tokens needed for future text generation. In this paper, we propose an approach that enhances LLM efficiency without token loss by reducing the memory and computational load of less important tokens, rather than discarding them. 
    We address two challenges: 1) investigating the distribution of important tokens in the context, discovering recent tokens are more important than distant tokens in context, and 2) optimizing resources for distant tokens by sharing attention scores across layers. The experiments show that our method saves $35$% KV cache without compromising the performance.

Decision: reject
Reviewer scores: 5, 8, 3, 3

Reviews and discussion:
<thread>
creation date:
1730129779837

thread_type:
official_review

<thread_object>
creator: reviewer_ixy9

<content>
summary:
This paper introduces **POD** (Proximal Over Distant), a method to boost the efficiency of LLM inference with long contexts. 

It addresses the memory and computation bottlenecks of the decoding by leveraging two insights: (1) Proximal Tokens Importance, where initial and recent tokens are more critical than distant tokens, allowing prioritized computation for these tokens; and (2) Intra-Layer Attention Similarity, where attention patterns between consecutive layers are similar, enabling shared attention scores for distant tokens across layers. POD optimizes inference by grouping layers based on this similarity, using shared attention for distant tokens, and minimizing memory usage in KV caching. Experiments show that POD can reduce KV cache memory by 35% while maintaining comparable accuracy to baseline models, effectively optimizing long-context LLMs without compromising on essential token information.

strengths:
1. The paper creatively combines layer-similarity and post-trianing and proposes a novel approach **POD**, which improves inference efficiency by reducing memory requirements for less important distant tokens through inter-layer attention sharing, achieving a 35% reduction in KV cache while retaining all tokens for performance stability.

2. It provides evaluations across multiple long-context benchmarks and relevant ablation studies. **POD** outperforms baselines in most cases on provided experiments. In addition, the authors described the pipeline and details of **PoD** in a clear way.

weaknesses:
1. The paper didn't explain the motivation in a reasonable way. 

For example, (1) the observation 1 from Figure 1(a) is NOT really valid given the algorithm of StreamingLLM is to evict the tokens except for sink tokens (i.e. initial tokens) and most recent ones. In this case, "A equals B" inserted in the middle will be evicted and thus not used by StreamingLLM, which reasonably fails to return correct answer in the input. (2) The trend of similar attention scores shown in Figure 1(d) might not be immediately obvious to readers.

2. The setup of empirical results of the paper is relatively limited. 

For example, (1) the paper only used LLaMA-3-8B model in experiment section, it would be better to how this approach performs when the LLM scales up; (2) the paper should also include some other SOTA token-selection-based approaches such as Quest: https://arxiv.org/abs/2406.10774

3. The training stage of **PoD** may add more computational costs the application of the approach.

questions:
1. How does the approach apply to larger-scaled and up-to-date models such LLaMA-3-70B and LLaMA-3.1 models? 

2. Does the approach outperform other token-selection-based sparse attention mechanisms other than pure eviction-based baselines?

3. Can authors justify the weakness 1. and elaborate the observation from Figure 1(a)?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730358544745

thread_type:
official_review

<thread_object>
creator: reviewer_jnu7

<content>
summary:
The paper proposes the POD (Proximal over Distant) approach to improve inference efficiency in large language models (LLMs) by optimizing memory and computational resources for distant tokens instead of discarding them. It introduces a strategy of sharing attention scores across layers for distant tokens, based on the observation that proximal tokens (recent and initial tokens) are more crucial. By selectively reusing attention scores in deeper layers for distant tokens, POD achieves a 35% reduction in KV cache without significant performance loss. Extensive experiments demonstrate that POD maintains performance while enhancing efficiency across benchmarks.

strengths:
The model demonstrates no performance drop with POD, and its speed is improved compared to dense methods.

weaknesses:
1. The baseline methods for comparison are insufficient. Although the authors mention other acceleration methods in the Related Work section, these were not included in the experimental comparisons.
2. Compared to other training-free acceleration methods, this approach requires a continued training phase, adding computational cost.
3. Although not essential, it remains unclear how this method performs with longer context lengths, such as the 128k length achieved by LLaMA 3.1.

questions:
I do not agree with the assumption that middle tokens are less important than edge tokens. This contradicts our prior knowledge, as the model may need to use information from any position during prediction. LM-Infinite and similar works have highlighted this issue, yet this is a limitation of the model rather than an optimization direction. Additionally, many papers achieve full accuracy in "needle-in-a-haystack" experiments, indicating that well-trained long-text models can overcome the "lost in the middle" issue.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730676500964

thread_type:
official_review

<thread_object>
creator: reviewer_h78a

<content>
summary:
This paper proposes a method to reduce the size of KV cache to improve memory efficiency for long-context LLM. The proposed method involves sharing KV cache for “distant” tokens across layers, and continued pre-training the LM to adapt to this inference paradigm. Experiments are conducted with an extended LLaMA-3-8B model and evaluated on NIAH and two downstream tasks (LongBench and L-Eval). Results show that the proposed method saves 35% of KV cache while retaining performance, compared to previously proposed KV cache eviction method (H2O and StreamingLLM).

strengths:
* This paper aims to improve efficiency for long-context language models, which is a practical and important problem.
* The paper is written clearly, with comprehensive analysis and experiments.

weaknesses:
* **Baseline set-up**: The proposed method involves continuing pre-training the model to adapt to the paradigm of KV cache sharing across layers, yet all of the baselines (StreamingLLM, H2O) except for the “window attention”, are inference-time methods, making the comparison a bit unfair. It would be good to show what is the performance of adopting the proposed method as an inference-time method.
* **Experiment results**: The paper reported performance for all methods in Table 1 but only reported memory footprint of the proposed method (PoD) in table 2. What is the memory footprint saving for the baseline methods (StreamingLLM, etc.), and what is the performance-efficiency trade-off for different methods?

questions:
The proposed method involves using the attention scores to set the gate for combining the attention output for the distant and proximal tokens. However, FlashAttention does not explicitly write out attention scores during the attention computation. I am curious how this is handled on the implementation side?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730693371488

thread_type:
official_review

<thread_object>
creator: reviewer_k46j

<content>
summary:
This paper proposes POD, a method to improve LLM inference efficiency by differentiating between proximal tokens (recent and initial tokens) and distant tokens. Instead of discarding tokens like previous approaches, POD shares attention scores across layers for distant tokens while keeping full computation for proximal tokens. The authors show their method can save 35% KV cache without significant performance degradation on tasks like LongBench and LEval.

strengths:
- The work studies an important problem in LLM deployment, as handling long contexts efficiently is becoming increasingly crucial. 
- The method is relatively straightforward to implement and requires minimal adaptation of existing models. As shown in Section 2.2, it mainly involves grouping similar layers and modifying the attention computation, without needing extensive finetuning or architectural changes.

weaknesses:
- The paper misses comparison with some important recent baselines, particularly SnapKV [1] and PyramidKV [2], which also address KV cache optimization. While the related work section mentions them briefly, not including them in the experimental comparison makes it difficult to assess the relative advantages of POD.
- The evaluation seems limited in scope given the current state of the field. The authors only test on LLaMA3-8B with 32K context, while recent models routinely handle 128K tokens, e.g., Llama 3.1 and 3.2, Mistral Nemo, Phi 3.5, and so on. This raises questions about how POD would scale to longer contexts and whether the benefits would hold at larger scales.
- Also, the methods were only evaluated on LLaMA3-8B. It is unknown whether the method only works for this specific model or whether the method can be generally applied to most LLMs. 
- The implementation details raise concerns about compatibility with modern attention optimizations. The head-wise different grouping of layers (shown in Figure 2) suggests that each attention head would need different attention patterns, which may make it incompatible with efficient implementations like FlashAttention. Have you explored the compatibility of POD with FlashAttention or similar optimized attention implementations? 
- The memory savings claims could be better substantiated. While the paper reports 35% KV cache savings, Table 2 shows somewhat inconsistent practical gains across different batch sizes, and there's limited analysis of the overhead introduced by maintaining separate attention patterns per head. Could you clarify the computational overhead of maintaining different attention patterns for each head? 

[1] Yuhong Li, Yingbing Huang, Bowen Yang, Bharat Venkitesh, Acyr Locatelli, Hanchen Ye, Tianle Cai, Patrick Lewis, and Deming Chen. Snapkv: Llm knows what you are looking for before generation. arXiv preprint arXiv:2404.14469, 2024.
[2] Yichi Zhang, Bofei Gao, Tianyu Liu, Keming Lu, Wayne Xiong, Yue Dong, Baobao Chang, Junjie Hu, Wen Xiao, et al. Pyramidkv: Dynamic kv cache compression based on pyramidal information funneling. arXiv preprint arXiv:2406.02069, 2024.

questions:
See the Weaknesses above.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732530368676

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Response to all reviewers (Part 2)

comment:
**Table 1: Evaluation results of different methods on LongBench [1: dense; 2-4: token-selection-based methods; 5-9: token-eviction-based methods; 10-12: layer-sharing-based methods]**

| Order | Model                  | Window      | Avg.  | SQA   | MQA   | Summ  | Few-Shot | Code  |
| ---------- | --------------------------- | ---------------- | ---------- | ---------- | ---------- | ---------- | ------------- | ---------- |
| 1     | LLaMA3-8B-32K          | 32K         | 45.28 | 32.94 | 32.23 | 25.41 | 69.30    | 66.54 |
| 2     | SnapKV                 | 4K          | 44.17 | 31.76 | 31.85 | 21.92 | 68.62    | 66.72 |
| 3     | PyramidKV              | 4K          | 44.78 | 33.34 | 31.51 | 23.76 | 68.91    | 66.36 |
| 4     | Quest                  | 4K          | 44.82 | 32.14 | 32.19 | 24.27 | 69.05    | 66.43 |
| 5     | LM-Infinite            | 4+4092      | 42.84 | 28.83 | 28.95 | 21.74 | 68.12    | 66.54 |
| 6     | StreamingLLM           | 4+4092      | 42.80 | 28.68 | 28.95 | 21.64 | 68.14    | 66.6  |
| 7     | H2O                    | 96+4000     | 43.24 | 29.36 | 29.51 | 22.73 | 68.45    | 66.17 |
| 8     | Window Attention       | 4K          | 14.76 | 8.90  | 3.63  | 9.05  | 11.13    | 41.08 |
| 9     | Window Attention + CPT | 4K          | 41.98 | 26.94 | 27.95 | 22.29 | 66.60    | 66.10 |
| 10    | CLA                    | 32K         | 37.87 | 24.02 | 22.58 | 22.50 | 60.92    | 59.35 |
| 11    | PoD (ours)             | 16+4080+28K | 44.76 | 30.97 | 32.43 | 24.82 | 67.30    | 68.26 |
| 12    | PoD+SnapKV(ours)       | 4K          | 44.25 | 30.98 | 32.68 | 22.90 | 66.90    | 67.79 |

---

**Table 2: Evaluation results of different methods on LEval [1: dense; 2-4: token-selection-based methods; 5-9: token-eviction-based methods; 10-12: layer-sharing-based methods]**

| Order | Model                  | Window      | Avg.  | Closed | QA    | Summ  |
| ---------- | --------------------------- | ---------------- | ---------- | ----------- | ---------- | ---------- |
| 1     | LLaMA3-8B-32K          | 32K         | 27.44 | 42.10  | 24.68 | 15.55 |
| 2     | SnapKV                 | 4K          | 25.76 | 39.86  | 23.90 | 13.53 |
| 3     | PyramidKV              | 4K          | 25.90 | 42.10  | 22.63 | 12.96 |
| 4     | Quest                  | 4K          | 26.94 | 40.55  | 25.59 | 14.69 |
| 5     | LM-Infinite            | 4+4092      | 24.68 | 37.32  | 22.8  | 13.91 |
| 6     | StreamingLLM           | 4+4092      | 24.57 | 37.12  | 22.79 | 13.81 |
| 7     | H2O                    | 96+4000     | 24.60 | 37.15  | 23.18 | 13.48 |
| 8     | Window Attention       | 4K          | 9.77  | 20.95  | 5.57  | 2.79  |
| 9     | Window Attention + CPT | 4K          | 22.53 | 32.94  | 22.09 | 12.55 |
| 10    | CLA                    | 32K         | 14.70 | 19.05  | 13.52 | 11.52 |
| 11    | PoD (ours)             | 16+4080+28K | 27.18 | 43.59  | 22.95 | 15.00 |
| 12    | PoD+SnapKV(ours)       | 4K          | 26.50 | 43.07  | 22.12 | 14.32 |
</content>
</thread_object>

</thread>
<thread>
creation date:
1732530411436

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Response to all reviewers (Part 1)

comment:
Thank you to all the reviewers for the effort you put into this work. Your suggestions have truly helped improve our work.

One of the concerns you raised was **the lack of baselines and unclear comparisons**. In response, we have further included the baselines you suggested and provided a detailed discussion between these baselines and our approach. In particular, we categorize these methods into three groups:

1. **Token-selection-based methods**: SnapKV [1], PyramidKV [2], Quest [3]
2. **Token-eviction-based methods**: LM-Infinite [4], StreamingLLM [5], H2O [6], Window Attention [7]
3. **Layer-sharing-based methods**: CLA [8], PoD (ours)

We have supplemented the results in Tables 1 and 2 (**part 2**). From these, we can conclude the following:

- Token-eviction-based methods (Lines 5-9 vs. Line 1) degrade performance.
- Both token-selection-based methods (Lines 2-4 vs. Line 1) and our method (PoD, a layer-sharing-based method, Line 11 vs. Line 1) maintain performance. Furthermore, _our method is orthogonal to token-selection-based methods_. Specifically, by integrating the typical token-selection-based method SnapKV with our PoD, we obtain a more efficient model (saving KV cache at both token and layer levels) that maintains performance (Line 12 vs. Line 1).
- Our method outperforms the classical layer-sharing-based method CLA (Line 11 vs. Line 10). When post-trained with the same amount of data, CLA significantly degrades performance (Line 10 vs. Line 1). In the original work, CLA was pretrained from scratch. We attribute the performance degradation to the fact that a pretrained LLM cannot effectively adapt to CLA with a relatively small amount of data.

We have updated the results in Table 1, as well as the comparison of all methods, which can be found in the highlighted lines 297-315 and 355-362 of the revised version.

---

[1] Li, Yuhong, et al. "Snapkv: Llm knows what you are looking for before generation." _arXiv preprint arXiv:2404.14469_ (2024).

[2] Cai, Zefan, et al. "Pyramidkv: Dynamic kv cache compression based on pyramidal information funneling." _arXiv preprint arXiv:2406.02069_ (2024).

[3] Tang, Jiaming, et al. "Quest: Query-Aware Sparsity for Efficient Long-Context LLM Inference." _arXiv preprint arXiv:2406.10774_ (2024).

[4] Han, Chi, et al. "LM-Infinite: Zero-Shot Extreme Length Generalization for Large Language Models." _Proceedings of the 2024 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies (Volume 1: Long Papers)_. 2024.

[5] Xiao, Guangxuan, et al. "Efficient streaming language models with attention sinks." _arXiv preprint arXiv:2309.17453_ (2023).

[6] Zhang, Zhenyu, et al. "H2o: Heavy-hitter oracle for efficient generative inference of large language models." _Advances in Neural Information Processing Systems_ 36 (2023): 34661-34710.

[7] Beltagy, Iz, Matthew E. Peters, and Arman Cohan. "Longformer: The long-document transformer." _arXiv preprint arXiv:2004.05150_ (2020).

[8] Brandon, William, et al. "Reducing Transformer Key-Value Cache Size with Cross-Layer Attention." _arXiv preprint arXiv:2405.12981_ (2024).
</content>
</thread_object>

</thread>
<thread>
creation date:
1732530706204

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Response to Reviewers K46j, jNu7 and iXy9 (part 2)

comment:
**Table 3: Evaluation results of different methods over 7 sub-tasks in the InfiniteBench under different context sizes** **[1: dense; 2-4: token-selection-based methods; 5-9: token-eviction-based methods; 10-12: layer-sharing-based methods]. OOM: out of memory over one A800-80G GPU**

| Order | Model                  | Window          | Avg.  | 32K   |       |        |           |            |              |             | 64K   |       |        |           |            |              |             | 128K  |       |        |           |            |              |             |
|-------|------------------------|-----------------|-------|-------|-------|--------|-----------|------------|--------------|-------------|-------|-------|--------|-----------|------------|--------------|-------------|-------|-------|--------|-----------|------------|--------------|-------------|
|       |                        |                 |       | En.QA | En.MC | En.Sum | Math.Find | Code.Debug | Retr.PassKey | Retr.Number | En.QA | En.MC | En.Sum | Math.Find | Code.Debug | Retr.PassKey | Retr.Number | En.QA | En.MC | En.Sum | Math.Find | Code.Debug | Retr.PassKey | Retr.Number |
| 1     | LLaMA3.1-8B-128K       | 128K            | 36.37 | 29.56 | 25.33 | 15.51  | 32.29     | 27.41      | 27.12        | 27.12       | 37.26 | 28.82 | 15.14  | 28.00     | 27.92      | 54.24        | 54.24       | 31.24 | 41.78 | 14.86  | 21.14     | 25.38      | 100.00       | 99.49       |
| 2     | SnapKV                 | 4K              | 35.43 | 28.93 | 26.20 | 13.26  | 32.28     | 27.41      | 27.12        | 26.61       | 35.14 | 27.95 | 11.69  | 28.00     | 27.92      | 54.24        | 53.05       | 29.34 | 34.50 | 13.65  | 21.14     | 26.65      | 100.00       | 98.98       |
| 3     | PyramidKV              | 4K              | OOM*   | 29.88 | 26.20 | 15.01  | 32.29     | 27.39      | 27.12        | 27.12       | 36.53 | 29.26 | 15.18  | 28.00     | 27.92      | 54.24        | 54.24       | OOM   |       |        |           |            |              |             |
| 4     | Quest                  | 4K              | 35.17 | 28.49 | 25.76 | 12.87  | 32.29     | 27.41      | 27.12        | 27.12       | 35.81 | 27.51 | 11.22  | 27.89     | 27.90      | 54.24        | 54.24       | 29.04 | 34.50 | 8.95   | 21.14     | 26.90      | 100.00       | 98.14       |
| 5     | LM-Infinite            | 4+4092          | 18.08 | 25.28 | 26.64 | 12.46  | 32.11     | 27.31      | 3.39         | 3.39        | 27.25 | 28.82 | 12.64  | 26.89     | 27.91      | 3.39         | 3.39        | 22.19 | 29.26 | 13.01  | 21.14     | 26.63      | 3.39         | 3.22        |
| 6     | StreamingLLM           | 4+4092          | 18.05 | 25.39 | 26.20 | 12.05  | 32.29     | 27.41      | 3.39         | 3.39        | 27.14 | 27.95 | 13.29  | 28.00     | 27.92      | 3.39         | 3.39        | 22.66 | 27.95 | 12.79  | 21.14     | 26.65      | 3.39         | 3.22        |
| 7     | H2O                    | 96+4000         | OOM*   | 25.42 | 26.20 | 13.08  | 32.29     | 27.41      | 4.92         | 3.39        | 27.55 | 27.95 | 14.31  | 28.00     | 27.91      | 7.12         | 3.56        | OOM   |       |        |           |            |              |             |
| 8     | Window Attention       | 4K              | 2.69  | 3.52  | 3.49  | 0.53   | 0.00      | 10.41      | 0.00         | 1.36        | 3.36  | 3.06  | 0.65   | 0.00      | 10.66      | 0.00         | 1.19        | 3.32  | 3.49  | 0.70   | 0.00      | 9.64       | 0.00         | 1.19        |
| 9     | Window Attention + CPT | 4K              | 13.77 | 12.59 | 18.78 | 11.31  | 18.57     | 28.43      | 3.39         | 3.39        | 13.12 | 17.90 | 10.78  | 17.71     | 28.68      | 3.39         | 3.39        | 12.81 | 20.96 | 11.26  | 17.71     | 28.17      | 3.39         | 3.39        |
| 10    | CLA                    | 32K             | 32.20 | 22.07 | 34.06 | 13.58  | 12.57     | 27.16      | 24.07        | 25.76       | 22.74 | 31.88 | 12.74  | 12.57     | 27.66      | 50.85        | 52.54       | 21.62 | 34.50 | 12.96  | 12.86     | 29.44      | 97.80        | 96.78       |
| 11    | PoD (ours)             | 16+4080+28K     | 36.22 | 27.43 | 38.43 | 17.89  | 23.68     | 25.61      | 26.61        | 27.12       | 29.62 | 41.92 | 15.52  | 20.78     | 26.29      | 53.73        | 54.24       | 26.64 | 42.79 | 15.52  | 22.16     | 25.76      | 99.83        | 99.15       |
| 12    | PoD+SnapKV (ours)      | 16+4080+28K; 4K | 34.77 | 28.61 | 38.86 | 13.90  | 23.68     | 25.61      | 26.61        | 23.39       | 28.48 | 43.23 | 14.03  | 20.78     | 26.29      | 53.73        | 49.49       | 24.72 | 40.61 | 12.36  | 22.16     | 25.76      | 99.83        | 87.97       |


\* ***We run the official code.***
</content>
</thread_object>

</thread>
<thread>
creation date:
1732530914722

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Response to Reviewers K46j, jNu7 and iXy9 (part 1)

comment:
**128K context over LLaMA3.1-8B** To explore the generality of our method, we conducted experiments on LLaMA3.1-8B, which can handle longer (128K) contexts. We sampled 5B tokens from the ProLong-data-512K [1] dataset and applied the same hyperparameter configuration used for training LLaMA3-8B-32K to post-train LLaMA3.1-8B with a sequence length of 128K. The evaluation results over 4 practical sub-tasks in the InfiniteBench [2] under different context sizes are shown in the following Table 3 (**part 2**).

Consistent with the conclusions found in Tables 1 and 2 (Section **Response to all reviewers (part 2)**), our method causes less performance degradation compared to token-eviction-based methods. However, a notable difference is that token-selection-based methods appear to struggle in maintaining model performance in longer context scenarios. This limitation is also reflected in the combined model (PoD+SnapKV), which integrates our method with token-selection-based methods, showing a decline in performance. This to some extent indicates that our method is more robust to the context length.

We have updated the results in Table 4, as well as the conclusion, which can be found in the highlighted lines 422-446 of the revised version.

[1] Gao, Tianyu, et al. "How to train long-context language models (effectively)." _arXiv preprint arXiv:2410.02660_ (2024).

[2] Zhang, Xinrong, et al. "∞ Bench: Extending long context evaluation beyond 100k tokens." _Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)_. 2024.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732531038928

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Response to Reviewers jNu7 and iXy9

comment:
**Training cost of PoD** To adapt a pre-trained LLM to PoD, additional post-training is indeed required. However, this post-training is very lightweight and efficient, as evidenced by two key aspects:

1. The amount of data required for post-training (5B tokens) constitutes **less than 0.04%** of the pretraining data (over 15T tokens).
2. Compared to the classical layer-sharing-based method CLA, our method PoD can restore performance comparable to the dense model using only 5B tokens, whereas CLA's performance still lags significantly behind the dense model (Lines 1, 10, and 11 of Tables 1 and 2 in Section **Response to all reviewers (part 2)**).
</content>
</thread_object>

</thread>


Metareview: