Cottention:基于余弦注意力的线性 Transformer

注意力机制,特别是 softmax 注意力,在 GPT 等基于 Transformer 的模型的成功中发挥了关键作用。然而,softmax 注意力相对于序列长度的二次内存复杂度在处理更长序列时带来了重大挑战。我们引入了 Cottention,一种用余弦相似度替代 softmax 操作的全新注意力机制。通过利用余弦相似度的性质并重新排列注意力方程,Cottention 实现了相对于序列长度的原生线性内存复杂度,使其在内存效率上本质上优于 softmax 注意力。我们证明 Cottention 可以被重构为具有有限隐藏状态的循环神经网络(RNN),从而在推理期间实现常数内存消耗。我们在双向 BERT 和因果 GPT 任务上评估了 Cottention,结果表明其性能与 softmax 注意力相当,同时显著降低了内存需求。为确保高效计算,我们为 Cottention 开发了定制的 CUDA 内核。我们的结果表明,得益于其原生线性内存复杂度及在推理期间保持常数内存占用,Cottention 是 softmax 注意力的一种极具前景的替代方案,能够在不牺牲性能的情况下处理更长序列。