TaylorShift：利用泰勒Softmax将自注意力复杂度从平方级降至线性（及反向）
机器学习
2025-06-06 v2 人工智能
摘要
注意力机制的二次复杂度是使用Transformer处理长序列的最大障碍之一。当前依赖稀疏表示或有状态循环的方法牺牲了token之间的交互，最终导致性能折衷。本文介绍了TaylorShift，一种泰勒Softmax的新颖重新表述，使得能够在线性时间和空间内计算完整的token间交互。我们分析确定了采用TaylorShift比传统注意力更高效的交叉点，与经验测量结果高度吻合。具体而言，我们的发现表明，TaylorShift对于短至800个token的序列即可提升内存效率，并加速约1700个token及以上输入的推理。对于更短的序列，TaylorShift的缩放与普通注意力相当。此外，在涉及长序列的五项任务上的分类基准测试显示，使用配备TaylorShift的Transformer时准确率没有下降。为便于复现，我们在https://github.com/tobna/TaylorShift提供代码。
引用
@article{arxiv.2403.02920,
title = {TaylorShift: Shifting the Complexity of Self-Attention from Squared to Linear (and Back) using Taylor-Softmax},
author = {Tobias Christian Nauen and Sebastian Palacio and Andreas Dengel},
journal= {arXiv preprint arXiv:2403.02920},
year = {2025}
}