Core Context Aware Transformers for Long Context Language Modeling
核心上下文感知 Transformer 用于长上下文语言建模

Most existing attention-based large language models (LLMs), such as GPT (Brown et al., 2020; OpenAI, 2023) and LLaMA (Touvron et al., 2023a), employ the next-token prediction (Vaswani et al., 2017) paradigm to generate text. Given a sequence of tokens $\mathbf{X}\small{=}[\mathbf{x}_{1};\mathbf{x}_{2};\dots;\mathbf{x}_{L}]$, where each token $\mathbf{x}_{i}\in\mathbb{R}^{1\times d}$, the model $\theta$ predicts the next token $\mathbf{x}_{t}$ by conditioning on all preceding tokens as its context $C(\mathbf{x}_{t}){=}[\mathbf{x}_{1:t-1}]$. Specifically, the model generates the next token with the highest probability as:
大多数现有的基于注意力的大型语言模型（LLMs），如 GPT（Brown 等，2020；OpenAI，2023）和 LLaMA（Touvron 等，2023a），采用下一个标记预测（Vaswani 等，2017）范式来生成文本。给定一个标记序列 $\mathbf{X}\small{=}[\mathbf{x}_{1};\mathbf{x}_{2};\dots;\mathbf{x}_{L}]$ ，其中每个标记 $\mathbf{x}_{i}\in\mathbb{R}^{1\times d}$ ，模型 $\theta$ 通过条件化在所有先前标记作为其上下文 $C(\mathbf{x}_{t}){=}[\mathbf{x}_{1:t-1}]$ 上来预测下一个标记 $\mathbf{x}_{t}$ 。具体来说，模型通过以下方式生成下一个概率最高的标记：

However, as the context length $L$ grows, the context inevitably exhibits redundant information. This redundancy stems from the inherent nature of natural language: not all contextual information is equally important for the representation of the target token. These redundant context (e.g., $C^{1}(\mathbf{x}_{t})$ w.r.t. $\mathbf{x}_{t}$ in Figure 1) have weak semantic relevance to the token $\mathbf{x}_{t}$, while introducing significant computational overhead. In contrast, the core context (e.g., $C^{2}(\mathbf{x}_{t})$ w.r.t. $\mathbf{x}_{t}$ in Figure 1) refers to the contextual information that is highly relevant to the token $\mathbf{x}_{t}$. This information is crucial for the token’s representation. Therefore, in long-context modeling, the model should prioritize the core context over redundant parts. To this end, most existing methods (Beltagy et al., 2020; Ding et al., 2023; Chen et al., 2024) employ sparse attention with predefined and fixed patterns. Unfortunately, they often overlook the importance of maintaining comprehensive information exchange among tokens, which may hinder the performance of long-context modeling tasks.
然而，随着上下文长度 $L$ 的增长，上下文不可避免地会表现出冗余信息。这种冗余源于自然语言的固有性质：并非所有上下文信息对目标标记的表示同等重要。这些冗余上下文（例如， $C^{1}(\mathbf{x}_{t})$ 相对于 $\mathbf{x}_{t}$ 在图 1 中）与标记 $\mathbf{x}_{t}$ 的语义相关性较弱，同时引入了显著的计算开销。相比之下，核心上下文（例如， $C^{2}(\mathbf{x}_{t})$ 相对于 $\mathbf{x}_{t}$ 在图 1 中）指的是与标记 $\mathbf{x}_{t}$ 高度相关的上下文信息。这些信息对于标记的表示至关重要。因此，在长上下文建模中，模型应优先考虑核心上下文而非冗余部分。为此，大多数现有方法（Beltagy et al., 2020；Ding et al., 2023；Chen et al., 2024）采用预定义和固定的稀疏注意力模式。不幸的是，它们往往忽视了在标记之间保持全面信息交换的重要性，这可能会阻碍长上下文建模任务的性能。

In this paper, we seek to reduce the context redundancy associated with full self-attention.
在本文中，我们寻求减少与全自注意力相关的上下文冗余。
To achieve this, we propose a Core Context Aware Attention (CCA-Attention), which employs globality-aware pooling and locality-preserving modules to capture both global and local dependencies within a long context. As shown in Figure 2, the globality-aware pooling module 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 context redundancy and computational cost (see Section 3.2). However, the globality-aware pooling module mainly focuses on long-range and coarse-grained information and overlooks local context.
为实现这一目标，我们提出了一种核心上下文感知注意力（CCA-Attention），该注意力机制采用全局感知池化和局部保持模块来捕捉长上下文中的全局和局部依赖关系。如图 2 所示，全局感知池化模块通过从输入序列的分割组中生成代表性核心标记来工作。然后，它使用这些减少数量的核心标记来计算注意力，从而减少上下文冗余和计算成本（见第 3.2 节）。然而，全局感知池化模块主要关注长距离和粗粒度信息，而忽略了局部上下文。
To address this limitation, the locality-preserving module is responsible for capturing the local information of the neighborhood tokens to ensure comprehensive coverage (see Section 3.3). Furthermore, we devise a differentiable fusion strategy to combine the insights from global and local modules (see Section 3.4). This is crucial as it retains the comprehensive understanding ability of the full self-attention within our CCA-Attention. The pseudo-code for our proposed CCA-Attention is presented in Algorithm 1.
为了解决这一限制，局部保持模块负责捕捉邻域标记的局部信息，以确保全面覆盖（见第 3.3 节）。此外，我们设计了一种可微的融合策略，以结合全局和局部模块的见解（见第 3.4 节）。这一点至关重要，因为它保留了我们在 CCA-Attention 中全自注意力机制的全面理解能力。我们提出的 CCA-Attention 的伪代码在算法 1 中给出。

The context redundancy aforementioned indicates that computational resources can be dynamically allocated to core contexts while reducing emphasis on the remaining ones. This could approximate the full self-attention with both reduced redundancy and computational overhead.
上述上下文冗余表明，可以将计算资源动态分配给核心上下文，同时减少对其他上下文的重视。这可以近似全自注意力，同时减少冗余和计算开销。
Motivated by this, we propose a globality-aware pooling module that dynamically identifies prominent tokens and encapsulates them 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 $g$ tokens, in total $m\small{=}\lfloor L/g\rfloor$ groups. For simplicity, we denote the $i$-th group by $\mathbf{X}^{\text{G}}_{\mathcal{I}_{i}}\small{\in}\mathbb{R}^{g\times d}$, where $\mathcal{I}_{i}{=}\{(i{-}1)g+1,(i{-}1)g+2,\ldots,ig\}$ with ${\bf x}_{ig}$ denoting the last token in the $i$-th group. To identify prominent tokens in the $i$-th group, we devise a group-wise weighted pooling strategy that employs the last token ${\bf x}_{ig}$ to evaluate the importance. This is inspired by attention map visualizations (Section C.4), which show that important tokens consistently receive high attention scores from subsequent tokens, indicating their significant influence regardless of position within the group. Formally, we derive one core token $\mathbf{c}_{i}$ from each group by
给定一个标记序列 $\mathbf{X}{=}[\mathbf{x}_{1};\mathbf{x}_{2};\dots;\mathbf{x}_{L}]$ ，我们将输入序列 $\mathbf{X}$ 分段，每个组包含 $g$ 个标记，总共 $m\small{=}\lfloor L/g\rfloor$ 个组。为了简便，我们用 $\mathbf{X}^{\text{G}}_{\mathcal{I}_{i}}\small{\in}\mathbb{R}^{g\times d}$ 表示第 $i$ 组，其中 $\mathcal{I}_{i}{=}\{(i{-}1)g+1,(i{-}1)g+2,\ldots,ig\}$ 与 ${\bf x}_{ig}$ 表示第 $i$ 组中的最后一个标记。为了识别第 $i$ 组中的突出标记，我们设计了一种组内加权池化策略，该策略使用最后一个标记 ${\bf x}_{ig}$ 来评估重要性。这受到注意力图可视化（第 C.4 节）的启发，这些可视化显示重要标记始终从后续标记中获得高注意力分数，表明它们在组内无论位置如何都具有显著影响。正式地，我们从每个组中推导出一个核心标记 $\mathbf{c}_{i}$ 。

where $\mathbf{Q}_{ig}$ is the query vector for the last token in the $i$-th group of $\mathbf{Q}_{\mathcal{I}_{i}}{=}{\bf X}^{\text{G}}_{\mathcal{I}_{i}}{\bf W}^{Q}$ and $\mathbf{K}^{{}^{\prime}}_{\mathcal{I}_{i}}{=}{\bf X}^{\text{G}}_{\mathcal{I}_{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. With $m$ groups in the input sequence $\mathbf{X}$, we derive $m$ core tokens in total, i.e., $\mathbf{C}{=}[\mathbf{c}_{1};\mathbf{c}_{2};\dots;\mathbf{c}_{m}]$.
其中 $\mathbf{Q}_{ig}$ 是第 2#和第 3#组中第 1#组的最后一个查询向量， $\mathbf{W}^{Q}$ 和 $\mathbf{W}^{K}$ 是可学习参数。这样，核心标记 $\mathbf{c}_{i}$ 封装了相应组的关键信息。在输入序列 $\mathbf{X}$ 中有 $m$ 组，我们总共导出 $m$ 个核心标记，即 $\mathbf{C}{=}[\mathbf{c}_{1};\mathbf{c}_{2};\dots;\mathbf{c}_{m}]$ 。

To reduce the redundancy, we use the sequence of core tokens $\mathbf{C}{=}[\mathbf{c}_{1};\mathbf{c}_{2};\dots;\mathbf{c}_{m}]$ instead of the original tokens $\mathbf{X}$ for attention computation. This substitution reduces the dimensionality from $\mathbf{X}\small{\in}\mathbb{R}^{L\times d}$ to $\mathbf{C}\small{\in}\mathbb{R}^{m\times d}$, thereby reducing the computational and storage complexity. For each query $\mathbf{Q}_{i}$, tokens that are distant from it typically exhibit lower relevance and are more likely to be redundant. Formally, we adopt core tokens to calculate the key and value matrices $\mathbf{K}^{\text{G}}$, $\mathbf{V}^{\text{G}}$ for these tokens as follows
为了减少冗余，我们在注意力计算中使用核心标记序列 $\mathbf{C}{=}[\mathbf{c}_{1};\mathbf{c}_{2};\dots;\mathbf{c}_{m}]$ 而不是原始标记 $\mathbf{X}$ 。这种替换将维度从 $\mathbf{X}\small{\in}\mathbb{R}^{L\times d}$ 降低到 $\mathbf{C}\small{\in}\mathbb{R}^{m\times d}$ ，从而减少了计算和存储的复杂性。对于每个查询 $\mathbf{Q}_{i}$ ，距离它较远的标记通常具有较低的相关性，更可能是冗余的。正式来说，我们采用核心标记来计算这些标记的关键矩阵 $\mathbf{K}^{\text{G}}$ 和值矩阵 $\mathbf{V}^{\text{G}}$ ，如下所示

where ${\bf W}^{K}$ and ${\bf W}^{V}$ is learnable parameters. In contrast, tokens in close proximity to the query $\mathbf{Q}_{i}$ are likely to be more relevant. We retain $s$ nearest tokens for a fine-grained attention computation (discussed in detail in Section 3.3). Thus, the index $j$ in Eqn. (3) can be calculated as ${j\small{=}\max(0,\lfloor(i\small{-}s)/g\rfloor)}$. When the context is short (i.e., $i\small{<}(g+s)$), the key and value $\mathbf{K}^{\text{G}}$, $\mathbf{V}^{\text{G}}$ would be excluded from attention calculation since the redundancy in the context is negligible. During inference, as tokens are generated sequentially, we derive a new core token via Eqn. (2) once the number of generated tokens reaches $g$. Different from the full self-attention, we cache $\widetilde{\mathbf{K}}^{\text{G}}$ and $\widetilde{\mathbf{V}}^{\text{G}}$ for inference.
其中 ${\bf W}^{K}$ 和 ${\bf W}^{V}$ 是可学习的参数。相比之下，靠近查询 $\mathbf{Q}_{i}$ 的标记更有可能是相关的。我们保留 $s$ 个最近的标记进行细粒度注意力计算（在 3.3 节中详细讨论）。因此，方程（3）中的索引 $j$ 可以计算为 ${j\small{=}\max(0,\lfloor(i\small{-}s)/g\rfloor)}$ 。当上下文较短（即 $i\small{<}(g+s)$ ）时，由于上下文中的冗余可以忽略不计，因此会排除关键 $\mathbf{K}^{\text{G}}$ 和值 $\mathbf{V}^{\text{G}}$ 的注意力计算。在推理过程中，随着标记的顺序生成，一旦生成的标记数量达到 $g$ ，我们就会通过方程（2）推导出一个新的核心标记。与全自注意力不同，我们在推理时缓存 $\widetilde{\mathbf{K}}^{\text{G}}$ 和 $\widetilde{\mathbf{V}}^{\text{G}}$ 。

As mentioned above, the globality-aware pooling module effectively captures long-range dependencies by compressing input tokens into core tokens. It focuses on coarse-grained global information, potentially overlooking fine-grained local context. However, recent studies (Manakul & Gales, 2021; Yang et al., 2021) demonstrate that local context plays a critical role in many language modeling tasks. To address this, we introduce a locality-preserving module that complements the globality-aware pooling module by focusing on neighboring tokens to capture detailed local dependencies.
如上所述，全局感知池化模块通过压缩输入标记为核心标记，有效地捕捉了长距离依赖关系。它关注粗粒度的全局信息，可能忽略了细粒度的局部上下文。然而，最近的研究（Manakul & Gales，2021；Yang 等，2021）表明，局部上下文在许多语言建模任务中起着关键作用。为了解决这个问题，我们引入了一个保留局部性的模块，通过关注邻近标记来捕捉详细的局部依赖关系，以此补充全局感知池化模块。

To be specific, the locality-preserving module ensures that each query $\mathbf{Q}_{i}$ attends to the preceding at least $s$ tokens to capture local dependencies. During the generation process, it is challenging to maintain the number of tokens as a multiple of the group size $g$. To address this, we set the local window size to $s\small{+}((i-s)\bmod g)$. This strategy ensures that each key token participates in the attention computation with the query $\mathbf{Q}_{i}$. Consequently, the key and value matrices for a specific query $\mathbf{Q}_{i}$ in the locality-preserving module are defined as follows:
具体来说，局部保持模块确保每个查询 $\mathbf{Q}_{i}$ 至少关注前 $s$ 个标记以捕获局部依赖。在生成过程中，保持标记数量为组大小的倍数 $g$ 是具有挑战性的。为了解决这个问题，我们将局部窗口大小设置为 $s\small{+}((i-s)\bmod g)$ 。这种策略确保每个关键标记都参与与查询 $\mathbf{Q}_{i}$ 的注意力计算。因此，局部保持模块中特定查询 $\mathbf{Q}_{i}$ 的关键和值矩阵定义如下：

where $k\small{=}\max(1,i\small{-}s\small{-}((i\small{-}s)\bmod g))$. Note that the locality-preserving module shares the projection parameters ${\bf W}^{Q}$, ${\bf W}^{K}$, and ${\bf W}^{V}$ with the globality-aware pooling module, thereby incurring no additional projection parameters.
其中 $k\small{=}\max(1,i\small{-}s\small{-}((i\small{-}s)\bmod g))$ 。请注意，局部保持模块与全局感知池化模块共享投影参数 ${\bf W}^{Q}$ 、 ${\bf W}^{K}$ 和 ${\bf W}^{V}$ ，因此不会产生额外的投影参数。

Both globality-aware pooling and locality-preserving modules involve only a portion of tokens in the attention computation, leading to a limited comprehensive understanding of the context.
全局感知池化和局部保持模块仅涉及注意力计算中的一部分标记，导致对上下文的全面理解有限。
To address this limitation, we seek to combine the involved tokens of these two attentions to integrate the insights they provide. Specifically, we concatenate the key and value matrices from both attentions, i.e., $[{\widetilde{\mathbf{K}}_{\mathcal{T}_{i}}^{\text{G}}};{\widetilde{\mathbf{K}}_{\mathcal{U}_{i}}^{\text{L}}}]$ and $[\widetilde{\mathbf{V}}^{\text{G}}_{\mathcal{T}_{i}};\widetilde{\mathbf{V}}_{\mathcal{U}_{i}}^{\text{L}}]$, to leverage the combined information. Formally, the proposed CCA-Attention is computed as follows:
为了解决这一局限性，我们寻求结合这两种注意力的相关标记，以整合它们提供的见解。具体来说，我们将来自两种注意力的键和值矩阵（即 $[{\widetilde{\mathbf{K}}_{\mathcal{T}_{i}}^{\text{G}}};{\widetilde{\mathbf{K}}_{\mathcal{U}_{i}}^{\text{L}}}]$ 和 $[\widetilde{\mathbf{V}}^{\text{G}}_{\mathcal{T}_{i}};\widetilde{\mathbf{V}}_{\mathcal{U}_{i}}^{\text{L}}]$ ）连接起来，以利用组合信息。形式上，所提出的 CCA-Attention 的计算如下：

We represent the final output of our CCA-Attention as $\mathbf{Att}=[\mathbf{Att}_{1};\mathbf{Att}_{2};\dots;\mathbf{Att}_{L}]$. In practice, we implement our attention mechanism with Triton (Tillet et al., 2019) to accelerate both training and inference processes. This implementation enables parallel computation of each $\mathbf{Att}_{i}$ with high computational efficiency. After integrating the global and local attention, we can formalize $\mathbf{Att}$ in Eqn. (5) element-wise into the structure of full attention, as detailed in Proposition 1 in the supplementary material A. The Proposition 1 shows that each token accesses all preceding tokens, ensuring full information exchange among tokens and thus enhancing capturing long-range dependencies.
我们表示我们 CCA-Attention 的最终输出为 $\mathbf{Att}=[\mathbf{Att}_{1};\mathbf{Att}_{2};\dots;\mathbf{Att}_{L}]$ 。在实践中，我们使用 Triton（Tillet 等人，2019）实现我们的注意力机制，以加速训练和推理过程。这种实现使得每个 $\mathbf{Att}_{i}$ 可以以高计算效率进行并行计算。在整合全局和局部注意力之后，我们可以将 $\mathbf{Att}$ 在式（5）中逐元素形式化为全注意力结构，如补充材料 A 中的命题 1 所述。命题 1 表明每个标记可以访问所有前面的标记，确保标记之间充分的信息交换，从而增强捕捉长距离依赖关系。

More importantly, our CCA-Attention demonstrates dynamic flexibility through adjustable group size $g$ and local window size $s$ during inference. This architectural flexibility allows the generation of multiple model variants tailored to varying user traffic, offering a substantial advantage over the full self-attention mechanisms in real-world deployment scenarios (see results and analysis in Section 4.3).
更重要的是，我们的 CCA-Attention 通过可调整的组大小 $g$ 和局部窗口大小 $s$ 在推理过程中展现出动态灵活性。这种架构灵活性允许生成多个针对不同用户流量的模型变体，在现实部署场景中相对于全自注意力机制具有显著优势（参见第 4.3 节的结果和分析）。

Our proposed CCA-Attention is fully compatible with existing attention-based LLMs, such as the LLaMA series models (Touvron et al., 2023b; AI, 2024), serving as a plug-and-play module that can replace the full self-attention. CCA-Attention maintains alignment with full self-attention in terms of input, output, and parameter dimensions.
我们提出的 CCA-Attention 与现有的基于注意力的 LLMs（如 LLaMA 系列模型（Touvron 等，2023b；AI，2024））完全兼容，作为一个即插即用的模块，可以替换全自注意力。CCA-Attention 在输入、输出和参数维度上与全自注意力保持一致。
This ensures that only a minimal training cost is able to preserve long-context modeling capabilities while reducing computational costs. In contrast, existing linear attention approaches (Katharopoulos et al., 2020; Sun et al., 2023) introduce kernel functions for attention and necessitate training from scratch, making them less practical for real-world applications due to their inability to leverage the extensive knowledge embedded in pretrained LLMs.
这确保了只需最小的训练成本就能保留长上下文建模能力，同时降低计算成本。相比之下，现有的线性注意力方法（Katharopoulos 等人，2020 年；Sun 等人，2023 年）为注意力引入核函数，需要从头开始训练，由于它们无法利用预训练 LLMs 中嵌入的丰富知识，因此在实际应用中不太实用。

We replace the self-attention module in existing attention-based LLMs with CCA-Attention, enabling compatibility with three different training strategies: 1) Training from Scratch: This strategy involves training CCA-Attention from scratch on large-scale corpora.
我们将现有基于注意力的 LLMs 中的自注意力模块替换为 CCA-Attention，使其能够兼容三种不同的训练策略：1）从头开始训练：这种策略涉及在大型语料库上从头开始训练 CCA-Attention。
While this approach may yield superior performance by fully adapting the model to the proposed attention mechanism, it is computationally intensive and requires significant resources.
虽然这种方法可能通过完全适应所提出的注意力机制而获得更好的性能，但它计算密集，需要大量资源。
2) Full Finetuning: A more efficient alternative is to finetune all parameters of the model based on the parameters of existing pretrained LLMs. This strategy leverages the pre-trained knowledge while allowing the model to adapt fully to our attention mechanism.
2）完全微调：一种更有效的方法是基于现有预训练 LLMs 的参数对模型的所有参数进行微调。这种策略利用了预训练知识，同时允许模型完全适应我们的注意力机制。
3) Partial Finetuning: For scenarios where efficiency is critical, we can finetune only the learnable parameters of CCA-Attention, i.e., $\mathbf{W}^{Q}$, $\mathbf{W}^{K}$, and $\mathbf{W}^{V}$. This strategy requires only a modest finetuning effort on a small number of corpora, making it computationally efficient for maintaining long-context modeling capabilities.
3) 部分微调：对于效率至关重要的场景，我们只需微调 CCA-Attention 的可学习参数，即 $\mathbf{W}^{Q}$ 、 $\mathbf{W}^{K}$ 和 $\mathbf{W}^{V}$ 。这种策略只需要在小规模语料库上进行适度的微调努力，因此在维持长上下文建模能力方面具有计算效率。
We provide empirical evidence to guide the selection of the most appropriate finetuning strategy in Table 10.
我们在表 10 中提供了经验证据，以指导选择最合适的微调策略。

Compared with the full self-attention, our CCA-Attention offers significant benefits in terms of computational complexity and key-value cache storage, as analyzed below.
与全自注意力相比，我们的 CCA-Attention 在计算复杂度和键值缓存存储方面提供了显著的优势，如下所述。

Acceleration via 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. For the globality-aware pooling module, the query and key matrices encompass $L$ and $m$ tokens, respectively, resulting in a computational complexity of $\mathcal{O}(Lm)$. Regarding the locality-preserving module, each token only attends preceding $s+g$ tokens at most. With $L$ tokens in total, the upper bound of complexity amounts to $O(L(s+g))$.
通过降低计算复杂度实现加速。我们的 CCA-Attention 根据任务类型表现出不同的计算复杂度。对于具有固定长度序列的任务（如多选题问答），我们的 CCA-Attention 表现出线性计算复杂度 $O(Lm+Ls)$ ，这比全自注意力的复杂度 $O(L^{2})$ 有显著提升。在这里，我们将组数 $m$ 定义为常数。对于全局感知池化模块，查询和键矩阵分别包含 $L$ 和 $m$ 个标记，导致计算复杂度为 $\mathcal{O}(Lm)$ 。至于局部保持模块，每个标记最多只关注前面的 $s+g$ 个标记。总共有 $L$ 个标记，复杂度的上限达到 $O(L(s+g))$ 。

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 $g$ as a constant, ensuring that our CCA-Attention is able to leverage key-value caching during autoregressive token generation. Once one group has certain $g$ tokens, the corresponding core token is also determined and cached. Thus, our CCA-Attention achieves a computational complexity of $O(L^{2}/g+Ls)$. The complexity analysis follows a similar pattern to the tasks with fixed-length sequences.
对于具有可变长度序列的任务（如开放式问答），模型以自回归的方式生成后续标记。在这种情况下，我们将组大小 $g$ 设为一个常数，确保我们的 CCA-Attention 能够在自回归标记生成过程中利用键值缓存。一旦一个组有了一定的 $g$ 个标记，相应的核心标记也会被确定并缓存。因此，我们的 CCA-Attention 实现了 $O(L^{2}/g+Ls)$ 的计算复杂度。复杂度分析遵循与固定长度序列任务相似的模式。

Acceleration through Reduced Key-Value (KV) Cache. In attention-based LLMs, the KV cache leverages the autoregressive nature to store and reuse key-value pairs, thereby significantly boosting the efficiency.
通过减少键值（KV）缓存加速。在基于注意力的 LLMs 中，KV 缓存利用自回归特性来存储和重用键值对，从而显著提高效率。
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 would consume considerable memory and significant memory IO resources.
KV 缓存的大小与输入序列的长度成线性关系，在推理过程中消耗了大部分内存占用。扩大的 KV 缓存将消耗大量内存和显著的内存 I/O 资源。
Compared with full attention’s complexity of $O(L)$, our CCA-Attention has a storage complexity of $O(L/g+s)$. For the globality-aware pooling module, 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/g)$. Besides, the locality-preserving module only maintains the key and value matrices for the preceding $s$ tokens. The storage complexity for this component is $O(s)$.
与全注意力机制的复杂度 $O(L)$ 相比，我们的 CCA-Attention 具有 $O(L/g+s)$ 的存储复杂度。对于全局感知池化模块，我们只保留核心标记的关键和值矩阵，而不是所有原始标记。这降低了内存需求至 $O(L/g)$ 。此外，局部性保持模块仅维护前 $s$ 个标记的关键和值矩阵。该组件的存储复杂度为 $O(s)$ 。

We apply our CCA-Attention and compared efficient attention methods to existing pretrained LLMs. We report the performance in long context modeling and computational efficiency. We put more implementation details and ablation studies of our method in the supplementary materials.
我们应用了我们的 CCA-Attention，并将其与现有的预训练 LLMs 中的高效注意力方法进行了比较。我们报告了在长上下文建模和计算效率方面的性能。我们将更多实现细节和我们的方法的消融研究放在了补充材料中。

Dataset & Evaluation Metrics. We quantitatively evaluate our models and compare them with other considered models on the following benchmark and metric: 1) 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.
数据集与评估指标。我们对我们的模型进行定量评估，并在以下基准和指标上与其他考虑的模型进行比较：1）LongBench（Bai 等，2023）是用于评估大型语言模型长上下文理解能力的双语、多任务和全面评估的开创性基准。
It covers multiple languages like Chinese and English, consisting of 6 major categories and 21 tasks involving various application areas. 2) Exact Match Score (EM Score) (Liu et al., 2024b) is a metric for measuring the model’s ability to find the key information within a long context in a multi-document question-answering task.
它涵盖了多种语言，如中文和英语，包括 6 个大类和 21 个涉及不同应用领域的任务。2) 精确匹配得分（EM 得分）（刘等，2024b）是衡量模型在多文档问答任务中寻找长文本中关键信息能力的指标。
In this task, each test sample comprises a certain number of documents to reach the specified context length, followed by a question about the key information inserted in context.
在这个任务中，每个测试样本包含一定数量的文档以达到指定的上下文长度，随后是一个关于上下文中插入的关键信息的问题。

Implementation Details¹¹1The source code for this project is publicly available at https://github.com/chenyaofo/CCA-Attention
该项目的源代码公开可获取，链接为 https://github.com/chenyaofo/CCA-Attention.
实现细节 ¹ . We apply our proposed CCA-Attention to LLaMA2-7B-32K, LLaMA2-7B-80K (Fu et al., 2024), LLaMA3.1-8B-128K (AI, 2024) and Qwen2.5-7B-128K (Yang et al., 2024) models. We replace the full self-attention in the above LLMs with our proposed CCA-Attention. In the continuous finetuning, we adopt the SlimPajama (Cerebras, 2024) dataset, an open-source replication of the LLaMA pretraining data mixture. The number of groups in globality-aware Attention is shared across different model sizes. We finetune the full model on 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.
. 我们将提出的 CCA-Attention 应用于 LLaMA2-7B-32K、LLaMA2-7B-80K（Fu 等人，2024 年）、LLaMA3.1-8B-128K（AI，2024 年）和 Qwen2.5-7B-128K（Yang 等人，2024 年）模型。我们将上述 LLMs 中的全自注意力替换为我们提出的 CCA-Attention。在持续微调中，我们采用 SlimPajama（Cerebras，2024 年）数据集，这是 LLaMA 预训练数据混合的开源复制。全局感知注意力中的组数在不同模型大小之间共享。我们使用 1 个微批大小和 8 个梯度累积，在 A800 GPU 上对完整模型进行微调，总共 1000 个训练步骤。此训练配置适用于所有模型大小和上下文长度。
To scale the models to long contexts, we modified the “base frequency” in RoPE from 10,000 to 500,000, following (Cerebras, 2024; Xiong et al., 2024). See Section B.2 for more implementation details.
为了将模型扩展到长上下文，我们将 RoPE 中的“基本频率”从 10,000 修改为 500,000，遵循（Cerebras，2024；Xiong 等，2024）。更多实现细节请参见 B.2 节。

Compared Methods. We conduct comprehensive comparisons between our proposed CCA-Attention and several state-of-the-art methods, including LLaMA-2 with vanilla attention, StreamingLLM (Xiao et al., 2024b), LM-infinite (Han et al., 2023), and MInference (Jiang et al., 2024), across the LongBench (Bai et al., 2023) and RULER (Hsieh et al., 2024) benchmarks. Our experiments are based on the LLaMA-2 7B model fine-tuned on sequences of length 32K and 80K (Fu et al., 2024). For StreamingLLM (Xiao et al., 2024b), we use the official implementation, adjusting the attention sink to 4 and setting the attention context size to 2000. Similarly, for LM-infinite (Han et al., 2023), we follow the official code, configuring the local branch size to 1024 and the global branch size to 16. In the case of MInference (Jiang et al., 2024), we also employ the official code implementations, configured with the official settings.
与比较方法。我们对提出的 CCA-Attention 与几种最先进的方法进行了全面比较，包括使用 vanilla attention 的 LLaMA-2、StreamingLLM（Xiao et al., 2024b）、LM-infinite（Han et al., 2023）和 MInference（Jiang et al., 2024），在 LongBench（Bai et al., 2023）和 RULER（Hsieh et al., 2024）基准测试中进行比较。我们的实验基于在长度为 32K 和 80K 的序列上微调的 LLaMA-2 7B 模型（Fu et al., 2024）。对于 StreamingLLM（Xiao et al., 2024b），我们使用官方实现，将注意力下沉设置为 4，并将注意力上下文大小设置为 2000。同样，对于 LM-infinite（Han et al., 2023），我们遵循官方代码，配置本地分支大小为 1024，全局分支大小为 16。在 MInference（Jiang et al., 2024）的情况下，我们也采用了官方代码实现，并配置了官方设置。

To further validate the effectiveness of our method across different model architectures, we conduct extensive experiments on more recent foundation models: LLaMA3.1-8B-Instruct-128K and Qwen2.5-7B-128K. As shown in Table 2. our CCA-Attention demonstrates superior performance over the state-of-the-art baseline MInference across three key metrics: computational efficiency, memory reduction, and model performance.
为了进一步验证我们的方法在不同模型架构上的有效性，我们在更近期的基座模型上进行了广泛的实验：LLaMA3.1-8B-Instruct-128K 和 Qwen2.5-7B-128K。如表 2 所示，我们的 CCA-Attention 在三个关键指标上优于最先进的基线 MInference：计算效率、内存减少和模型性能。
The consistent improvements across different model architectures suggest that our approach is not limited to specific model designs.
不同模型架构的一致性改进表明，我们的方法并不仅限于特定的模型设计。

Comparisons on Long-document QA. We evaluate the performance of our CCA-Attention against other methods, including StreamingLLM (Xiao et al., 2024b), LM-Infinite (Han et al., 2023), and MInference (Jiang et al., 2024), on LLaMA2-7B-80K models using the multi-document EM Score metric. As shown in Table 3, we conduct comparisons across a range of context lengths: 4K, 8K, 16K, 32K, 64K, and 128K. For the short context (e.g., 4K), our CCA-LLM consistently achieves the highest EM score across all methods, showcasing significant capability for short sequence modeling. For example, our CCA-LLM outperforms StreamLLM (39.3 vs. 33.6) and MInference (39.3 vs. 39.0) in terms of EM score. The reason is that our CCA-Attention captures context dependencies without discarding crucial tokens. In contrast, StreamingLLM (Xiao et al., 2024b) prioritizes attention on the initial and final tokens, effectively discarding intermediate tokens, which may contain essential information. Similarly, MInference (Jiang et al., 2024) employs predefined sparse attention patterns, selectively attending to tokens and potentially overlooking critical parts of the input sequence.
在长文档问答的比较中。我们评估了我们的 CCA-Attention 与其他方法的性能，包括 StreamingLLM（Xiao 等，2024b）、LM-Infinite（Han 等，2023）和 MInference（Jiang 等，2024），在 LLaMA2-7B-80K 模型上使用多文档 EM 分数指标进行评估。如表 3 所示，我们在一系列的上下文长度上进行比较：4K、8K、16K、32K、64K 和 128K。对于短上下文（例如，4K），我们的 CCA-LLM 在所有方法中始终实现最高的 EM 分数，展示了显著的短序列建模能力。例如，我们的 CCA-LLM 在 EM 分数方面优于 StreamLLM（39.3 比 33.6）和 MInference（39.3 比 39.0）。原因是我们的 CCA-Attention 能够捕捉上下文依赖关系而不丢弃关键标记。相比之下，StreamingLLM（Xiao 等，2024b）优先关注初始和最终标记，实际上丢弃了中间标记，这些标记可能包含重要信息。同样，MInference（Jiang 等，2024）采用预定义的稀疏注意力模式，选择性地关注标记，并可能忽略输入序列的关键部分。
Both approaches risk losing important contextual information, leading to suboptimal performance in tasks requiring comprehensive understanding.
两种方法都可能导致丢失重要的上下文信息，导致在需要全面理解的任务中表现不佳。
Our method, by preserving both local and global contexts, ensures that no critical information is overlooked, thereby achieving superior performance.
我们的方法通过保留局部和全局上下文，确保不会遗漏任何关键信息，从而实现卓越的性能。

For the extremely long context (e.g., 64K and 128K), Our CCA-LLM shows much better performance than vanilla self-attention in terms of EM score (35.3 vs. 34.6) 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.
对于极长上下文（例如，64K 和 128K），我们的 CCA-LLM 在 EM 分数（35.3 比 34.6）和 128K 上下文长度下的 7.9 倍推理速度提升方面，比传统的自注意力机制表现出更好的性能。随着上下文长度的增加，我们方法的优势变得更加明显，而传统的自注意力机制的性能甚至可能在上下文长度变得非常大时下降。
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-LLM alleviates the redundant context issue and improves the long-context modeling performance.
原因是，在极其长的上下文中，非核心上下文（即无关的上下文）将被所提出的加权池化压缩。通过这种方式，CCA-LLM 缓解了冗余上下文问题，并提高了长上下文建模性能。

In real-world applications, user traffic exhibits diurnal variations. During peak traffic periods, the full self-attention struggles with throughput limitations, necessitating additional servers to accommodate the increased demand.
在现实世界的应用中，用户流量表现出日间变化。在高峰流量期间，全自注意力机制因吞吐量限制而苦苦挣扎，需要额外的服务器来满足增加的需求。
Instead, our CCA-LLM can dynamically adjust the group size $g$ and the local window size $s$ during inference to improve throughput. During off-peak traffic periods, our CCA-LLM is able to enhance model performance, albeit with a minor compromise in throughput. This strategy exhibits remarkable elasticity to address the variable user traffic challenges.
相反，我们的 CCA-LLM 可以在推理过程中动态调整组大小 $g$ 和局部窗口大小 $s$ 以提高吞吐量。在非高峰流量期间，我们的 CCA-LLM 能够提高模型性能，尽管在吞吐量上有所妥协。这种策略表现出显著的弹性，以应对可变用户流量的挑战。

To verify this, we conduct experiments with different group sizes $g$ and local window sizes $s$ during inference. In Figure 3, our CCA-LLM demonstrates an improved Longbench-E Score with a concomitant increase in computational overhead as $g$ decreases. With a reduction in the group size $g$ from 16 to 2 and the local window size $s$ from 4096 to 1024, the Longbench-E Score escalates from 20.54 to 22.69, while the latency rises from 2.80 seconds to 3.89 seconds. Despite training CCA-LLM only once, we obtain a spectrum of models, each with distinct performance and computational demands. This flexibility comes from two complementary modules: The globality-aware pooling module dynamically compress core tokens based on semantic relevance via intra-group attention, enabling adaptation to different group sizes.
为了验证这一点，我们在推理过程中对不同组大小 $g$ 和局部窗口大小 $s$ 进行了实验。如图 3 所示，随着 $g$ 的减小，我们的 CCA-LLM 在 Longbench-E 分数上有所提高，同时计算开销也随之增加。当组大小 $g$ 从 16 减少到 2，局部窗口大小 $s$ 从 4096 减少到 1024 时，Longbench-E 分数从 20.54 上升到 22.69，而延迟从 2.80 秒增加到 3.89 秒。尽管 CCA-LLM 只训练了一次，我们获得了多种模型，每种模型都有不同的性能和计算需求。这种灵活性来自于两个互补的模块：全局感知池化模块通过组内注意力动态压缩核心标记，基于语义相关性，从而能够适应不同的组大小。
The locality-preserving module uses rotary embeddings to encode relative positions, achieving translation invariance and maintaining local context across scales.
局部保持模块使用旋转嵌入来编码相对位置，实现平移不变性，并保持不同尺度下的局部上下文。

We compare our CCA-LLM with LLaMA2-7B-80K equipped with full self-attention, LM-Infinite (Han et al., 2023), and MInference (Jiang et al., 2024) in terms of inference latency and memory footprint during forward-propagation on a single NVIDIA A800 GPU. The efficiency performance was assessed across a range of input sequence lengths, i.e., $\{$4K, 8K, 16K, 32K, 64K, 128K$\}$. In Figure 4, our CCA-Attention achieves a 7.9$\times$ inference speedup than LLaMA2-7B-80K with the full self-attention (15.89s vs. 124.85s in 128K context) MInference (15.89s vs. 44.50s in 128K context) in the pre-filling stage. Our CCA-Attention also exhibits a reduced KV cache memory usage than LLaMA2-7B-80K (4.5GB vs. 64GB in 128K context) and MInference (4.5GB vs. 64GB in 128K context). Note that Minference (Chen et al., 2024) only accelerates the pre-filling stage and adopts full self-attention in the decoding stage. Moreover, it does not reduce KV cache, leading to the same memory usage as the full self-attention.
我们将我们的 CCA-LLM 与配备全自注意力机制的 LLaMA2-7B-80K、LM-Infinite（Han 等，2023）和 MInference（Jiang 等，2024）在单个 NVIDIA A800 GPU 上正向传播时的推理延迟和内存占用进行比较。效率性能在一系列输入序列长度上进行了评估，即 $\{$ 4K、8K、16K、32K、64K、128K $\}$ 。在图 4 中，我们的 CCA-Attention 在预填充阶段比配备全自注意力的 LLaMA2-7B-80K（128K 上下文中 15.89 秒 vs. 124.85 秒）和 MInference（128K 上下文中 15.89 秒 vs. 44.50 秒）实现了 7.9 $\times$ 倍的推理速度提升。我们的 CCA-Attention 在 128K 上下文中 KV 缓存内存使用量也比 LLaMA2-7B-80K（4.5GB vs. 64GB）和 MInference（4.5GB vs. 64GB）低。请注意，Minference（Chen 等，2024）仅加速预填充阶段，在解码阶段采用全自注意力机制。此外，它没有减少 KV 缓存，导致与全自注意力机制相同的内存占用。
Conversely, our CCA-LLM is able to accelerate both pre-filling and decoding stages with reduced KV cache, which is more practical in real-world applications.
相反，我们的 CCA-LLM 能够通过减少 KV 缓存来加速预填充和解码阶段，这在实际应用中更为实用。