Core Context Aware Attention for Long
Context Language Modeling

The full self-attention excels in reachability, where the attention weights $a_{ij}$ for $i\geq j$ are positive. However, the computational complexity of the full self-attention grows in quadratical time as the sequence length $L$ increases, i.e., $O(L^{2})$. This poses a significant challenge in handling long sequences, affecting both processing time and memory footprint. Most existing methods (Beltagy et al., 2020; Ding et al., 2023; Chen et al., 2024) mitigate the computational cost by applying predefined and fixed sparse attention patterns. Unfortunately, they often overlook the importance of maintaining reachability between tokens. This oversight may lead to inadequate information transfer between tokens, which hinders the performance of complex tasks involving long sequences.

In this paper, we seek to reduce the computational complexity associated with full self-attention without sacrificing token reachability. To achieve this, we propose a Core Context Aware Attention (CCA-Attention), which employs globality-pooling and locality-preserved attention to capture both global and local dependencies within a long context. As shown in Figure 2, globality-pooling attention operates by generating representative core tokens from segmented groups of the input sequence. It then computes attention using these reduced-number core tokens, thereby reducing the computational cost (see Section 3.2). However, the autoregressive nature inherently limits each token to access core token within the same group (i.e., local context), failing to achieve token reachability. To address this, locality-preserved attention is responsible for capturing the local information of the neighborhood to ensure comprehensive coverage (see Section 3.3). Furthermore, we devise an adaptive fusion strategy to integrate the insights from both attentions (see Section 3.4). This strategy is crucial as it retains the reachability of the full self-attention within our CCA-Attention framework. The pseudo-code for our proposed CCA-Attention is presented in Algorithm 1.

Empirical studies in Figure 1 have revealed that attention weight maps exhibit a non-uniform distribution, with a minority of tokens being prominent while the majority are overlooked (aligning with findings in recent literature (Bondarenko et al., 2023; Xiao et al., 2023; 2024b)). Thus, it is possible to dynamically prioritize computational resources to prominent tokens and neglect the remaining ones. This could approximate the full self-attention with reduced computational overhead. Motivated by this, we propose globality-pooling attention that dynamically identifies prominent tokens and encapsulates the critical information into a smaller set of core tokens for attention.

Given an input sequence of tokens $\mathbf{X}{=}[\mathbf{X}_{1};\mathbf{X}_{2};\dots;\mathbf{X}_{L}]$, we segment the input sequence $\mathbf{X}$, each group containing $k$ tokens, in total $m{=}\lfloor L/k\rfloor$ group. For simplicity, we denote the $i$-th group by $\mathbf{X}^{\text{global}}_{i}{\in}\mathbb{R}^{k\times d}$, where ${\bf X}^{\text{global}}_{i}{=}[{\bf X}_{(i-1)k+1};{\bf X}_{(i-1)k+2};\dots;{\bf X}_{ik}]$ with ${\bf X}_{ik}$ representing the last token in the group. To identify prominent tokens in the group $\mathbf{X}^{\text{global}}_{i}$, we devise a group-wise attention pooling strategy that leverages the last token ${\bf X}_{ik}$ to evaluate the importance, as it has full access to all the tokens in $\mathbf{X}^{\text{global}}_{i}$. Formally, we derive a core token $\mathbf{c}_{i}$ from each group $\mathbf{X}^{\text{global}}_{i}$ by

where $\mathbf{Q}_{ik}$ is the query vector for the last token of $\mathbf{Q}^{{}^{\prime}}_{i}{=}{\bf X}^{\text{global}}_{i}{\bf W}^{Q}$ and $\mathbf{K}^{{}^{\prime}}_{i}{=}{\bf X}^{\text{global}}_{i}{\bf W}^{K}$, $\mathbf{W}^{Q}$ and $\mathbf{W}^{K}$ are learnable parameters. In this way, the core token $\mathbf{c}_{i}$ encapsulates crucial information of the corresponding group. Then, the core token sequence will be $\mathbf{C}{=}[\mathbf{c}_{1};\mathbf{c}_{2};\dots;\mathbf{c}_{m}]$.

To compute globality-pooling attention, we use the core token sequence $\mathbf{C}{=}[\mathbf{c}_{1};\mathbf{c}_{2};\dots;\mathbf{c}_{m}]$ as the key instead of the original sequence $\mathbf{X}$. This substitution reduces the dimensionality from $\mathbf{X}{\in}\mathbb{R}^{L\times d}$ to $\mathbf{C}{\in}\mathbb{R}^{m\times d}$, thereby reducing the computational complexity. Formally, the output $\mathbf{Att}^{\text{global}}_{i}$ with globality-pooling attention for each token ${\bf X}_{i}$ can be computed by

where $\mathbf{A}^{\text{global}}_{i}{=}\text{softmax}({{\bf Q}_{i}{{\bf K}^{\text{global}}}^{\top}}/{\sqrt{d}}){\in}\mathbb{R}^{1\times m}$, ${\bf Q}={\bf X}{\bf W}^{Q}{\in}\mathbb{R}^{L\times d}$, ${\bf K}^{\text{global}}={\bf C}{\bf W}^{K}{\in}\mathbb{R}^{m\times d}$, $\mathbf{V}^{\text{global}}{=}{\bf C}{\bf W}^{V}{\in}\mathbb{R}^{m\times d}$ with $\mathbf{V}^{\text{global}}_{\cdot,j}$ being its $j$-th column, $\mathbf{W}^{K}$ and $\mathbf{W}^{V}$ being learnable parameters. Conveniently, we can implement this by the attention operation $\mathbf{Att}^{\text{global}}=\text{Attention}(\mathbf{Q},\mathbf{K}^{\text{global}},\mathbf{V}^{\text{global}})\in\mathbb{R}^{L\times d}$. Importantly, $\mathbf{Att}^{\text{global}}$ retains the same dimensions as the input sequence, which aligns with the full self-attention. Since language models typically follow the next-token prediction paradigm during both training and inference, individual tokens cannot directly attend to the core tokens generated from their own group. In our globality-pooling attention, each token can only attend to the core tokens from preceding groups, aligning with the autoregressive nature of the language models.

As mentioned above, our globality-pooling attention prevents each token from accessing the core token in the same group due to the autogressive nature. Consequently, individual token is unable to fully capture the local dependencies of preceding neighboring tokens. However, numerous studies (Manakul & Gales, 2021; Yang et al., 2021) have demonstrated the critical role of local context in a variety of language modeling tasks. To bridge this gap, we introduce a supplementary attention mechanism termed locality-preserved attention. It serves as a complement to globality-pooling attention, enabling the model to effectively capture local dependencies that are essential for effective language modeling. To be specific, our proposed attention ensures that each token attends to the preceding $s(s\geq k)$ tokens to capture local dependencies, as defined by the following equation:

where $\mathbf{A}_{i}^{\text{local}}{=}\text{softmax}({{\bf Q}_{i}{{\bf K}^{\text{local}}}^{\top}_{[\max(1,i-s):i]}}/{\sqrt{d}}){\in}\mathbb{R}^{1\times{\max(1,i-s)}}$, $\mathbf{Q}=\mathbf{X}\mathbf{W}^{Q}$, ${\bf K}^{\text{local}}={\bf X}{\bf W}^{K}$, $\mathbf{V}^{\text{local}}{=}{\bf X}{\bf W}^{V}$ with $\mathbf{V}^{\text{local}}_{\cdot,j}$ being its $j$-th column. Similar to Eqn. (3), we can calculate this by the attention operation $\mathbf{Att}^{\text{local}}_{i}=\text{Attention}(\mathbf{Q}_{i},\mathbf{K}^{\text{local}}_{[\max(1,i-s):i]},\mathbf{V}^{\text{local}}_{[\max(1,i-s):i]})\in\mathbb{R}^{1\times d}$. $\mathbf{Att}^{\text{local}}{\in}\mathbb{R}^{L\times d}$ maintains the same dimensions as the input $\mathbf{X}$, consistent with the full self-attention. In practice, the attention for each $\mathbf{Att}^{\text{local}}_{i}$ is computed parallelly with high computational efficiency.

Note that the locality-preserved attention shares the linear projection parameters ${\bf W}^{Q}$, ${\bf W}^{K}$, and ${\bf W}^{V}$ with globality-pooling attention, thereby incurring no additional projection parameters. Furthermore, our globality-pooling and locality-preserved attentions are not required to be trained from scratch. By reusing the pretrained parameters from existing language models, our CCA-Attention can seamlessly replace the full self-attention in existing models with a modest fine-tuning effort. This facilitates efficient inference at a lower computational cost compared to full self-attention.

Compatibility with Attention-based Pretrained LLMs. We design our CCA-Attention to efficiently replace full self-attention, offering a plug-and-play module for existing attention-based LLMs. Our CCA-Attention aligns with full self-attention in input, output, and parameter dimensions. This compatibility ensures that only modest finetuning is able to maintain the long-context modeling capability with reduced computational cost. Conversely, existing linear attention approaches (Katharopoulos et al., 2020; Sun et al., 2023) introduce kernel function for attention and require training from scratch. This renders these methods less practical for real-world applications, as they do not reuse the extensive knowledge embedded in pretrained LLMs.

Compared with full self-attention, our CCA-Attention offers significant advantages in terms of computational complexity and key-value cache memory usage. These benefits substantially enhance the running speed and memory usage efficiency for our CCA-Attention (see results in Figure 4).

Acceleration through Reduced Computational Complexity. Our CCA-Attention exhibits varying computational complexities depending on the type of task. For tasks with fixed-length sequences (such as multi-choice question answering), our CCA-Attention exhibits a linear computational complexity of $O(Lm+Ls)$, marking a significant enhancement over the full self-attention with a complexity of $O(L^{2})$. Here, we define the number of group $m$ as a constant. Specifically, for globality-pooling attention, the query and key matrices encompass $L$ and $m$ tokens, respectively, resulting in a computational complexity of $\mathcal{O}(Lm)$. Regarding locality-preserved attention, each token only attends preceding $s$ tokens. With $L$ tokens in total, the complexity amounts to $O(Ls)$.

For tasks with variable-length sequences (such as open-ended question answering), models generate subsequent tokens in an autoregressive manner. In this case, we set the group size $k$ as a constant, ensuring that our CCA-Attention is able to leverage key-value caching during autoregressive token generation. Once one group has certain $k$ tokens, the corresponding core token is also determined and can be cached. Thus, our CCA-Attention achieves a computational complexity of $O(L^{2}/k+Ls)$. The complexity analysis follows a similar pattern to the tasks with fixed-length sequences.

Acceleration through Reduced Key-Value (KV) Cache. In attention-based transformers, the KV cache leverages the autoregressive nature to store and reuse key-value pairs (i.e., $\mathbf{K}$ and $\mathbf{V}$ in Eqn. (1)), thereby significantly boosting the efficiency of attention calculations. The size of the KV cache scales linearly with the length of the input sequence, consuming a major part of the memory footprint during inference. The expanded KV cache not only consumes considerable memory but also potentially hinders inference speed due to the increased demand for IO operations.

Compared with full attention’s complexity of $O(L)$, our proposed CCA-Attention achieves a storage complexity of $O(L/k+s)$. For globality-pooling attention, we only retain the key and value matrices for core tokens, rather than for all original tokens. This reduces the memory requirement to $O(L/k)$. Besides, locality-preserved attention confines the attention of each token to the preceding $s$ tokens. Consequently, the storage complexity for this component is $O(s)$. Experimentally, we have set $k{=}16$ and $s{=}1024$. The results show a significant reduction in storage complexity compared to traditional full attention models (see results in Figure 4).

We apply our CCA-Attention and considered efficient attention methods to pretrained LLMs, and report long context modeling performance. We provide the details of the considered models, evaluation metrics and implementation, more details are put in the supplementary materials.

Models. We apply our proposed CCA-Attention to LLaMA2-7B and LLaMA2-13B (Touvron et al., 2023b) models. We extend the context sizes of LLaMA2-7B and LLaMA2-13B up 16K and 32K, respectively. Following (Xiong et al., 2023), we modify the “base frequency $b$” of RoPE positional encoding (Su et al., 2024) from 10,000 into 500,000.

Compared Methods. We compare our CCA-Attention with LongLoRA (Chen et al., 2024), StreamingLLM (Xiao et al., 2024b), LM-Infinite (Han et al., 2023), InfLLM (Xiao et al., 2024a), and MInference (Jiang et al., 2024) in our experiments. We summarize differences between our CCA-Attention and compared methods in Table 1.

Evaluation Metrics. We quantitatively evaluate our models and compare them with other considered models in twofold: 1) EM Score (Liu et al., 2024) measures the ability to find the key information within a long multi-document context. 2) MMLU (Hendrycks et al., 2021) evaluates the ability to answer a broad spectrum of common knowledge questions. 3) LongBench Bai et al. (2023) is a pioneering benchmark for the bilingual, multi-task and comprehensive assessment of large language models’ long context understanding capabilities. It covers multiple languages like Chinese and English, consists of 6 major categories and 21 tasks involving various application areas.

Implementation Details. For the continuous pretraining, we adopt the SlimPajama (Cerebras, 2024) dataset, an open-source replication of the LLaMA pretraining data mixture. We replace the full self-attention in LLaMA-2 with our proposed CCA-Attention. The number of groups in Globality-aware Attention is shared across different model sizes. Training is conducted on 8 $\times$ A800 GPUs using a micro-batch size of 1 and a gradient accumulation of 8, with a total of 1000 training steps. This training configuration is applicable to all model sizes and context lengths.

Comparisons on Long-document QA. We finetune LLaMA2-7B and LLaMA2-13B models and report the evaluation metrics for EM Score and MMLU in Table 2. We compare our CCA-Attention with LongLoRA (Chen et al., 2024), StreamingLLM (Xiao et al., 2024b), and LM-Infinite (Han et al., 2023) across varying context lengths: 4K, 8K, 16K, and 32K. The results indicate that CCA-LLM consistently achieves the highest EM Score across all methods, showcasing enhanced capability for long sequence modeling. Additionally, our CCA-LLM also achieves comparable MMLU with other methods. These results show that our CCA-LLM maintains robustness in general knowledge QA and short context modeling. This is attributed to our globality-pooling attention, which leverages core tokens for attention, ensuring comprehensive information capture. This contrasts with other methods that selectively attend to tokens, overlooking critical information. We find CCA-LLM trained on 8K context obtains worse MMLU than that on 16K context. This could potentially be ascribed to the training bias arising from the truncation of data from diverse domains. Our training samples are generated by sampling within or concatenating across domains to form 80K-length sequences following (Cerebras, 2024; Fu et al., 2024). When truncating these sequences to the target context length (e.g., 8K) and discarding the remaining parts, it leads to a shift in data distribution. Such a shift due to truncation might have caused the initial decrease in MMLU.

In practical scenarios, it is conceivable that a model trained with a shorter context length (e.g., 4K context) could be deployed in environments with longer contexts (e.g., 16K context). Our CCA-LLM is designed to perform effectively during inference, even when extended to accommodate longer contexts. From the results in Table 2, our CCA-LLM surpasses the full self-attention baseline in EM Score, with a significant margin (25.05 vs. 0.04 under 8K context with LLaMA-2 7B). Furthermore, our CCA-LLM not only matches but also exceeds the performance of other models such as LongLoRA (25.05 vs. 16.46 under 8K context with LLaMA-2 7B).

Our CCA-LLM shows much better performance than vanilla self-attention in terms of EM score (31.50 vs. 17.5) and 7.9x inference speedup with a context length of 128K. The advantages of our method become more prominent as the length of the context increases, while the performance of vanilla self-attention may even decrease when the context length becomes very large. The reason is that in an extremely long context, non-core contexts (i.e., the irrelevant context) will be compressed by the proposed weighted pooling. In this way, CCA-Attention not only alleviates the redundant context issue and thus improves the long-context modeling performance.

Comparisons on Longbench-E. We conduct further experiments on Longbench-E (Bai et al., 2023) using our CCA-Attention and baseline methods, such as StreamingLLM (Xiao et al., 2024b), LM-Infinite (Han et al., 2023), InfLLM (Xiao et al., 2024a), and MInference (Jiang et al., 2024). As shown in Table 3, 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 than vanilla self-attention) and lower memory consumption (46% less than vanilla self-attention).

We compare our CCA-Attention with full self-attention and SinkAttention (Xiao et al., 2024b) in terms of running speed and memory footprint during forward-propagation on NVIDIA A800. The efficiency performance was assessed across a range of input sequence lengths, i.e., $\{$4K, 8K, 16K, 32K$\}$. In Figure 4, our CCA-Attention achieves a 3.5$\times$ inference speed than full self-attention (e.g., 2.6s vs. 9.2s in 32K context). Our CCA-Attention also exhibits a reduced GPU memory footprint (19.1GB vs. 35.5GB in 32K context). If the context length extends to 64K, our CCA-Attention achieves a 5.7$\times$ inference speed (e.g., 5.7s vs. 32.4s) and significantly reduced memory usage (e.g., 33.9GB vs. 60.0GB in 64K context) compared with full self-attention. Note that LongLoRA (Chen et al., 2024) only employ its $S^{2}$-Attention in training, and adopt vanilla self-attention in inference. In this sense, its inference speed and memory consumption are the same as vanilla self-attention.