CoCA：在 Transformer 中融合位置嵌入与共线约束注意力以扩展长上下文窗口
机器学习
2024-02-29 v3 人工智能
计算与语言
摘要
自注意力与位置嵌入是基于 Transformer 的大语言模型（LLMs）的两个关键模块。然而，二者之间的潜在关系远未被充分研究，尤其在长上下文窗口扩展方面。事实上，我们的工作揭示了旋转位置嵌入（RoPE）与原始自注意力之间存在损害长上下文外推的异常行为。为解决该问题，我们提出一种新颖注意力机制 CoCA（共线约束注意力，Collinear Constrained Attention）。具体而言，我们对 与 施加共线约束，以无缝集成 RoPE 与自注意力。该集成仅增加极小计算与空间复杂度，却显著增强了长上下文窗口外推能力。我们提供了优化实现，使其成为任何现有基于 Transformer 模型的即插即用替换。大量实验表明 CoCA 在扩展上下文窗口方面表现极为优异。一个基于 CoCA 的 GPT 模型以 512 的上下文长度训练，可在无任何微调下无缝扩展上下文窗口至 32K（60）。此外，将 CoCA 植入 LLaMA-7B，我们在仅 2K 训练长度内实现了至 32K 的外推。我们的代码公开于：https://github.com/codefuse-ai/Collinear-Constrained-Attention
引用
@article{arxiv.2309.08646,
title = {CoCA: Fusing Position Embedding with Collinear Constrained Attention in Transformers for Long Context Window Extending},
author = {Shiyi Zhu and Jing Ye and Wei Jiang and Siqiao Xue and Qi Zhang and Yifan Wu and Jianguo Li},
journal= {arXiv preprint arXiv:2309.08646},
year = {2024}
}
备注
16 pages, 7 figures