Core Context Aware Attention for Long Context Language Modeling

We deeply appreciate your valuable feedback and constructive comments on improving our work. We would like to address your questions below.

Q1. The paper only compares CCA-Attention with StreamingLLM and LM-infinite, without including existing methods that retrieve relevant information in long contexts, such as MInference [A] and InfLLM [B]. Additionally, it does not compare the proposed method with models directly trained on longer texts. The evaluation of long-text capabilities is limited to multi-document question answering. It would be beneficial for the authors to evaluate the methods on a wider range of tasks, such as those in the RULER [C] and LongBench [D] benchmarks.

A1. Thanks for your suggestions. We conduct further comparisons between our CCA-Attention and baseline models on Longbench-E. As shown in Table I, our CCA-Attention attains the highest average score on Longbench-E. For example, the average score of our CCA-LLM is higher than that of LM-Infinite (22.12 vs. 21.20) and LM-MInference (22.12 vs. 22.08). Regarding InfLLM, we utilize its official implementation to evaluate its LongBench performance. Nevertheless, InfLLM consistently generates repeated and meaningless characters, resulting in an average score of merely 0.1.

Furthermore, we report the inference speed and memory footprint with respect to a 32K context. The reason for choosing 32K to showcase the inference speed and memory is that the longest input within Longbench is approximately 32K. Our CCA-LLM demonstrates a faster inference speed (3.5 times that of vanilla self-attention) and lower memory consumption (46% less than vanilla self-attention). These results confirm the effectiveness and efficiency of our CCA-Attention.

We have included these results in Section C.1 and disscussed them in related works of the revised manuscript. Due to the time constraint of the rebuttal, we are currently unable to provide the results on RULER benchmark. We would conduct and experiments on RULER in the future.

Table I. Comparisons with state-of-the-art methods in terms of LongBench-E. We report the inference latency and memory usage in the pre-filling phase on a single A800 GPU.

[A] MInference 1.0: Accelerating Pre-filling for Long-Context LLMs via Dynamic Sparse Attention. NeurIPS 2024.

[B] InfLLM: Training-Free Long-Context Extrapolation for LLMs with an Efficient Context Memory. NeurIPS 2024.

[C] RULER: What's the Real Context Size of Your Long-Context Language Models? COLM 2024.

[D] LongBench: A Bilingual, Multitask Benchmark for Long Context Understanding. ACL 2024.

Q2. The attention mechanisms are applied to the query of the last token and the key and value of the core tokens. However, during the formation of core tokens, the importance of each token in the group is determined by the query of the last token. Could information that is relevant to the query but less important within the group be overlooked?

A2. Our CCA-Attention approach of assessing token importance within each group using the attention from the last token is both rational and effective. This is supported by two points:

• Attention Pattern Insight from the Visualization in Appendix C.2: The attention map visualization reveals a distinct pattern where tokens that are important to the query receive consistently high attention scores from all subsequent tokens. This indicates that important tokens, regardless of their position within a group, have a notable influence on the attention distribution, suggesting that our method of importance assessment is capable of capturing these crucial tokens.
• Empirical Performance Validation: Our experimental outcomes demonstrate the effectiveness of this strategy. The consistent high performance in long-context modeling tasks, as evidenced by our perplexity and EM scores, confirms that our CCA-Attention mechanism not only maintains computational efficiency but also effectively captures global and local dependencies within long texts. This effectiveness is a direct result of our pooling strategy, which ensures that information relevant to the query is not overlooked, even when tokens are grouped and evaluated within a local context.

We have included these discussions in the revised manuscript.

Thank you for your feedback! I will increase my rating to 6.

We deeply appreciate your constructive comments. We would like to address your questions below.

Q1. How do models based on other sparse attention or linear attention approximations perform under the same training configuration?

A1. Thanks for your suggestions. We conduct further experiments on Longbench-E with our CCA-Attention and baseline models by applying them on pretrained LLMs. The baseline models include sparse attention methods (StreamingLLM[A], LM-Infinite[B], InfLLM[C] and Minference[D]). For fair comparisons, we do not compare with linear attention methods, since they introduce kernel function for attention and often require training from scratch. In this sense, linear attention methods are hard to applied to existing pretrained LLMs.

As shown in Table I, our CCA-LLM attains the highest average score on Longbench-E. For example, the average score of our CCA-LLM is higher than that of LM-Infinite (22.12 vs. 21.20) and LM-MInference (22.12 vs. 22.08). Regarding InfLLM, we use its official implementation to evaluate its LongBench performance. Nevertheless, InfLLM consistently generates repeated and meaningless characters, resulting in an average score of merely 0.1.

Furthermore, we report the inference speed and memory footprint with respect to a 32K context. The reason for choosing 32K to showcase the inference speed and memory is that the longest input within Longbench is approximately 32K. Our CCA-Attention demonstrates a faster inference speed (3.5 times that of vanilla self-attention) and the lowest memory consumption (46% less than vanilla self-attention). These results confirm the effectiveness and efficiency of our CCA-Attention.

We have included these results in Section C.1 the revised manuscript. Due to the time constraint of the rebuttal, we are currently unable to provide the results on RULER benchmark. We would conduct and experiments on RULER in the future.

Table I. Comparisons with state-of-the-art methods in terms of LongBench-E. We report the inference latency and memory usage in the pre-filling phase on a single A800 GPU.

[A] Efficient Streaming Language Models with Attention Sinks. ICLR 2024.

[B] LM-Infinite: Simple On-the-Fly Length Generalization for Large Language Models. arXiv 2024.

[C] InfLLM: Training-Free Long-Context Extrapolation for LLMs with an Efficient Context Memory. NeurIPS 2024.

[D] MInference 1.0: Accelerating Pre-filling for Long-Context LLMs via Dynamic Sparse Attention. NeurIPS 2024.

Q2. Table 3 reflects the different trends of PPL and MMLU metrics. Althought the PPL of Max Pooling is slighly lower than that of Mean Pooling, its MMLU score is significantly better than the other one's. This raises concerns about whether adopting LoRA+ will decrease LongLoRA's performance as presented in Table 2.

A2. We have incorporated more comparisons involving two variants of LongLoRA, namely LongLoRA (LoRA+) and LongLoRA (Full finetuning). As shown in Table II, LongLoRA with full finetuning attains a better performance in long-context modeling than LongLoRA with LoRA+. Furthermore, our CCA-LLM invariably outperforms the two variants of LongLoRA under all cases. Regarding the inference efficiency, our CCA-LLM achieves an inference speed that is 3.5$\times$ faster than that of LongLoRA in a 32K context. We have included these results and discussions in the revised manuscript.

Table II. Comparisons in terms of EM score under different contexts (%) between LongLoRA and our CCA-LLM.

Q3. Why is the performance of SinkAttention even worse than that of Vanilla Self-Attention, especially in the aspect of Inference Speed and Inference Memory Usage?

A3. We analyze the suboptimal performance of SinkAttention from two perspectives:

• Ineffectiveness in long-context modeling. SinkAttention solely concentrates on the initial and the most recent tokens, thereby neglecting the crucial information within the intermediate tokens. Consequently, it is difficult for SinkAttention to extract useful information in long-document question-answering tasks. Similar experimental outcomes have also been identified in [A, B].
• Lower computational efficiency. In the efficiency comparisons presented in Figure 4, we employ its official implementation for SinkAttention to evaluate its inference speed and memory usage. The inferior performance compared to vanilla self-attention can be mainly ascribed to two factors:
  • During the pre-filling phase, SinkAttention necessitates the computation of attention across all tokens. It inevitably results in inference speed and memory usage that are at least equivalent to those of vanilla self-attention.
  • The official implementation of SinkAttention does not integrate with FlashAttention[C], an acceleration technique adopted in both vanilla self-attention and our CCA-Attention. This contributes to its reduced efficiency compared to vanilla self-attention and our CCA-Attention.

We have re-run SinkAttention with FlashAttention and updated the results in Figure 4 of the revised manuscript.

[A] MInference 1.0: Accelerating Pre-filling for Long-Context LLMs via Dynamic Sparse Attention. NeurIPS 2024.

[B] InfLLM: Training-Free Long-Context Extrapolation for LLMs with an Efficient Context Memory. NeurIPS 2024.

[C] FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. ICLR 2024.

Q4. Also, what is the performance of LongLoRA in this context?

A4. In Figure 4, We have already incorporated comparisons with LongLoRA. LongLoRA's $S^2$-Attention is designed for the training phase and is not compatible with autoregressive generation. Consequently, LongLoRA reverts to the vanilla self-attention during inference. In this sense, its inference speed and memory consumption are the same as those of the vanilla self-attention (which is alreasy reported in Figure 4). We would further clarify this in the revised manuscript.

We sincerely hope our clarifications above have addressed your questions.

We are grateful for your time and effort. We would like to answer your questions below.

Q1. The method evaluates on too few benchmarks; recent long-document evaluation tasks, such as Longbench, Ruler, and Infinitebench, are available. Many papers have indicated that PPL is not an accurate metric.

A1. Thanks for your suggestions. We conduct further experiments on Longbench-E with our CCA-LLM and baseline models. As shown in Table I, our CCA-LLM attains the highest average score on Longbench-E. For example, the average score of our CCA-LLM is higher than that of LM-Infinite (22.12 vs. 21.20) and LM-MInference (22.12 vs. 22.08). Regarding InfLLM, we utilize its official implementation to evaluate its LongBench performance. Nevertheless, InfLLM consistently generates repeated and meaningless characters, resulting in an average score of merely 0.1.

Furthermore, we report the inference speed and memory footprint with respect to a 32K context. The reason for choosing 32K to showcase the inference speed and memory is that the longest input within Longbench is approximately 32K. Our CCA-LLM demonstrates a faster inference speed (3.5 times that of vanilla self-attention) and lower memory consumption (46% less than vanilla self-attention). These results confirm the effectiveness and efficiency of our CCA-Attention.

We have included these results in Section C.1 the revised manuscript. Due to the time constraint of the rebuttal, we are currently unable to provide the results on RULER and Infinitebench benchmark. We would conduct and experiments on RULER and Infinitebench in the future.

Table I. Comparisons with state-of-the-art methods in terms of LongBench-E. We report the inference latency and memory usage in the pre-filling phase on a single A800 GPU.

Q2. In Table 2, why wasn’t LLaMA-2 used with continued training for comparison?

A2. We have further included results where LLaMA-2 7B and LongLoRA are continued training on 8K and 16K context lengths, as shown Tables II. Additionally, we report the inference speed and memory footprint of both LLaMA-2 and our proposed CCA-Attention across different context lengths in Table III. Although our CCA-Attention slightly lags behind full self-attention in terms of EM scores, it achieves significant improvements in inference speed and memory efficiency, e.g., 5.71$\times$ inference speed with only 56.41% GPU memory footprint under 64K context. Moreover, when compared to the LongLoRA method, our approach not only outperforms it in terms of long-context modeling accuracy but also achieves faster processing times and lower memory usage. We have included thses results in the revised manuscript.

Table II. Comparisons in terms of EM score under different contexts (%) and MMLU (%) between our CCA-Attention and existing attention methods.

Table III. Comparisons in terms of inference latency (seconds) and memory usage (GB) in pre-filling phase between our CCA-Attention and vanilla self-attention. Since $S^2$-Attention in LongLoRA does not support inference-stage evaluation, its inference speed is identical to vanilla self-attention.

(a) comparisons on inference latency (seconds)

(b) comparisons on memory usage (GB)

We are truly grateful to the reviewer for bringing up the new questions. Their insights will be invaluable in enhancing our manuscript.

Q1. Limiting the testing range to 32k makes it difficult for me to assess the value of this paper, especially when many open-source models (e.g., LLaMA-3.1, LLaMA-3.2) already support a context length of 128k.

A1. Thank you for your questions and suggestions. During the rebuttal, we evaluate the performance of our CCA-Attention on LLaMA2-7B with context lengths of 64K and 128K on the Multi-document QA task [r1]. As shown in Tables I and II, our CCA-Attention exhibits a substantially better EM Score and a significant inference speedup compared to vanilla self-attention at these context lengths. In particular, CCA-Attention shows much better performance than vanilla self-attention in terms of EM score (31.45 vs. 17.52) and 7.9x inference speedup with a context length of 128K.

Table I. Performance comparisons of CCA-Attention and vanilla self-attention on LLaMA2-7B in terms of EM score (%) on the Multi-document QA task with different context lengths.

Table II. Performance comparisons of CCA-Attention and vanilla self-attention on LLaMA2-7B in terms of inference time (seconds) on a A800 GPU with different context lengths. We also report the speedup of our method compared to vanilla self-attention.

More critically, from Tables I and II, the advantages of our method become more prominent as the length of the context increases (in terms of both performance and speedup), while the performance of vanilla self-attention may even decrease when the context length becomes very large.

The reasons for the prominent performance of CCA-Attention towards long-context modeling. As discussed in the paper, we discovered that self-attention may face severe redundant context issue with an extremely long context in sequence modeling p(x_t|context). This not only hampers the modeling performance, but incurs unbearable waste of computational overhead. To address this, we propose the core context-aware attention mechanism, in which non-core contexts (i.e., the irrelevant context to any x_t) will be compressed by weighted pooling. In this way, CCA-Attention not only alleviates the redundant context issue and thus improves the long-context modeling performance, but also enhances computational efficiency significantly. In particular, the KV cache of our CCA-Attention is remarkably smaller than the vanilla self-attention, e.g., 4.5GB vs. 64GB with a context length of 128K on LLaMA2-7B.

At this moment we may not be able to provide the results on the latest LLaMA-3.1/3.2 models for two reasons: 1) Although our CCA-Attention can be adopted as a plug-and-play module to replace the vanilla self-attention, finetuning (or even full training which may need thousands of GPUs) is required to learn the parameters of CCA-LLM. Unfortunately, the training data of LLaMA-3 so far is not publicly available, and some key learning hyperparameters (such as learning rate and weight decay strategy) are unknown. 2) Our method is currently designed for improving the vanilla self-attention that is widely adopted in mainstream LLMs (such as OPT[r2], LLaMA-2[r3], and Qwen-1.5[r3]), while LLaMA-3.1/3.2 adopt grouped query attention (GQA) [r5]. As a result, applying our method to GQA may need further adjustments (such as modifying the pooling strategy of non-core tokens). We leave these for future exploration.

[r1] Lost in the middle: How language models use long contexts. TACL, 2024.

[r2] OPT: Open Pre-trained Transformer Language Models. arXiv, 2022.

[r3] Llama 2: Open Foundation and Fine-Tuned Chat Models. arXiv, 2023.

[r4] Introducing Qwen1.5. arXiv, 2024.

[r5] GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints. EMNLP, 2023.

Thank you for your response.

Q1: Vanilla self-attention is highly related to the training data and RoPE base size. I believe that the performance of CCA-Attention surpassing vanilla self-attention is unexpected. Many studies [1][2] have already demonstrated the strong performance of vanilla self-attention when applied to long texts. It seems the authors might not have fine-tuned vanilla self-attention properly.

Q2: Extending the model's length should not come at the cost of its performance on short texts. Similar discussions can also be found in [1][2].

Q3: The authors focused excessively on speed comparisons, but improving speed should not sacrifice too much performance.

The authors' response did not address my concerns, so I lowered my score from 5 to 3.

[1] Fu Y, Panda R, Niu X, et al. Data engineering for scaling language models to 128k context.

[2] Gao T, Wettig A, Yen H, et al. How to train long-context language models (effectively).

Q1: Vanilla self-attention is highly related to the training data and RoPE base size. I believe that the performance of CCA-Attention surpassing vanilla self-attention is unexpected. Many studies [r1][r2] have already demonstrated the strong performance of vanilla self-attention when applied to long texts. It seems the authors might not have fine-tuned vanilla self-attention properly.

A1. We feel regret that the reviewer may misunderstand our experimental setting and results, thus lower the review score. We try to make the following clarifications on the finetuning training strategy and the experimental results.

We would highlight the comparisons between our CCA-Attention and vanilla self-attention are fair. We adopt the same and widely used settings for finetuning both models: following [r3], we finetune for 1000 steps with a total of 1.05 billion tokens using the data from [r1]. We set the base size of RoPE to 500,000 following the common settings in [r4]. From 4K to 128K, we use the test data from [r5] with different context lengths and same experimental settings. It should be noted that we use the finetuning instead of pretaining for validation of our method mainly due to the limitations of computing resources. Additionally, the relatively poor performance of vanilla self-attention at 128K might be attributed to the insufficiency of data (as previously stated, 1.05 billion tokens). Currently, we are in the process of conducting finetuning with more data to further improve and validate our method. To enhance the reliability of our research, we would release our code and models to ensure reproducibility.

Why CCA-Attention shows better performance than vanilla self-attention in long-context modeling (e.g., 64K and 128K). As stated in the prior rebuttals, we discovered that redundant contexts have adverse effects for self-attention particularly in an extremely long context. In our CCA-Attention, redundant/non-core contexts will be compressed by weighted pooling. In this way, CCA-Attention not only alleviates the redundant context issue and but also improves the long-context modeling performance.

[r1] Data Engineering for Scaling Language Models to 128K Context. ICML 2024.

[r2] How to Train Long-Context Language Models (Effectively). arXiv 2024.

[r3] LongLoRA: Efficient Fine-tuning of Long-Context Large Language Models. ICLR 2024.

[r4] Effective Long-Context Scaling of Foundation Models. NAACL-HLT 2024.

[r5] Lost in the middle: How language models use long contexts. TACL, 2024.

Q2: Extending the model's length should not come at the cost of its performance on short texts. Similar discussions can also be found in [r1][r2].

A2. We would like to highlight that extending a model's context length often leads to performance degradation on shorter text tasks, as your mentioned the works [r1][r2]. For instance, in [r1], LongLoRA-7B suffers a notable decline of 7.4 in terms of MMLU (from 45.3 to 37.9), while LongChat-v1.5-7B decreases by 3.0 (from 45.3 to 42.3). All these evidences verify the trade-offs inherent in extending model lengths, particularly the challenge of maintaining performance on shorter texts.

[r1] Data Engineering for Scaling Language Models to 128K Context. ICML 2024.

[r2] How to Train Long-Context Language Models (Effectively). arXiv 2024.

Q3: The authors focused excessively on speed comparisons, but improving speed should not sacrifice too much performance.

A3 Our motivation is not merely aimed at acceleration. Instead, we discovered that self-attention may face severe redundant context issue with an extremely long context in sequence modeling. This not only hampers the modeling performance, but incurs unbearable waste of computational overhead. To address this, we propose the core context-aware attention mechanism, in which non-core contexts will be compressed by weighted pooling, thereby improving performance in long-context modeling. Moreover, the reduction of redundant contexts greatly improves inference speed.

Our method shows significant advantages in long-context scenarios, where the redundant context issue is more critical. For shorter contexts (e.g., 4K or 8K), the redundancy issue is less severe. CCA-Attention substantially revises the structure of self-attention, necessitating fine-tuning for effective integration. Due to constraints in computational resources and time, we trained our model on a very small subset of the data using a single A800 server with 8 GPUs. This results in fewer training samples (1B tokens compared to the 2T tokens used in the original LLaMA2 pretraining), which may impact performance in shorter contexts (e.g., 4k or 8k). We plan to train our CCA-LLMs on more data to further enhance the performance on shorter contexts.

We sincerely hope that you can understand our motivation and core contributions. We would be grateful if you could kindly reconsider the evaluation of our paper.

We are grateful for your great effort for reviewing our work. After the last reply, we further finetune LLaMA-2 model using More Data to evaluate the long-context performance of the vanilla self-attention.

First, we would like to clarify that the previous results of LLaMA-2 on 128K context length were obtained by finetuning with limited data (following settings in LongLoRA[r1] with 1.05 billion tokens) using only one A800 GPU station with 8 GPUs. This may result in the suboptimal performance of vanilla self-attention. Nevertheless, we use the same finetuning strategy on the same data set for both CCA-Attention and vanilla self-attention, thus our comparison is completely Fair. Moreover, at the context length of 64K, the vanilla self-attention achieves slightly worse performance than CCA-Attention in terms of EM Score (31.33 vs. 28.91), but exhibits much slower inference speed than CCA-Attention.

Second, after the last reply, we try to further finetune the LLaMA-2 model with 128K context on more data (i.e., 5 billion tokens of SlimPajama dataset [r2]). Under this setting, the vanilla self-attention shows slightly better performance than the last experiment in terms of EM Score (19.02 v.s. 17.50). So, with this observation, we believe our last results are reasonable. However, we shall explore more fintuning strategies to extend LLaMA-2 model on larger datasets to 128K context, for which we leave our future study.

Last, we wish to emphasize that CCA-Attention achieves 5.7x and 7.9x inference speedup at context lengths of 64K and 128K than vanilla self-attention, respectively.

We would be sincerely grateful if you could reconsider the evaluation of our paper.

[r1] LongLoRA: Efficient Fine-tuning of Long-Context Large Language Models. ICLR 2024.

[r2] Data Engineering for Scaling Language Models to 128K Context. ICML 2024.