Fast RoPE Attention: Combining the Polynomial Method and Fast Fourier Transform

Introduction

Large language models (LLMs) are among the most impactful tools in modern machine learning. LLMs such as Transformer , BERT , GPT-3 , PaLM , OPT , GPT-4 , Gemini , Gemini 1.5 , Claude3 , GPT-4o , o1 , can process natural language more effectively than smaller models or traditional algorithms. This means that they can understand and generate more complex and nuanced language, which can be useful for a variety of tasks such as language translation, question answering, and sentiment analysis. LLMs can also be adapted to multiple purposes without needing to be retained from scratch.

The straightforward algorithm for this problem runs in roughly quadratic time. Moreover, there are known complexity-theoretic lower bounds proving that the problem cannot be solved in truly subquadratic time in the case when the input matrices $Q,K,V$ have large entries, assuming a popular conjecture from fine-grained complexity theory called the Strong Exponential Time Hypothesis ($\mathsf{SETH}$ ) which we discuss more shortly.

Bounded entries in practice. The theoretical results of offer an explanation for a phenomenon commonly observed in practice: attention computation becomes significantly more efficient when the input matrices have smaller entries. Previous work on LLM implementations has noted similar observations; algorithmic techniques like quantization and low-degree polynomial approximation , which result in bounded or low-precision entries, can dramatically accelerate LLM operations. See, for example, the discussions in and .

RoPE: Rotary Position Embedding. This work mainly explores the efficient computation of an emerging type of attention, namely RoPE attention, which enables improved attention expressiveness while resulting in a more difficult computational problem. This property makes the efficient computation of RoPE attention more challenging, since a wide range of previous works (e.g., the algorithm in ) cannot be applied in this new setting. Various industrial LLMs have adopted RoPE attention as their model components, making RoPE a standard approach in attention computation. Examples include Meta’s open-source Llama family models , Anthropic’s private commercial model Claude 3 , and Apple’s LLM architecture The use of RoPE attention can be found in the technical reports of these LLMs. See page 3 of , page 5 of , page 7 of , and page 3 of ..

The inherent intuition of RoPE is to enhance the attention expressivity via rotating the queries and keys. Specifically, the rotation depends on the sequence positions, thereby ensuring that the inner product of vectors with position encoding can express the actual relative positions. Instantiation of RoPE attention is based on the $R_{j-i}$ matrices, which we will define below. These matrices perform position-aware rotations to the embeddings, which makes token pairs with smaller relative distances have larger correlations.

We now briefly describe the mathematical definition of the RoPE method. We will make use of $2\times 2$ rotation matrices, which for an angle of rotation $\theta$, can be written as

As above, we denote the length of input sequences by $n$, and represent the dimension of embedding vectors by $d$. We assume here that $d$ is even.

The angle frequencies are given by $\theta_{k}=\alpha^{-2(k-1)/d}$ for $k\in[d/2]$. Here one thinks of the angle $\alpha$ as a fixed constant for all $i$ and $j$; in the original RoPE it is $10^{4}$ (see details in Equation (15) in page 5 of ).

Formulation of RoPE Attention. In this paper, we give a new algorithm for RoPE attention. We now formally define the problem we will solve. Notably, our algorithm actually solves the following generalization of RoPE attention, which captures RoPE (as we described it above) as well as many natural variants on RoPE that future work may want to consider. We emphasize that changing the many parameters which go into the RoPE definition would still be captured by our generalization below.

Our main result is a new algorithm which computes General Approximate RoPE Attention Computation in almost linear time:

Suppose $\epsilon=1/\operatorname{poly}(n)$, $B=o(\sqrt{\log n})$, and $d=O(\log n)$. There is an $n^{1+o(1)}$ time algorithm to approximate $\mathsf{ARAttC}$ up to $\epsilon$ additive error.

In other words, although RoPE attention is more complicated than the usual attention, we are able to achieve the same running time for this more expressive version. This is, to our knowledge, the first fast algorithm for RoPE attention with provable guarantees. As we will discuss more shortly, there is a substantial barrier to using prior algorithmic techniques for attention in the setting of RoPE attention, and we overcome this barrier using a novel approach combining the polynomial method with Fast Fourier transforms.

Furthermore, we prove that the bound of $B=o(\sqrt{\log n})$ used by our algorithm is necessary, since when $B$ is any bigger, it is impossible to design a truly subquadratic time algorithm:

Assuming $\mathsf{SETH}$, for every $q>0$, there are constants $C,C_{a},C_{b}>0$ such that: there is no $O(n^{2-q})$ time algorithm for the problem $\mathsf{ARAttC}(n,d=C\log n,B=C_{b}\sqrt{\log n},\epsilon=n^{-C_{a}})$.

To emphasize, our Theorem 1.4 doesn’t just prove that our algorithmic approach cannot give a nontrivial algorithm when $B=\Omega(\sqrt{\log n})$, but more generally that it is impossible to design a nontrivial algorithm, no matter what algorithmic techniques one uses.

Our Theorem 1.4 closely matches the parameters of prior lower bounds on the usual attention problem (and it is not too difficult to prove given these prior lower bounds). Because of the increased complexity of RoPE attention, it previously seemed conceivable that one could prove a stronger lower bound for $\mathsf{ARAttC}$; perhaps surprisingly, our Theorem 1.3 shows that it is actually tight. Since the proof of Theorem 1.4 is so similar to prior work, we provide it in Section B.2 in the Appendix.

Technique Overview: Limitation of Prior Techniques

Although this approach has been successful in prior work on designing faster algorithms for many problems related to attention, it fundamentally cannot apply to RoPE attention. The key issue is that in RoPE attention, the underlying matrix which $\exp$ is applied to no longer needs to have low rank. Indeed, let $A$ denote the RoPE attention matrix (defined in Equation (1) above) and let $M$ denote $A$ before it was entry-wise exponentiated. Even in the simplest case $d=1$, one can see that by picking the $R_{j-i}$ entries appropriately (and the entries of all other matrices in Equation (1) to equal 1), one can choose $M$ to be any Toeplitz matrix (i.e., matrix whose $(i,j)$ entry depends only on the difference $j-i$). The polynomial method then cannot be used to argue that $A$ is approximately low-rank, since $M$ itself is not low-rank.

Technique Overview: Combining the Polynomial Method and Fast Fourier Transform

Although Toeplitz matrices are typically not low-rank matrices, there is a vast literature on algorithms for manipulating them using the Fast Fourier transform. (The reader may be more familiar with this fact for circulant matrices; this same algorithm can be applied by first embedding the Toeplitz matrix into a circulant matrix with twice the side-length.) Notably, it is not hard to notice that applying any function entry-wise to a Toeplitz matrix results in another Toeplitz matrix, so if $M$ were indeed a Toeplitz matrix as described in the previous paragraph, one could use the Fast Fourier transform to perform operations with the resulting matrix $A$.

However, even in the case of $d=1$, the matrix $M$ can actually be a more general type of matrix which we call a rescaled Toeplitz matrix (because of the $X$ matrices in Equation (1)). This is a matrix of the form $D_{1}CD_{2}$ for diagonal matrices $D_{1},D_{2}$ and Toeplitz matrix $C$. Unfortunately, applying a function entry-wise to a rescaled Toeplitz matrix need not result in another rescaled Toeplitz matrix.

Our main algorithmic idea is a new version of the polynomial method: we prove that if $M$ is a rescaled Toeplitz matrix, or even a sum of a small number of rescaled Toeplitz matrices, and one applies a function $f$ entry-wise to $M$ such that $f$ has a low-degree polynomial approximation, then the resulting matrix can be approximated by a sum of a relatively small number of rescaled Toeplitz matrices. In our case, we use this to write the RoPE attention matrix as a sum of rescaled Toeplitz matrices, each of which is then manipulated using the Fast Fourier transform to yield our final algorithm.

We believe our new approach, of applying polynomial approximations entry-wise to structured matrices other than low-rank matrices, may be broadly applied in other settings as well. Although the polynomial method has been applied in many algorithmic contexts, to our knowledge, it was always previously used to find a low-rank approximation of the underlying matrix, and not another structured decomposition like this.

Algorithmic techniques in practice. We emphasize that our two core techniques, the polynomial method and Fast Fourier transform, are both prevalent in practice. The polynomial method is particularly used in numerous practical algorithms for attention . For example, see detailed discussions in . Our new algorithm improves on these approaches in part by using theoretically optimal polynomials for exponentials, and combining them with the Fast Fourier transform, to give provable guarantees about their correctness and near linear running time. To our knowledge, the Fast Fourier transform has not been used in this way in prior attention algorithms.

Roadmap. In Section 2, we present our related work. In Section 3, we define certain basic notations for linear algebra. In Section 4, we commence by solving the linear case. Finally, we provide a conclusion in Section 5.

Related Work

utilize polynomial kernel approximation techniques proposed by to speed up both training and inference of a single attention layer, achieving almost linear time complexity. This method is further applied to multi-layer transformer , tensor attention , LoRA , Hopfield model , differentially private cross attention , and Diffusion Transformer , adapters , calibration approaches , multitask fine-tuning strategies , prompt tuning techniques , scratchpad approaches , instruction tuning methodologies , symbol tuning , black-box tuning , reinforcement learning from the human feedback (RLHF) , chain-of-thought reasoning and various other strategies. We will also use the polynomials of here.

The Fast Fourier transform algorithm can multiply the $n$ by $n$ Discrete Fourier transform matrix times an input vector in $O(n\log n)$ time. This algorithm is impactful in many areas, including image processing, audio processing, telecommunications, seismology, and polynomial multiplication. Due to its fundamental importance, a significant body of modern research has been dedicated to further accelerating the Fast Fourier transform. These efforts include decreasing the number of required arithmetic operations , reducing the sample complexity in the sparse setting , and improving the running time in the sparse setting . Some recent studies have explored leveraging machine learning techniques to optimize FFT performance in practical scenarios. Other works have investigated hardware-specific optimizations to further enhance computational efficiency, particularly in large-scale applications.

Other Algorithms for Computing Attention.

Due to its quadratic time complexity with respect to context length , the attention mechanism has faced criticism. To address this issue, various approaches have been employed to reduce computational overhead and improve scalability, including sparse attention , low-rank approximations , and kernel-based methods . Additionally, linear attention has emerged as a significant fast alternative to softmax attention, prompting substantial research in this area . Moreover, other related works examine various aspects of attention computation, including I/O complexity , circuit complexity , differential privacy , weights pruning , half-space reporting , graph neural network , regression problems , and quantum algorithms . A recent work has investigated the significance of selecting large weights in approximating attention computation to enhance expressiveness.

Accelerated Computation in Machine Learning.

Due to the increasing scale of training data in various applications of machine learning, including but not limited to human language , images , audio , and social networks , accelerated computation of modern ML models has been a central concern of today’s AI community . Regression models have long been a simple yet effective solution to many ML problems, such as optimization , neural network training , and signal processing . A wide range of techniques has been applied to accelerate regression computation, such as pre-conditioning and sketching . Diffusion models have recently become a fundamental game changer in content generation, producing realistic and aesthetically desirable images and videos that meet high standards. These successful stories also extend to many non-visual applications, such as text generation , drug discovery , recommender systems , and time series forecasting . A recent work has explored the intersection of diffusion models and socially aware recommender systems, aiming to mitigate the social heterophily effect through diffusion-based social information enhancement. Recent works have revealed that some specific types of diffusion modes can be approximated in almost linear time with provably efficient criteria . To accelerate the inference and training of diffusion models, enabling real-time content generation for users and fast model updates for model owners, recent progress includes shortcut models , pre-conditioning , lazy learning , and weight pruning . Graph neural networks (GNNs) are essential tools for modeling relational data , powering a wide range of applications, including traffic forecasting , fake news detection , social network analysis , human action recognition , and e-commerce . Recent advances in acceleration include model quantization , lazy learning , and sketching . A recent study accelerated GNNs using both lazy propagation and variance-reduced random sampling of finite sums, resulting in a linear-time GNN with broad applications in e-commerce.

Preliminaries

In Section 3.1, we define several notations. We discuss some backgrounds for fast circulant transform. In Section 3.2, we provide a tool from previous work about how to control error by using low-degree polynomial to approximate exponential function. In Section 3.3, we discuss some backgrounds about fast circulant transform. In Section 3.4, we formalize the toeplitz matrix and introduce the tools we will use. In Section 3.5, we define rescaled circulant matrix and provide some basic tools for it.

2 Polynomial Approximation of Exponential

To control the error dependence of our proposed approximate algorithm, we present a standard technical lemma used in many previous works .

Furthermore, $P$ can be computed efficiently: its coefficients are rational numbers with $\operatorname{poly}(g)$-bit integer numerators and denominators which can be computed in $\operatorname{poly}(g)$ time.

3 Fast Circulant Transform

Circulant matrices have been widely used in applied mathematics , compressive sensing and regression literature . Here we provide the formal definition.

Thus, we can multiply $\mathsf{Circ}(a)$ with an input vector of length $n$ in $O(n\log n)$ time using the Fast Fourier transform algorithm.

4 Toeplitz Matrix

In other words, $\mathsf{Toep}(a)_{i,j}:=a_{i-j}$.

Facts 3.3 and 3.5 imply that the matrix-vector product of a Toeplitz matrix can be computed in $O(n\log n)$ time.

5 Rescaled Toeplitz Matrix

Our algorithm will critically involve manipulating a certain kind of structured matrix we call a rescaled Toeplitz matrix. In this section we define these matrices and prove basic properties which we will use.

Suppose $M=D_{1}CD_{2}$, we first compute $D_{2}v$ straightforwardly in $O(n)$ time. Then we compute $C\cdot(D_{2}v)$ in $O(n\log n)$ time. Finally, we compute $D_{1}\cdot(CD_{2}v)$ in $O(n)$ time. ∎

If $A$ and $B$ are rescaled Toeplitz matrices, then $A\circ B$ is also a rescaled Toeplitz matrix.

Suppose $A=\operatorname{diag}(a_{1})A_{2}\operatorname{diag}(a_{3})$ where $A_{2}$ is a Toeplitz matrix, and $B=\operatorname{diag}(b_{1})B_{2}\operatorname{diag}(b_{3})$ where $B_{2}$ is a Toeplitz matrix. We can show

Therefore, we know $A\circ B$ is also a rescaled Toeplitz matrix. ∎

If $A_{1},\cdots,A_{t}$ are rescaled Toeplitz matrices, then for any vector $v$, we have $(A_{1}\circ A_{2}\circ\cdots\circ A_{t})v$ can be computed in $O(tn\log n)$ time.

The proof directly follows from applying Lemma 3.9 and Fact 3.8, $t$ times. ∎

How to Compute the Linear Attention under RoPE

Before starting to work on RoPE softmax attention, here we consider the simpler problem of computing RoPE linear attention. This linear attention does not have entry-wise $\exp$.

We define $D:=\operatorname{diag}(A{\bf 1}_{n})$. The attention computation is going to output an $n\times d$ matrix

Given a collection of weight matrices $W_{-(n-1)},\cdots W_{-1},W_{0},W_{1},\cdots W_{n-1}$, we use $S$ to denote their support such that $\forall i\in\{-(n-1),\cdots,n-1\}$, $\operatorname{supp}(W_{i})=S$.

where the second step follows from Definition 4.3. ∎

Conclusion

In this work, we provide an almost linear time algorithm for RoPE attention. RoPE attention is used as a more expressive variant on attention in many applications, but the usual polynomial method approach inherently cannot work for calculating it quickly. We introduced a new way to combine the polynomial method with our “rescaled Toeplitz matrices” and the Fast Fourier transform in order to solve this problem more efficiently. As future work introduces more variants on attention, it will be exciting to explore whether these and other linear algebraic tools can still be used to perform fast computations.

In Section A, we introduce some theoretical foundations and other basic preliminaries. In Section B, we introduce some background and present our hardness result. In Section C, we explain how to handle the exp units and present the proofs of our main result. In Section D, we discuss the limitations of the paper. In Section E, we present the impact statement.

Appendix A Preliminaries

In this Section, we introduce concepts about nearly-linear time and almost-linear time.We include this discussion into the paper due to the request from ICLR 2025 reviewer https://openreview.net/forum?id=AozPzKE0oc.

We say $O(n\operatorname{poly}(\log n))$ is nearly-linear time, and $O(n^{1+o(1)})$ is almost-linear time.

Then we introduce the relationship between $O(n\log n)$ and $O(n^{1+o(1)})$.

We can show that $O(n\operatorname{poly}(\log n))\leq O(n^{1+o(1)})$.

First, observe that $\operatorname{poly}(\log n)=n^{O(\log(\log n))/\log n}$. Since $\log(\log n)/\log n\to 0$ as $n\to\infty$, we have that $\log(\log n)/\log n=o(1)$.

This directly shows that $O(n\operatorname{poly}(\log n))\leq O(n^{1+o(1)})$. ∎

Appendix B Background on Hardness and Complexity

In Section B.1, we introduce some background and the low bound existence. In Section B.2, we present our hardness result.

In computational complexity theory, algorithmic hardness refers to the inherent difficulty of solving computational problems, measured by the resources (such as time and space) required for their resolution. As established in Garey and Johnson’s foundational work ”Computers and Intractability” , understanding this hardness helps researchers and practitioners determine whether efficient solutions exist for given problems. Particularly significant in this context, lower bounds serve as a critical theoretical tool for establishing the minimum resources required to solve specific computational problems. This naturally leads us to examine how lower bounds play a fundamental role in computational complexity theory, establishing fundamental limits on the resources required to solve computational problems. As discussed in ”Introduction to the Theory of Computation” (Chapter 9) and ”Computational Complexity: A Modern Approach” (Chapter 3), proving lower bounds helps us understand the inherent difficulty of problems and provides insights into computational hierarchies.

$\mathcal{A}$ is the set of all possible algorithms

$\text{Resources}(A)$ denotes the resource usage of algorithm $A$

$\text{Succeeds}(A,P)$ indicates that algorithm $A$ correctly solves problem $P$

$\text{LB}_{\mathcal{C}}$ represents the lower bound for class $\mathcal{C}$

$f(n)$ is a function of the input size $n$

For proving computational complexity lower bounds, we can establish the following: Let $\mathcal{C}$ be a class of computational problems. To prove that all algorithms solving problems in $\mathcal{C}$ require at least $f(n)$ resources (time or space), it is sufficient to demonstrate that there exists a single problem instance $P\in\mathcal{C}$ for which no algorithm using less than $f(n)$ resources can correctly solve $P$.

B.2 Hardness

In this section, we show The Strong Exponential Time Hypothesis. Over 20 years ago, Impagliazzo and Paturi introduced The Strong Exponential Time Hypothesis (SETH). It is a stronger version of the $\mathsf{P}\neq\mathsf{NP}$ conjecture, which asserts that our current best $\mathsf{SAT}$ algorithms are roughly optimal:

For every $\epsilon>0$ there is a positive integer $k\geq 3$ such that $k$-$\mathsf{SAT}$ on formulas with $n$ variables cannot be solved in $O(2^{(1-\epsilon)n})$ time, even by a randomized algorithm.

SETH is a popular conjecture which has been used to prove fine-grained lower bounds for a wide variety algorithmic problems, as discussed in depth in the survey .

Assuming $\mathsf{SETH}$, for every $q>0$, there are constants $C,C_{a},C_{b}>0$ such that: there is no $O(n^{2-q})$ time algorithm for the problem $\mathsf{ARAttC}(n,d=C\log n,B=C_{b}\sqrt{\log n},\epsilon=n^{-C_{a}})$.

Appendix C How to Handle the Exp Terms

We now give our full algorithm for general RoPE attention. In Section C.1, we study matrices which are the entry-wise products of a number of rescaled Toeplitz matrix, and how to use that decomposition to quickly multiply such matrices with a vector. In Section C.2, we show how to decompose the RoPE attention matrix into summation of a number of such structured matrices using the polynomial method. In Section C.3, we show how to put everything together to get our main result.

Using Lemma 3.9, we know the entry-wise product between any two rescaled Toeplitz matrix is still a rescaled Toeplitz matrix. Thus, applying Lemma 3.9 to the above equations for $\sum_{i=1}^{|S|}t_{i}$ times, we can show that $A^{(m)}$ is still a rescaled Toeplitz matrix.

Using Lemma 3.10, we know that for any vector $v$, $A^{(m)}v$ can be computed in $O((\sum_{i=1}^{|S|}t_{i})\cdot n\log n)$ time. ∎

C.2 Expanding polynomials into summation of several rescaled Toeplitz matrices

where the first step follows from definition of $M$, the second step follows from Eq. (2), and the last step follows from definition of $N^{(m)}$. Thus, we can see $M=\sum_{m\in{\cal M}}\alpha_{m}\cdot N^{(m)}.$

C.3 Main Result

Finally, we are ready to put all our techniques together.

Suppose $d=O(\log n)$ and $B=o(\sqrt{\log n})$. There is an $n^{1+o(1)}$ time algorithm to approximate $\mathsf{ARAttC}$ up to $\epsilon=1/\operatorname{poly}(n)$ additive error.

We use the polynomial of Lemma 3.1 in Lemma C.2 with choice of $k=|S|=O(d)=O(\log n)$ and $\widetilde{d}=o(\log n)$ is the degree of the polynomial from Lemma 3.1 for error $1/\operatorname{poly}(n)$. We can thus upper bound

The total running time consists of three parts: first, approximating $A{\bf 1}_{n}$ which gives an approximation to diagonal matrix $D$; second, approximating $Av$ for $d$ different columns vectors $v$, this will approximate $AV$; third, combining approximation of $D^{-1}$ with approximation of $AV$, to obtain an approximation of $D^{-1}AV$. Combining Lemma C.1 and C.2. The dominating running time for above three parts is

Due to the choice of $|{\cal M}|=n^{o(1)}$, $|S|=O(d)$, $d=O(\log n)$.

The error analysis remains identical to prior attention algorithms using the polynomial method , thus we omit the details here. ∎

Appendix D Limitations

This work presents an almost linear-time algorithm for RoPE attention, supported by theoretical analysis. However, we do not include any empirical evaluations to validate the practical performance of the proposed method.

Appendix E Impact Statement

This work introduces the first almost linear-time algorithm for RoPE attention, providing a novel solution and new insights into the computational bottlenecks of RoPE-based attention mechanisms. It has the potential to accelerate future large language model (LLM) training and evaluation. As this is a purely theoretical contribution, we do not foresee any negative social impacts.

References