ACL 2024 | CoCA：自注意力的缺陷与改进

近年来，在大语言模型（LLM）的反复刷屏过程中，作为其内核的 Transformer 始终是绝对的主角。然而，随着业务落地的诉求逐渐强烈，有些原本不被过多关注的特性，也开始成为焦点。例如：在 Transformer 诞生之初，被视为天然具备的长度外推能力，随着相关研究的深入，人们发现，传统的 Transformer 架构在训练长度之外无一例外表现出糟糕的推理性能。

在本文中，作者从一个全新的视角，剖析了造成这种糟糕表现的可能原因，并给出了相应的解决方案。文章主要聚焦于 Self-Attention (Vaswani et al., 2017) 与 RoPE (Su et al., 2021) 的碰撞，后者是近年较多开源模型所采用的位置编码方式，例如：LLaMA (Touvron et al., 2023a) 和 Qwen (Bai et al., 2023)。

论文已被ACL 2024接收，技术细节可以查看预印版本：https://arxiv.org/abs/2309.08646

引言：

在自注意力 (Vaswani et al., 2017) 诞生之初，长度外推被认为是一种理所当然的能力。然而，随着实际应用的不断验证，这在事实上是有难度的，进而产生了一系列相关的优化工作。

现有工作通常聚焦于2个方向：注意力模块和位置编码，并有一系列优秀的工作产生。如：Longformer (Beltagy et al., 2020)、StreamingLLM (Xiao et al., 2023)、LM-Infinite (Han et al., 2023)、Alibi (Press et al., 2021)、Position Interpolation (PI) (Chen et al., 2023)、NTK-aware Scaled RoPE (bloc97, 2023)、CLEX (Chen et al., 2024) 等。

本文从一个全新的视角，揭示了自注意力与位置编码之间的内在联系（尤其是如今广泛应用的RoPE）。自注意力之中，查询和键之间天然存在的夹角，将位置编码引入了意料之外的困境，尤其是对具有关键信息的邻近位置的估计，存在不符合预期的异常行为。文章以此为切入，提出了相应的解决方案。

主要贡献如下：

揭示了自注意力与位置编码之间的一种异常行为
提出了 Collinear Constrained Attention (CoCA) 以解决上述问题
实验表明 CoCA 在长上下文处理能力比常规自注意力具有显著优势
开源了一份 CoCA 高效实现，不会增加现有计算和空间复杂度

Fig. 1. CoCA model architecture.

01 背景

旋转位置编码

理论完备性和简洁的实现，使 RoPE 成为了多数开源模型的选择。RoPE 通过旋转矩阵来编码每一个 Token 的位置信息，并利用查询和键的旋转复合，来实现相对位置的表达。

Fig. 2. rotary position embedding.(Su et al., 2021)

异常行为

在 Transformer 模型中，核心思想是计算 query 和 key 之间的关系。注意力机制使用这些关系来决定模型应该“关注”输入序列中的哪些部分。而 RoPE 利用旋转矩阵来编码位置信息的过程中，存在以下潜在的异常行为，如图 3 所示：

Fig. 3. Anomalous Behavior between RoPE and Attention Matrices.

情况（b）和（c）：这是符合预期的行为，因为 query 和 key 之间注意力得分随着 m 和 n 的距离变大而逐渐减小，符合“近大远小”的先验假设。

情况（a）和（d）：这是发生异常的行为，因为在最邻近的 Token 处，注意力得分预期之外的衰减，模型为了补偿这种衰减，必须在训练阶段给邻近 Token 补偿额外的增益，进而在长度外推过程中产生训练/推理的不一致。

02 CoCA实现

共线约束

基于上述观察，一个很自然的想法是让 Self-Attention 中的query和key初始夹角为0，这是论文中共线约束（Collinear Constrained Attention）的由来。

详细的推导和公式，这里不进行一一展开，读者可以阅读原文进行深入理解，这里只给出核心公式：

与原始的 Self-Attention 和 RoPE 相比，上述公式表达了CoCA 的核心：即在第 m 个 query 和第 n 个 key 之间建立联系，使它们的任意一个二维切片共线，从而保证 query 和 key 初始夹角为 0 。

松弛约束

然而，上述共线约束所导出的精确解仅仅在理论上可行，实际操作过程中，由于空间复杂度的问题，并不能够实现。为此，文章中给出了一种“对偶”实现，并证明了两者的等价性。

核心公式如下：

文章中证明了“对偶”实现施加以下额外约束后，等价于理论精确解：

最后，文章移除该额外约束，并得到CoCA的最终实现，这是松弛约束（Slack Constraint）的由来。

03 实验结果

长文本能力

文章分别评估了重新训练和基于LLaMA微调2种方式，在PG-19 数据集 (Rae et al., 2019)和 (Mohtashami & Jaggi, 2023) 提出的密钥检索综合评估任务，均表明CoCA相比常规的Self-Attention在长文本能力上具有显著优势。

Fig. 4. Experiment Results.

消融实验

文章对比了松弛约束和非松弛版本的模型，得到了一些出人意料的结果：即尽管模型结构一致，但松弛约束具有更大的上下文窗口，且不影响模型表达能力。

Fig. 5. Ablation study.

04 总结

文章提出了一种新的自注意力架构，以解决 RoPE 和原始 Self-Attention 之间的异常行为。这是首次对自注意力机制中查询和键的相对位置的深入研究，并由此发现了此前被忽视的位置编码异常。文章进一步导出了 CoCA 的松弛实现，并在大量实验上表明了该方法在长文本扩展上的优越性。同时，CoCA 与其他优化方法的兼容性，也为其未来的实用价值提供了基础。

CoCA 开源地址：https://github.com/codefuse-ai/Collinear-Constrained-Attention