提出快速旋转位置编码注意力算法,实现近线性时间计算。
Fast RoPE Attention: Combining the Polynomial Method and Fast Fourier Transform
- 结合多项式方法与快速傅里叶变换,设计新算法
- 在有界输入条件下实现近线性时间复杂度
- 适用于需高效处理序列位置信息的Transformer模型
Transformer架构广泛应用于各类机器学习任务,其主要性能瓶颈在于注意力计算。Alman和Song(NeurIPS 2023)证明,在有界输入条件下,存在接近线性时间的注意力近似算法,且该假设在强指数时间假设下为必要条件。近期,采用位置编码的Transformer版本表现优异,其中最有效的是旋转位置编码(RoPE),能捕捉标记间的依赖关系并考虑其序列位置。然而,RoPE使注意力计算复杂化,导致已有快速算法失效。本文提出一种新算法,在有界输入条件下,通过创新结合多项式方法与快速傅里叶变换,实现RoPE注意力的近线性时间计算。已知下界表明,有界输入条件仍为必要。
原文摘要 · Abstract (English)
The transformer architecture has been widely applied to many machine learning tasks. A main bottleneck in the time to perform transformer computations is a task called attention computation. [Alman and Song, NeurIPS 2023] have shown that in the bounded entry regime, there is an almost linear time algorithm to approximate the attention computation. They also proved that the bounded entry assumption is necessary for a fast algorithm assuming the popular Strong Exponential Time Hypothesis. A new version of transformer which uses position embeddings has recently been very successful. At a high level, position embedding enables the model to capture the correlations between tokens while taking into account their position in the sequence. Perhaps the most popular and effective version is Rotary Position Embedding (RoPE), which was proposed by [Su, Lu, Pan, Murtadha, Wen, and Liu, Neurocomputing 2024]. A main downside of RoPE is that it complicates the attention computation problem, so that previous techniques for designing almost linear time algorithms no longer seem to work. In this paper, we show how to overcome this issue, and give a new algorithm to compute the RoPE attention in almost linear time in the bounded entry regime. (Again, known lower bounds imply that bounded entries are necessary.) Our new algorithm combines two techniques in a novel way: the polynomial method, which was used in prior fast attention algorithms, and the Fast Fourier Transform.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。