Core Context Aware Transformers for Long Context Language Modeling

本文提出了一种核心上下文感知注意力机制（CCA-Attention），通过全局感知池化和局部保持模块减少长上下文建模中的冗余信息，在保持性能的同时显著提升计算效率，实验表明在 128K 上下文下实现了 7.9 倍加速和约 45% 内存减少。

Large Language Model, Long Context, Efficiency, Transformer, Representation Learning

Yaofo Chen, Zeng You, Shuhai Zhang, Haokun Li, Yirui Li, Yaowei Wang, Mingkui Tan

South China University of Technology, Peng Cheng Laboratory, Pazhou Laboratory, Harbin Institute of Technology, Key Laboratory of Big Data and Intelligent Robot, Ministry of Education

Generated by grok-3

Background Problem

大型语言模型（LLMs）在处理长上下文任务（如文档级理解、复杂推理）时表现出色，但随着上下文长度（如128K tokens）的增加，自注意力机制的计算复杂度和存储开销呈平方级增长，同时上下文中的冗余信息可能损害模型的表示能力。本文旨在解决这一问题，通过提出一种高效的注意力机制，减少冗余上下文的影响，同时保持长上下文建模能力，并在计算效率上取得显著提升。

Method

本文提出了一种即插即用的核心上下文感知（Core Context Aware, CCA）注意力机制，主要包括两个互补模块：

• 全局感知池化模块：将输入 token 序列分组，每组通过基于重要性的加权池化压缩为一个核心 token（core token），使用这些核心 token 替代原始 token 进行注意力计算，从而减少计算复杂度并聚焦于核心上下文。具体方法是利用组内最后一个 token 的查询向量评估组内 token 重要性，公式为：$\mathbf{c}_{i} = \text{softmax}\left(\frac{\mathbf{Q}_{ig}\mathbf{K}_{\mathcal{I}_{i}}^{\prime\top}}{\sqrt{d}}\right)\mathbf{X}_{\mathcal{I}_{i}}^{\text{G}}$，其中核心 token 代表组内关键信息。
• 局部保持模块：通过关注邻近 token 保留局部上下文，捕捉细粒度的局部依赖，作为全局模块的补充，确保不丢失近距离语义信息。
• 融合策略：通过可微分的融合方法，将全局和局部模块的注意力结果结合，公式为：$\mathbf{Att}_i = \text{softmax}\Big(\frac{\mathbf{Q}_i[\widetilde{\mathbf{K}}^{\text{G}}_{\mathcal{T}_i}; \widetilde{\mathbf{K}}^{\text{L}}_{\mathcal{U}_i}]^\top)}{\sqrt{d}}\Big) [\widetilde{\mathbf{V}}^{\text{G}}_{\mathcal{T}_i}; \widetilde{\mathbf{V}}^{\text{L}}_{\mathcal{U}_i}]$，从而兼顾长距离和短距离依赖。

批判性思考：虽然 CCA-Attention 的设计理念在减少冗余和降低复杂度方面有创新，但核心 token 的提取依赖于组内最后一个 token 的重要性评估，这种方法可能在某些情况下（如关键信息不在组尾）导致信息丢失。此外，动态池化策略虽然灵活，但缺乏对不同任务适应性的深入探讨，可能存在过拟合特定上下文分布的风险。

Experiment

实验在多个长上下文基准数据集上进行，包括 LongBench-E 和多文档问答任务，使用了 LLaMA2-7B、LLaMA3.1-8B 和 Qwen2.5-7B 等模型，上下文长度覆盖 4K 到 128K。实验设置包括与 StreamingLLM、LM-Infinite 和 MInference 等方法的对比，评估指标涵盖准确性（如 EM Score）、计算效率（首 token 延迟 FTL）和内存占用。

• 结果：在 LongBench-E 上，CCA-Attention 在 LLaMA2-7B-32K 和 80K 模型中取得了接近或优于全自注意力的平均得分（例如 22.24 vs. 22.42），同时在 128K 上下文下实现了 7.9 倍的推理加速（15.89s vs. 124.85s）和 44%-49% 的内存减少。在多文档问答任务中，CCA-Attention 在长上下文（如 64K 和 128K）下表现出色，EM 分数高于全自注意力（35.3 vs. 34.6）。
• 分析与批判：实验设置较为全面，覆盖了多种模型和上下文长度，但数据集选择可能偏向于特定任务（如问答和总结），缺乏对其他长上下文任务（如生成任务）的验证，可能存在结果选择性报告的风险。此外，虽然效率提升显著，但性能在某些任务（如 Few-Shot Learning）上略有下降，表明方法在捕捉某些细粒度信息时可能存在局限。动态调整 group size 和 local window size 的灵活性实验虽然展示了性能-效率的权衡，但未深入探讨这种调整对模型泛化能力的影响，实验设计可以更全面。

Further Thoughts

CCA-Attention 的设计理念为长上下文建模提供了一种新思路，尤其是在计算资源受限的场景下，其线性复杂度和内存效率的提升具有重要意义。然而，我认为其核心 token 提取机制可能在某些任务中（如需要捕捉稀疏但关键的上下文信息）表现不佳，未来可以探索更自适应的池化策略，例如结合任务特定的注意力分布。此外，论文中提到的动态调整 group size 和 local window size 的灵活性启发了我思考如何将这种机制与其他领域（如多模态模型）结合，特别是在处理长序列视频或音频数据时，是否可以通过类似的核心信息提取来减少冗余计算？另外，与最近的一些工作（如基于状态空间模型的长序列建模）相比，CCA-Attention 在理论上是否能进一步优化其全局-局部融合策略，可能是值得深入研究的方向。