用元素级注意力替代传统自注意力,实现高效长序列训练与推理。
Element-wise Attention Is All You Need
- 改用元素平方距离计算相似性,避免点积运算
- 训练复杂度降为O(tLD),推理可转为RNN,仅O(tD)
- 保留自注意力性能,无信息压缩或尖锐性损失
自注意力机制在多个领域表现优异,但训练与推理复杂度高。现有低复杂度方案如线性注意力、线性RNN和状态空间模型虽降低开销,却存在信息压缩与尖锐性下降等性能退化问题。本文提出一种新型元素级注意力机制,以元素平方欧氏距离代替点积计算相似性,并用泰勒多项式近似指数项$\\(exp(q_{ic}k_{jc})$$。该方法训练复杂度为$\\mathcal{O}(tLD)$,推理可重构为RNN,复杂度仅为$\\mathcal{O}(tD)$,其中$L$为序列长度,$D$为特征维度,$t$为多项式最高阶数。该设计有效规避了现有方法的性能缺陷,在因果与非因果形式下均达到与自注意力相当的性能。
原文摘要 · Abstract (English)
The self-attention (SA) mechanism has demonstrated superior performance across various domains, yet it suffers from substantial complexity during both training and inference. The next-generation architecture, aiming at retaining the competitive performance of SA while achieving low-cost inference and efficient long-sequence training, primarily focuses on three approaches: linear attention, linear RNNs, and state space models. Although these approaches achieve reduced complexity than SA, they all have built-in performance degradation factors, such as diminished “spikiness†and compression of historical information. In contrast to these approaches, we propose a novel element-wise attention mechanism, which uses the element-wise squared Euclidean distance, instead of the dot product operation, to compute similarity and approximates the quadratic complexity term $\exp(q_{ic}k_{jc})$ with a Taylor polynomial. This design achieves remarkable efficiency: during training, the element-wise attention has a complexity of $\mathcal{O}(tLD)$, making long-sequence training both computationally and memory efficient, where $L$ is the sequence length, $D$ is the feature dimension, and $t$ is the highest order of the polynomial; during inference, it can be reformulated as recurrent neural networks, achieving a inference complexity of $\mathcal{O}(tD)$. Furthermore, the element-wise attention circumvents the performance degradation factors present in these approaches and achieves performance comparable to SA in both causal and non-causal forms.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。