SCCA: Shifted Cross Chunk Attention for long contextual semantic expansion

📄 arXiv: 2312.07305v1 📥 PDF

作者: Yuxiang Guo

分类: cs.CL, cs.AI

发布日期: 2023-12-12

备注: work in progress

💡 一句话要点

提出Shifted Cross Chunk Attention，扩展LLM长文本上下文能力

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 长文本建模 稀疏注意力 大型语言模型 上下文扩展 移位注意力 LLaMA2 位置插值

📋 核心要点

  1. 现有稀疏注意力方法依赖窗口自注意力，限制了全局信息的有效传递，阻碍了模型对长文本的理解。
  2. 论文提出Shifted Cross Chunk Attention (SCCA)，通过KV移位策略扩展注意力层的感受野，并结合Dilated Attention进一步提升性能。
  3. 实验结果表明，SCCA能有效扩展LLM的上下文长度，例如在LLaMA2 7B上，可将上下文从4k扩展到8k，且易于集成现有技术。

📝 摘要（中文）

本文提出Shifted Cross Chunk Attention (SCCA)，旨在解决稀疏注意力中窗口自注意力阻碍全局信息流动的问题。SCCA通过不同的KV移位策略来扩展每个注意力层的感受野。此外，结合Dilated Attention (DA)和Dilated Neighborhood Attention (DNA)，提出了Shifted Dilated Attention (SDA)。SCCA和SDA都能在多头注意力中累积注意力结果，以获得近似于完整注意力的感受野。论文通过语言建模实验，验证了SCCA的不同模式以及SCCA与SDA组合的有效性。结合位置插值(PI)和LoRA，SCCA能有效扩展大型语言模型(LLM)的上下文长度。实验表明，SCCA可以在单个V100上将LLaMA2 7B的上下文从4k扩展到8k。这种注意力模式提供了一种即插即用的微调方法，可以在保留原始架构的同时扩展模型上下文，并且与大多数现有技术兼容。

🔬 方法详解

问题定义：现有稀疏注意力方法虽然降低了计算成本，但由于其依赖于窗口自注意力，导致全局信息流动受阻，限制了模型处理长文本上下文的能力。这种局部性限制了模型对长距离依赖关系的建模，影响了语言建模等任务的性能。

核心思路：论文的核心思路是通过移位不同的Key和Value (KV) 来扩展每个注意力层的感受野。通过在不同的注意力头中使用不同的移位模式，SCCA能够使模型在多个注意力头中累积信息，从而近似于全局注意力。这种方法旨在在计算效率和全局信息获取之间取得平衡。

技术框架：SCCA可以作为现有Transformer架构中的一个模块进行替换或集成。整体流程包括：输入序列被分割成chunk，然后对每个chunk应用移位的KV注意力机制。不同的注意力头采用不同的移位模式。此外，论文还提出了SDA，它结合了SCCA、Dilated Attention (DA) 和 Dilated Neighborhood Attention (DNA)，以进一步扩展感受野。最后，通过多头注意力机制聚合来自不同移位模式的信息。

关键创新：SCCA的关键创新在于其移位的KV注意力机制，它允许模型在保持计算效率的同时，关注更广泛的上下文信息。与传统的窗口自注意力相比，SCCA通过移位操作打破了局部性限制，使得模型能够捕捉更长距离的依赖关系。SDA的提出进一步增强了感受野的扩展能力。

关键设计：SCCA的关键设计包括KV移位的具体模式。论文可能探讨了不同的移位策略，例如循环移位、随机移位等。此外，多头注意力的头数、每个头的维度、以及Dilated Attention和Dilated Neighborhood Attention的具体参数（例如dilation rate）也是重要的设计选择。损失函数通常采用标准的语言建模损失函数，例如交叉熵损失。

📊 实验亮点

实验结果表明，SCCA能够有效扩展LLM的上下文长度。例如，在LLaMA2 7B模型上，SCCA可以将上下文长度从4k扩展到8k，而无需进行大规模的重新训练。此外，SCCA与位置插值(PI)和LoRA等技术相结合，可以进一步提升性能。该方法在单个V100 GPU上即可实现，表明其具有良好的计算效率。

🎯 应用场景

SCCA具有广泛的应用前景，尤其是在需要处理长文本序列的任务中，如长文档摘要、机器翻译、代码生成、对话系统等。该方法可以作为一种即插即用的模块，方便地集成到现有的LLM中，扩展其上下文处理能力，提升模型在长文本任务上的性能。此外，SCCA的轻量级特性使其适用于资源受限的场景。

📄 摘要（原文）

Sparse attention as a efficient method can significantly decrease the computation cost, but current sparse attention tend to rely on window self attention which block the global information flow. For this problem, we present Shifted Cross Chunk Attention (SCCA), using different KV shifting strategy to extend respective field in each attention layer. Except, we combine Dilated Attention(DA) and Dilated Neighborhood Attention(DNA) to present Shifted Dilated Attention(SDA). Both SCCA and SDA can accumulate attention results in multi head attention to obtain approximate respective field in full attention. In this paper, we conduct language modeling experiments using different pattern of SCCA and combination of SCCA and SDA. The proposed shifted cross chunk attention (SCCA) can effectively extend large language models (LLMs) to longer context combined with Positional interpolation(PI) and LoRA than current sparse attention. Notably, SCCA adopts LLaMA2 7B from 4k context to 8k in single V100. This attention pattern can provide a Plug-and-play fine-tuning method to extend model context while retaining their original architectures, and is compatible with most existing techniques.