ICML 2025 | Transformer 性能大提升：CCA-Attention 替代自注意力模块，LLM长文本建模突破-优快云博客

一、导读

随着大语言模型（LLM）在自然语言处理中的广泛应用，长上下文建模已成为其关键能力之一。然而，随着上下文长度的增加，冗余信息积累带来的计算开销成了一个亟待解决的问题。南方科技大学与鹏城实验室的研究团队提出了 Core Context Aware Attention (CCA-Attention)，该方法通过创新的全局感知池化与局部保留模块，显著减少冗余信息并提高长上下文建模的效率。实验结果表明，CCA-Attention 在处理长上下文时的计算效率和性能均优于现有方法。

论文基本信息

论文标题：Core Context Aware Transformers for Long Context Language Modeling
作者：Yaofo Chen¹, Zeng You¹², Shuhai Zhang¹³, Haokun Li¹², Yiru Li¹, Yaowei Wang²⁴, Mingkui Tan¹³⁵

二、摘要

近年来，长上下文建模成为大语言模型（LLMs）中的一个重要研究方向。现有的自注意力机制面临着冗余信息和计算开销的问题，影响了模型性能。为了解决这一问题，本文提出了 Core Context Aware Attention (CCA-Attention) 机制。该机制通过 全局感知池化 和 局部保留模块，有效减少冗余信息，保留核心上下文，并提高计算效率。实验表明，CCA-Attention 在多个长上下文任务中显著提高了计算效率，并在性能上超过了传统方法。此外，本文还提出了三种微调策略，展示了该方法在现有大语言模型中的高效性和灵活性。

二、相关工作

这一部分回顾了基于自注意力的Transformer方法及其局限性。自注意力的计算复杂度随着上下文长度的增加呈二次增长，这对于长文本任务非常低效。现有的一些高效注意力方法，如稀疏注意力（Longformer、BigBird）和线性注意力（Linear Transformer、RetNet），虽然通过减少计算量来提高效率，但通常会牺牲信息交换的全面性，从而影响长上下文的建模能力。

现有方法包括：

稀疏注意力：Longformer和BigBird等方法使用稀疏注意力，通过只计算固定间隔的tokens之间的注意力，减少计算开销。
线性注意力：Linear Transformer等方法将自注意力计算重新定义为线性复杂度的近似计算。

这些方法虽然提高了计算效率，但在处理长文本时仍然面临冗余信息的计算问题。本文提出的CCA-Attention通过结合全球感知池化模块和保留局部性模块，在减少冗余信息的同时，保持了信息交换的全面性，解决了现有方法的不足。

三、方法与实现

3.1 Motivation and Method Overview

传统的自注意力机制随着上下文长度的增加，计算复杂度为，在处理非常长的文本时会导致计算和存储开销巨大。冗余信息成为问题，因为并非所有上下文信息对于目标token的表示都有重要作用。为了缓解冗余信息的影响，提出了核心上下文感知注意力（CCA-Attention），该方法包括两个关键模块：

全球感知池化模块：将输入tokens分组，并根据每组token的重要性计算核心token，用于后续的注意力计算，从而减少冗余信息。
保留局部性模块：该模块确保每个token与其邻近的tokens进行交互，从而捕捉局部上下文。

通过这两个模块，CCA-Attention能够有效降低冗余信息，减少计算开销，同时保持对长距离依赖的建模能力。

在这里插入图片描述

3.2 全球感知池化模块

为了有效减少上下文中的冗余，全球感知池化模块将输入tokens分成若干组，并根据每组的权重计算出一个核心token。该模块的关键创新是通过每组的最后一个token计算其重要性，并将其作为核心token，代替原始tokens进行后续的注意力计算。

计算公式：

其中，是第i组的查询token，是第i组的键，是该组的tokens。通过这种方式，核心token捕捉了该组中最重要的信息。

通过这种核心token的替代，模型将注意力计算的复杂度从原始tokens的降低为基于核心tokens的，显著减少了计算和存储开销。

3.3 保留局部性模块

保留局部性模块用于补充全球感知池化模块，专注于捕捉邻近tokens的局部信息。每个token在计算注意力时，会考虑其邻近的tokens，确保捕捉到局部的上下文依赖关系。这对于细粒度的上下文建模至关重要。

计算公式：

通过这一设计，保留了局部上下文的精细信息，确保模型不仅能够理解长距离的全局依赖，还能够处理局部的细节信息。

3.4 可微分融合策略

为了解决全球感知池化模块和保留局部性模块各自的局限性，提出了可微分融合策略。通过将这两种模块的输出（全局和局部的键值矩阵）进行拼接，模型能够有效融合全球信息和局部信息，从而增强模型对长上下文的理解。

最终的注意力计算公式为：

这种融合策略确保了全局和局部信息的互补，使得模型在推理时能够充分利用这两种信息源，从而有效提高长上下文建模的能力。

3.5 微调策略

为了便于集成到现有的预训练模型中，CCA-Attention设计为插件式模块，能够替换掉传统的自注意力模块。训练时，CCA-Attention能够与现有的预训练模型兼容，并只需最小的微调即可适应新的注意力机制。与需要重新训练的大型线性注意力方法相比，CCA-Attention的微调成本极低，使得其更适用于实际应用中。

本研究提出了三种微调策略：

从头训练：完全从零开始训练模型，适用于需要高性能且计算资源充足的场景。
完全微调：微调所有模型参数，适用于大多数应用场景。
部分微调：只微调可学习的参数（如），适用于需要快速训练或计算资源有限的场景。

3.6 计算和存储复杂度分析

在计算和存储复杂度分析中，CCA-Attention展示了相较于传统的自注意力方法显著的效率提升：

对于固定长度序列，计算复杂度为，其中为每组的核心token数量，为局部窗口大小。相比于传统方法的，减少了计算量。
对于变长序列，计算复杂度为，其中是组大小。

在存储方面，CCA-Attention仅存储核心tokens的键值矩阵，存储复杂度为，显著减少了内存占用。

四、实验

4.1 实验设置

为评估CCA-Attention的性能，本文通过与现有的高效注意力方法（如StreamingLLM、LM-Infinite、MInference）进行对比，进行了一系列实验。实验的目的是验证CCA-Attention在长上下文建模中的优势，尤其是在计算效率和精度方面。

基准数据集：

RealEstate10K数据集：该数据集用于评估相机姿态估计、深度估计和点云重建等任务。

4.2 长上下文建模

实验结果表明，CCA-Attention在长上下文建模和计算效率方面均优于现有方法。例如，在RealEstate10K数据集上的实验中，CCA-Attention的相机姿态估计AUC@30为85.3，显著超越了DUSt3R的67.7，且推理时间仅为0.2秒。

4.3 其推理灵活性

CCA-Attention展示了其推理灵活性。通过调整组大小和局部窗口大小，CCA-Attention能够根据不同的用户流量需求动态调整推理速度和精度。这一灵活性使得CCA-Attention在高峰时段能够优化吞吐量，而在低峰时段提升模型的性能。

灵活性

4.4 计算效率

通过与LLaMA-2模型进行对比，CCA-Attention在推理延迟和内存占用方面均显示出显著的优势。实验结果表明，CCA-Attention能够在保持高精度的同时，减少推理延迟和内存使用。

五、结论

本文提出的Core Context Aware Attention (CCA-Attention)，一种用于长上下文建模的高效注意力机制。实验结果表明，CCA-Attention在减少计算开销和存储需求的同时，显著提升了长上下文建模的性能。该方法能够无缝集成到现有的预训练语言模型中，并且仅需最小的微调。

六、快速开始

我们提供了一种简单的方法，可以将Transformer模型系列（如Llama、GPT等）中的自注意力模块替换为我们的CCA-Attention模块。以下是一个将Llama模型中的自注意力模块替换为CCA-Attention模块的示例。详见: https://github.com/chenyaofo/CCA-Attention

 # if using KV cache compression with cca-attention, please import 'codebase.core.patch_simple' at the first line

# import codebase.core.patch_simple

import transformers
from codebase.core.simple_pooling_attn import replace_attn_forward

# Replace the self-attention module with CCA-Attention module prior to all codes
replace_attn_forward(group_size, window_size, pool_func="cca")

# Load the pre-trained model
model = transformers.AutoModelForCausalLM.from_pretrained(
  model_name_or_path,
  torch_dtype=torch.bfloat16,
 )
tokenizer = transformers.AutoTokenizer.from_pretrained(
  model_name_or_path,
  use_fast=True,
 )

如何学习大模型 AI ？

由于新岗位的生产效率，要优于被取代岗位的生产效率，所以实际上整个社会的生产效率是提升的。

但是具体到个人，只能说是：

“最先掌握AI的人，将会比较晚掌握AI的人有竞争优势”。

这句话，放在计算机、互联网、移动互联网的开局时期，都是一样的道理。

我在一线互联网企业工作十余年里，指导过不少同行后辈。帮助很多人得到了学习和成长。

我意识到有很多经验和知识值得分享给大家，也可以通过我们的能力和经验解答大家在人工智能学习中的很多困惑，所以在工作繁忙的情况下还是坚持各种整理和分享。但苦于知识传播途径有限，很多互联网行业朋友无法获得正确的资料得到学习提升，故此将重要的AI大模型资料包括AI大模型入门学习思维导图、精品AI大模型学习书籍手册、视频教程、实战学习等录播视频免费分享出来。

在这里插入图片描述