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

Unreasonable motivation

• Invalid observation 1: We apologize for the misunderstanding caused by a typo. Indeed, Observation 1—"Proximal tokens (initial tokens + recent tokens) are more important than distant tokens"—is based on Figure 1-(c), not Figure 1-(a). Figure 1-(c) demonstrates that, even when attending to only the 256 proximal tokens, the model predicts the next token identically to the model attending to all tokens in 80% of the cases for the same input sequence. This phenomenon proves the observation 1. We have corrected the typo, which is highlighted in lines 80-81 of the revised version.

• Not obvious to the readers: To help readers directly grasp the similarity of attention scores between consecutive layers, we made following two adjustments to Figure 1-(d):

  • Replace it with a head that more clearly demonstrates the phenomenon
  • Reduce the number of colors
  • Distinguish only between values above and below the threshold (0.5) We have updated Figure 1-(d) in the revised version.

Limited baselines Please refer to Section Response to all reviewers.

Scaling of long context and other LLMs Please refer to Section Response to Reviewers K46j, jNu7 and iXy9.

Scaling to LLaMA3-70B Thank you for your valuable feedback and suggestions. We understand the importance of evaluating our approach on a more powerful model like LLaMA3-70B. However, due to the limited computational resources currently available to us, we are unable to conduct experiments on LLaMA3-70B within a short timeframe.

To explore the generality of our method as thoroughly as possible, we have conducted experiments on two models, LLaMA3-8B and LLaMA3.1-8B, with two different context lengths (32K and 128K). We validated our approach on a variety of benchmarks, including Needle-in-a-Haystack, LongBench, LEval, and InfiniteBench.

We hope this addresses your concern, but if you believe additional experiments on LLaMA3-70B are crucial, we are happy to include them in future work once the required resources become accessible.

Thank you for your understanding and consideration.

Training cost of PoD Please refer to Section Response to Reviewers jNu7 and iXy9.

Insufficient baselines Please refer to Section Response to all reviewers.

Training cost of PoD Please refer to Section Response to Reviewers jNu7 and iXy9.

Results over LLaMA3.1 with 128K context Please refer to Section Response to Reviewers K46j, jNu7 and iXy9.

The discussion on whether edge tokens are more important than middle tokens First, we acknowledge that, as you mentioned, the model may need information from any position during prediction. However, the frequency with which the model requires information from edge tokens versus middle tokens is significantly different. Two empirical results support this point:

1. Experimental results in Figure 1-(c) of our manuscript demonstrates that, even when attending to only the 256 edge tokens, the model predicts the next token identically to the model attending to all tokens (32K) in 80% of the cases for the same input sequence.
2. In Tables 1 and 2 of Section Response to all reviewers (part 2), the model based on window attention maintains 80%-90% of the performance of the dense model in practical tasks (Line 9 vs. Line 1).

Therefore, we consider the importance of edge tokens and middle tokens from the perspective of the frequency with which the model requires their information, rather than their intrinsic importance for completing the task itself. More specifically, unlike token-eviction-based methods that directly discard middle tokens, we do not ignore these middle tokens. Instead, we reduce their impact on KV cache consumption through attention layer sharing, while ensuring that the model's performance is not significantly compromised. We hope this explanation alleviates your concerns.

Performance-efficiency tradeoff We supplement and compare the performance-efficiency tradeoff of all baselines and our methods in Table 4. We focus primarily on the compression of the KV cache, so in Table 4, we compare the different methods in terms of the optimized stage (prefilling or decoding), the proportion of KV cache saved, and the resulting performance loss. From this, we can draw the following conclusions:

• Token-eviction-based methods (lines 4-8) are more efficient (suitable for both prefilling and decoding, and they save a relatively large amount of KV cache), but they incur a greater performance loss for the model.
• Token-selection-based methods (SnapKV & PyramidKV, lines 1-2) can significantly compress the KV cache with minimal impact on model performance, but they are only applicable to the prefilling stage. Quest (line 3), on the other hand, does not degrade model performance and can be applied to both prefilling and decoding stages. However, it does not reduce the size of the KV cache; instead, it reduces the KV cache load (from the GPU's storage to the computation area) by restricting the model's attention to a selectively filtered subset of tokens.
• Layer-selection-based methods can be applied to both the prefilling and decoding stages. Our method incurs much less performance degradation compared to the classical layer-sharing method, CLA (lines 9-10). Furthermore, our method is orthogonal to token-selection-based methods, and when combined (lines 1, 10, 11), the two approaches result in a model that excels in three aspects: optimized stage (applicable to both prefilling and decoding), KV cache saving (with a higher compression rate), and maintaining model performance.

We have updated this in Table 3 and the highlighted lines 378-401 in the revised version.

Table 4: Performance-efficiency tradeoff comparison for LLaMA3-8B-32K [1-3: token-selection-based methods; 4-8: token-eviction-based methods; 9-11: layer-sharing-based methods]

Implementation of the gate As demonstrated in Formulas (3) and (4) on page 5 of our manuscript, calculating the gate involves obtaining the exponential sum across each row of $QK^T$(with dimensions of batch size × head number × sequence length), rather than itself (with dimensions of batch size x head number x sequence length x sequence length). As you noted, FlashAttention does not return the complete matrix but outputs the exponential sum across each row [1].

We will release our code in the future, which will include more implementation details.

[1] https://github.com/Dao-AILab/flash-attention/blob/v2.5.7/flash_attn/flash_attn_interface.py#L824

Baseline-setup The additional lightweight post-training introduces a certain imbalance when comparing our method to token-eviction-based approaches (e.g., StreamingLLM, H2O). To address this, we could either apply post-training to these token-eviction-based methods or adapt our approach to be compatible with such methods. However, the former poses significant challenges, as the updated positional embeddings in StreamingLLM and the dynamic token eviction in H2O complicate training. Conversely, adapting our method to token-eviction-based strategies conflicts with our primary objective of retaining all tokens. Therefore, to ensure a more equitable comparison, we have included additional baselines based on token-selection and layer-sharing, as outlined in the Section Response to all reviewers.

Missing baselines Please refer to Section Response to all reviewers.

Longer context and other LLMs Please refer to Section Response to Reviewers K46j , jNu7 and iXy9.

Compatibility with modern attention optimizations

• First, sorry for the misunderstanding.

The attention pattern for each attention head is the same. The difference lies in the layers that share the attention scores for distant tokens.

For example, in Figure 2-①, all heads use the attention mask (pattern) shown in Figure 2-②. However, for head 1, the attention scores for distant tokens in layer 3 are derived from the results computed in layer 2. In contrast, for head H, the attention scores for distant tokens in layer 3 are derived from the results computed in layer 1.

• Second, the implementation is fully compatible with modern attention optimizations. As described in Section 2.2 of our manuscript, the computation for each token in the attention module is divided into two parts: computation for proximal tokens and distant tokens. Let $Q_{l}^h, K_{l}^h, V_l^h$ denote the query, key, and values states of the $l\text{-th}$ layer for head $h$. $H$ is the number of heads.

  • For proximal tokens, there is no sharing of information across layers, and the attention pattern follows the approach of StreamingLLM [1]. Since FlashAttention does not currently support this, we utilized PyTorch's FlexAttention [2], an efficient attention implementation technique that supports custom masks, i.e.,

    O_l^P = \text{FlexAttention}\left(\text{Concat}\left( \left\\{Q_l^h\right\\}_{1\le h\le H} \right) , \text{Concat}\left( \left\\{ K_l^h \right \\} _{1\le h\le H} \right) , \text{Concat}\left( \left \\{ V_l^h \right \\} _{ 1\le h\le H } \right) \right) (1)

  • For distant tokens, different heads may share attention scores with different layers. Since most modern optimized attention mechanisms do not support returning the attention scores (with dimensions: batch size × head number × sequence length × sequence length), we pass the query and key states from the lower layer to the upper layer for each head. The upper layer then concatenates these head-specific states and recalculates the attention, i.e.,

    \\tilde{Q} _l^h=\\begin{cases} Q _l^h, & \text{if no sharing between layer }l \text{ and } l-1 , \\\ Q _{l-1}^h, & \text{otherwise} . \\end{cases} (2)

    $\tilde{K} _l^h=\begin{cases} K _l^h, & \text{if no sharing between layer }l \text{ and } l-1 , \\\ K _{l-1}^h, & \text{otherwise} . \end{cases}$ (3)

    O_l^D = \text{FlexAttention}\left(\text{Concat}\left( \left\\{\\tilde{Q} _l^h\right\\} _{1\le h\le H} \right) , \text{Concat}\left( \left\\{ \tilde{K} _l^h \right \\} _{1\le h\le H} \right) , \text{Concat}\left( \left \\{ V_l^h \right \\} _{ 1\le h\le H } \right) \right) (4)

We will release our code in the future, which will include more implementation details.

Computational overhead of managing distinct attention patterns for each head As mentioned above, the attention pattern for each head is the same. Compared to standard dense attention, the computational overhead of our method primarily comes from the transmission of query states and key states between layers. During inference, the key states come from the KV cache, so only the transmission of query states between layers adds extra computational overhead.

In the prefilling phase, the main memory usage comes from the calculation of attention scores (batch size x number of heads x sequence length x sequence length) (with Flash Attention optimizing this part). The query states transmitted between layers have dimensions of (batch size x number of heads x sequence length x model dimension), which is much smaller compared to the attention score when the context is long. Similarly, in the decoding phase, only one token's query state is transmitted between layers at each step, which is minimal compared to the key and value states of all tokens already cached in the KV cache. Therefore, the computational overhead introduced by our method is relatively small.

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

[2] https://pytorch.org/blog/flexattention/

Thank you for providing responses.

InfiniteBench comprises 12 unique tasks. Why are there results for only four tasks, especially omitting the widely used Longbook QA? Additionally, the HELMET benchmark mentioned in reference [1] includes a wider variety of tasks and might be a better choice for testing. I also noticed that the Longbench results missed the Synthetic task.