Core Context Aware Transformers for Long Context Language Modeling

Q1. The paper only compares against training-free methods, important training-based approaches are missing. While the method achieves better latency and memory efficiency, the actual improvement in average score is marginal.

A1. We appreciate your emphasis on contextualizing our work against training-based methods. Below, we address the concerns raised:

Method Positioning and Performance Our work focuses on long-context modeling via efficient attention mechanism for pretrained LLMs. This is distinct from position interpolation methods or long-context adaptation techniques like LongLLaMA, which do not modify attention in LLMs. By reducing computational costs ( 7.9× on 128K context) while outperforming training-free baselines, our method offers immediate value to practitioners. Last, we emphasize our goal is to reduce redundancy in attention, which should be compared with those that also optimize attention (e.g., MInference).

Further Experiments To address your concerns, we further compare our method with a training-based model, LongLoRA[ICLR 2024] with PI techniques. In Table I, our CCA-LLM outperforms LongLoRA in modeling accuracy, inference speed and memory footprint. We will include these in the revised manuscript.

Table I: Results on LongBench-E. ''FTL'' denotes the latency to generate the first token.

** Note: During inference, LongLoRA exhibits the same computational and storage overhead as full self-attention. The reason is LongLoRA's S2-Attention is only applicable in training and uses full self-attention during inference.

Q2. More results on different architectures.

A2. According to your suggestions, we conduct more experiments on LLaMA3.1-8B and Qwen2.5-7B (128K context). The results in Tables II and III show CCA-Attention outperforms state-of-the-art counterpart MInference in terms of performance, computational efficiency, and memory reduction (see Table II and III).

We would include these results in the revised paper.

Table II. Comparisons of Llama3.1-8B-Instruct-128K on LongBench-E.

Table III. Comparisons of Qwen2.5-7B-128K on LongBench-E.

Q3. Benefits of using fixed-size grouping.

A3. This allows for optimized memory access and parallel computation. While dynamic group sizes would introduce significant implementation complexities.

Q4. If multiple tokens in a group have high scores with the last token, wouldn't using one core token to represent them introduce significant approximation errors?

A4. The mentioned issue can be significantly mitigated by the fine-tuning process, based on the finding: CCA attention with/without fine-tuning yields performance of 6.92 and 22.24 on LongBench-E, respectively. This confirms core token aggregation effectively captures group-wise information.

Q5. Why are the scores in Table 1 different from the official leaderboard?

A5. The leaderboard tests chat models with SFT on LongBench, while we test base models without SFT on LongBench-E.

Q6. Whether during the fine-tuning only fixed g and s are used? How can it perform well across varying g and s?

A6. We use fixed $g$ and $s$ during finetuning. The reasons for the promising performance with varying $g$ and $s$ are twofold:

• Globality-aware pooling module dynamically pools groups of tokens into core tokens based on their importance, regardless of $g$. Thus it generalizes to arbitrary $g$ during inference.
• Locality-preserving module employs rotary position embeddings that encode relative distances rather than absolute positions, which decouples from specific $s$, allowing adaptation to different $s$.

We would include these in the revised paper.

Q7. Was CCA-LLM fully or partially fine-tuned in Tables 1 & 2? Could you compare them?

A7. Full fine-tuning. In Section C.3, we compare different fine-tuning strategies. Full fine-tuning achieves better performance than partial fine-tuning (22.24 vs. 21.96) on LongBench-E with more training hours (15.4 vs. 11.2).

Q8. Implementation details in Figure 4.

A8. We use batch size=1, $g$=16 and $s$=1024.

Q9. Why does MInference barely reduce KV cache usage?

A9. Minference's sparse attention is only applicable in pre-filling. During decoding, it uses full self-attention, leading to same KV cache usage.

---

We sincerely hope our clarifications above have addressed your questions.

Q1. The proposed methods share similarities to previous context compression methods[r1-r6]. r1: Compressive Transformers, r2: Gist, r3: NUGGET, r4: Landmark, r5: Beacon, r6: DODO.

A1. Thanks for your valuable comments. We clarify the differences below:

• Problem Importance & Motivation The core challenge we tackle is efficient and effective long-context modeling in LLMs. While self-attention handles long sequences, its quadratic complexity and redundancy in ultra-long contexts (e.g., 128K) severely degrade efficiency and modeling performance. Unlike prior works either compress features via auxiliary networks [r1] or compress context with extra specific tokens [r2-r6]), our method aims to dynamically identify and strengthen task-relevant core-context while suppressing redundancy. This distinction is crucial: compression prioritizes shortening sequences, whereas our core-context awareness optimizes context redundancy in self-attention computation for long context modeling.
• Novelty & Differences Our CCA Attention introduces two innovative, complementary mechanisms:
  • Globality-aware Pooling Module: Unlike context compression methods [r1-r6], it dynamically groups tokens and pools them into core tokens based on their importance. This ensures the model focuses on task-relevant context while reducing redundancy.
  • Locality-preserving Module: While prior compression methods may have a risk of losing local details, our module selectively preserves neighboring tokens to maintain locally relevant information.

We would cite the above papers and include these discussions in the revised manuscript.

---

Q2. It might be interesting to explore more long-context tasks in addition to LongBench, such as RULER.

A2. According to your suggestions, we conduct more experiments on RULER on LLaMA2-7B-80K with varying context lengths (4K-64K). From the results in Table I, our method achieves superior performance over MInference, a state-of-the-art counterpart.

Table I: Comparisons of LLaMA2-7B-80K on RULER.

Q3. It would be interesting to explore whether the proposed approach can help more recent models (such as LLama 3 and Qwen models).

A3. We conduct more experiments on LLaMA3.1-8B and Qwen2.5-7B with context length 128K in Tables II and III, respectively. From the results, our method is able to help LLaMA3.1 and Qwen models improve computational efficiency and memory reduction in long-context modeling.

• Computational efficiency: CCA-Attention achieves 3.1× on LLaMA3.1-8B, and 3.9× on Qwen2.5-7B with 32K context.
• Memory reduction: CCA-Attention reduces 49% GPU memory footprint on LLaMA3.1-8B, and 45% on Qwen2.5-7B.
• Performance: CCA-Attention consistently outperforms MInference across all models, maintaining competitive accuracy with valina self-attention.

Table II. Comparisons of Llama3.1-8B-Instruct-128K on LongBench-E. ''FTL'' denotes the latency to generate the first token in the pre-filling stage.

Table III. Comparisons of Qwen2.5-7B-128K on LongBench-E.

Q4. There are some key hyper-parameters, such as the group size and local window size. It would be nice to include more analyses and ablation studies.

A4. In Section C.3 of supplementary, we show ablation studies on group size $g$ and local window size $s$. From the results, as $g$ increases, the efficiency will improve and the performance will decrease (in terms of PPL); as $s$ increases, the performance will improve and the efficiency will reduce. To achieve a balance, we select $g=16$ and $s=1024$ as default training setting.

Q5. How much data will be needed to fine-tune the model?

A5. In Section B.2 of the supplementary, we show the details for fine-tuning: We use SlimPajama dataset to fine-tune the model, with only 2.1 billion tokens for LLaMA2-7B-32K and 5.0 billion tokens for LLaMA2-7B-80K.

---

We sincerely hope our clarifications above have addressed your questions.

Q1. Figure 3 demonstrates stable performance with varying group sizes (g) and window lengths (s) during inference, what are the possible reasons?

A1. We deeply appreciate your valuable feedback. We design our CCA-Attention for flexibility, enabling promising performance across varying $g$ and $s$ settings without requiring additional fine-tuning. The reasons are twofold:

• Globality-aware pooling module dynamically pools groups of tokens into core tokens based on their importance, regardless of the group size $g$. The pooling weights are derived from intra-group attention scores, which depend on semantic relevance rather than fixed positional patterns. As a result, the module generalizes to arbitrary group sizes $g$ during inference.
• Locality-preserving module is inherently translation-invariant to window boundaries by employing rotary position embeddings that encode relative distances rather than absolute positions. This relative distance encoding naturally decouples the module's operation from specific window size $s$, allowing adaptation to different $s$ in inference while preserving local context relationships.

We would include these discussions in the revised paper.

Q2. Could the core token (Eqn.2) be biased toward later tokens in a group due to using the last token’s query?

A2. In our CCA-Attention, the strategy of deriving core tokens based on the attention scores from the last token is both rational and effective. This is supported by two points:

• The attention map visualization in Section C.2 of supplementary material reveals a distinct pattern: 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. This suggests that our pooling weights assessed by the last group token are able to capture these crucial tokens.
• The consistent high performance in long-context modeling tasks confirms that our CCA-Attention effectively captures global and local dependencies within long contexts. This is a direct result of our pooling strategy, which shows that information relevant to the query is not overlooked.

We would include these discussions in the revised paper.

Q3. Without requiring additional experiments, could the authors discuss: is there a principled way to set m for unseen tasks?

A3. We appreciate the reviewer's insightful comments. In practice, we are able to set different group sizes according to the nature of the task. For high-precision tasks (e.g., legal document analysis), smaller group sizes are preferable to ensure finer granularity and capture more detailed contextual dependencies. Conversely, for tasks prioritizing efficiency over precision (e.g., real-time dialogue systems), larger group sizes are able to reduce computational overhead. This trade-off aligns with the inherent flexibility of our proposed CCA-Attention, which allows dynamic adjustment of m during inference.

Q4. Could the authors explain why the model performance does not consistently degrade with increasing text length in Table 2?

A4. As discussed in the paper, we discovered that full self-attention may face redundant context issues in sequence modeling. This redundancy may hamper the modeling performance by 1) weakening the focus on semantically important content, and 2) introducing noise from irrelevant context. More importantly, this issue becomes more severe with longer sequences.

To address this, we propose a core context aware attention mechanism, in which non-core/irrelevant contexts will be compressed by weighted pooling. In this way, CCA-Attention not only alleviates the redundant context issue but thus improves the long-context modeling performance. As a result, the advantages of our method become more prominent as the length of the context increases.

Q5. Will the assumption of redundant text be related to the topic of testing cases? Math and general QA may have different degrees of redundancy.

A5. Our assumption of redundant context is not related to the topic. To verify this, we perform an analysis using LLaMA2-7B on 16K length corpora from two domains: math (MathPile dataset[1]) and general QA (WikiText dataset[2]).

By examining the attention weights of the final token relative to preceding tokens, we quantify the redundancy degree as the proportion of tokens receiving attention weights below 0.001. The results demonstrate comparable levels of redundancy between math (99.8% tokens' weights below 0.001) and general QA context (98.7% tokens' weights below 0.001), supporting our initial assumption.

[r1] A Billion-Token-Scale Pretraining Corpus for Math. NeurIPS 2024.

[r2] Pointer Sentinel Mixture Models. ICLR 2017.

---

We sincerely hope our clarifications above have addressed your concerns.