用余弦相似度替代softmax,让Transformer内存线性增长。
Cottention: Linear Transformers With Cosine Attention
- 用余弦相似度替换softmax,实现线性内存复杂度。
- 在BERT和GPT任务上性能接近softmax注意力,内存大幅降低。
- 可视为带有限状态的RNN,推理时内存恒定,适合长序列处理。
注意力机制,尤其是softmax注意力,在基于Transformer的模型(如GPT)中发挥了关键作用。然而,softmax注意力随序列长度呈二次增长的内存复杂度,给长序列处理带来显著挑战。我们提出Cottention,一种将softmax操作替换为余弦相似度的新注意力机制。通过利用余弦相似度的特性并重排注意力公式,Cottention实现了与序列长度成线性关系的原生内存复杂度,天然比softmax注意力更高效。我们证明Cottention可重构成具有有限隐藏状态的循环神经网络(RNN),从而在推理阶段实现恒定内存使用。我们在双向BERT和因果GPT任务上评估了Cottention,结果表明其性能与softmax注意力相当,同时显著降低了内存需求。为确保高效计算,我们开发了专用CUDA内核。实验显示,由于原生线性内存复杂度及推理时恒定内存占用,Cottention是softmax注意力的有力替代方案,能够无需牺牲性能即可处理更长序列。
原文摘要 · Abstract (English)
Attention mechanisms, particularly softmax attention, have been instrumental in the success of transformer-based models such as GPT. However, the quadratic memory complexity of softmax attention with respect to sequence length poses significant challenges for processing longer sequences. We introduce Cottention, a novel attention mechanism that replaces the softmax operation with cosine similarity. By leveraging the properties of cosine similarity and rearranging the attention equation, Cottention achieves native linear memory complexity with respect to sequence length, making it inherently more memory-efficient than softmax attention. We demonstrate that Cottention can be reformulated as a recurrent neural network (RNN) with a finite hidden state, allowing for constant memory usage during inference. We evaluate Cottention on both the bidirectional BERT and causal GPT tasks, demonstrating comparable performance to softmax attention while significantly reducing memory requirements. To ensure efficient computation, we develop a custom CUDA kernel for Cottention. Our results show that Cottention is a promising alternative to softmax attention, enabling the processing of longer sequences without sacrificing performance, due to its native linear memory complexity and ability to maintain a constant memory footprint during inference.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。