[Apple 2024] Exclusive Self Attention: 消除“自恋”偏见，让 Transformer 真正读懂上下文

TL;DR

Apple 的研究员 Shuangfei Zhai 提出了一种名为 Exclusive Self Attention (XSA) 的新型注意力机制。其核心逻辑非常简单：在计算注意力输出时，强制剔除掉属于“自己”的信息。实验证明，这种简单的正交化处理能有效解决 Attention Similarity Bias 问题，在 2.7B 规模的模型上实现了全方位的性能提升，且序列越长，优势越明显。

痛点深挖：注意力机制的“自恋”行为

在标准的 Transformer 架构中，Self Attention (SA) 负责聚合上下文信息，而 Feed Forward Network (FFN) 负责逐位置的特征变换。理想情况下，这两者应该各司其职。

然而，作者通过实验发现了一个令人惊讶的现象——Attention Similarity Bias。通过观察图1，我们可以发现：

1. 每一层注意力生成的输出向量 $y_i$ 与该位置原始的 Value 向量 $v_i$ 具有极高的余弦相似度。
2. 随着层数加深，这种相似度不断攀升。

这意味着注意力机制花费了大量的计算能力在“搬运”自己本就拥有的信息，而不是从周围吸取新知识。这不仅造成了计算浪费，还与后续的 FFN 层功能重叠，限制了模型对复杂上下文的建模能力。

图1：左：Value向量相似度；中：自位置注意力分值；右：输出与自Value的相似度。可以明显看到右图中相似度随层数剧增。

Methodology：两行代码实现正交化

为了纠正这种偏差，XSA 引入了一个简单的几何约束：从聚合结果中减去在自向量方向上的投影。

数学表达式如下： $z_i=y_i-(y_i^T{v}_i)\frac{v_i}{\Vert v_i{\Vert }_2^2}$

简单来说，这就是一个典型的向量去相关操作。通过减去投影分量，得到的 $z_i$ 与 $v_i$ 严格正交。这意味着注意力层的输出现在完全不包含当前位置原始的 Value 信息，强迫模型后续的 FFN 层必须依赖 Residual Path（残差路径）来获取自位置特征，而 SA 层则专注于整合“他山之石”。

在工程实现上，这极其优雅——仅需在标准的 Multi-Head Attention 代码中增加两行 Tensor 运算（详见下方的 PyTorch 风格伪代码）。

算法 1：XSA 的简洁实现。

实验与结果：小改动大收益

1. 几乎可以忽略的计算开销

尽管增加了一步正交化计算，但在 B200 GPU 上的实测显示，无论是内存占用还是时间延迟，XSA 与标准 SA 相比几乎没有差别（见图 2）。

图2：XSA 在不同序列长度和模型尺寸下均保持极低的计算开销。

2. 模型规模与长文本优势

在 100B Token 的训练测试中，XSA 在 0.7B、1.4B、2.7B 三种参数量级下均取得了更低的 Validation Loss。

最令人兴奋的发现在于序列长度的扩展性（图 5）。当序列长度从 512 增加到 16384 时，XSA 相对于 Baseline 的领先优势不断扩大。这说明在长文本场景下，上下文建模的矛盾更加突出，而 XSA 的强制分工策略精准地解决了这一痛点。

图5：序列越长，XSA 的性能增益越显著。

深度洞察与总结

Exclusive Self Attention 的成功给我们带来了两点重要的启发：

1. 架构设计的“减法”逻辑：有时候，限制模型的某些行为（如限制它关注自己）反而能迫使它学习更深层次的模式。这与对比学习中防止“模型坍缩”的逻辑有异曲同工之妙。
2. 隐式 Attention Sink：作者提到，XSA 实际上起到了“隐式注意力汇”的作用。它不需要显式添加特定的 Sink Token，而是通过语法结构上的约束，自然地将无效的注意力权重分配给自位置，从而维持了注意力分布的稳定性。

总结 (Takeaway)：XSA 是一个性价比极高的改进。它通过物理直觉明确了 Transformer 组件的分工，不仅提升了预训练效率，更为大模型的长文本扩展提供了一种低成本且稳健的方案。

展望：目前该研究主要集中在 2.7B 以内的语言模型。在更大的 Scale（如 70B+）下，这种正交约束是否会带来训练稳定性的挑战，或者是进一步放大增益，值得业界进一步探索。