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

Thank you for the thorough review. We will implement the suggested changes in the final version.

To answer your questions:

• A number of papers have been published about empirically evaluating the polynomial method for attention, including by researchers at companies which have very large-scale Transformer models. (One example is PolySketchFormer.)
• The final running time given here is performing a single RoPE attention computation, so if one were performing $L$ different computations (either in parallel, or in series, or otherwise over the course of a larger Transformer computation), the running time should scale linearly with $L$, since it would just perform this algorithm $L$ times. (Sorry in advance if we're misunderstanding the question.)

Thank you for the very thorough review. We will implement all the suggested changes, and wish to comment on just a few:

normalizations You're right that the normalizations are important, and our results are indeed for the correct normalization by \sqrt{d} (we will correct this typo). We ultimately use the same normalization/error analysis as Alman & Song (2023), which uses the proper normalization.

Section 4 You're right that the approach you give would also work for linear attention. However, this asymmetric approach (i.e., which treats i and j differently) ends up generalizing poorly in section 5 once one incorporates the entry-wise exponentiation. In short, when one gets down to line 497, it would replace d with d^2 in the binomial coefficient, and thus result in a worse rank bound which only applies for smaller B.

empirical results We agree that an empirical evaluation of our algorithm would be exciting future work. That said, we believe the theoretical results we provide here are more interesting than an empirical evaluation for two reasons.

First, since RoPE attention appears more complicated than regular attention, we previously believed it would be possible to prove a greater lower bound for RoPE attention. Our algorithm shows this intuition is wrong, and an improved lower bound is not possible, since the same approximation can be achieved quickly for RoPE attention. This main message, that RoPE attention is no harder than attention in theory, was very surprising to us.

Second, many different works have explored different ways to make use of positional information in attention and attention variants. One could empirically evaluate a particular version of RoPE attention, and a new variant would likely be introduced the very next day. However, our new principled way to combine the polynomial method and FFT should work for any such variant, including new ones that researchers develop in the future, and can hopefully guide future approaches to incorporating positional information.

Thank you for your thorough review.

algorithmic novelty: You're right that the polynomial method and FFT have both been used previously for attention. We also discuss this on page 5. That said, we believe we are the first to combine these two approaches in a single algorithm. See the section on page 4 called "Technique Overview: Limitation of Prior Techniques" where we discuss the barriers to using prior approaches for RoPE attention, and the subsequent section on page 4 where we describe the novel way we combine these techniques.

lower bound You're right that we're proving a lower bound on a special case of the general RoPE attention here (albeit, the case which appears the easiest). In fact, one can modify the lower bound proof for attention of Alman & Song, 2023 to work for the usual RoPE attention as well. The modification is straightforward but tedious, so we previously didn't think it was worth including. That said, you are absolutely correct that it would strengthen the lower bound statement, so we will include it in the next version.

implementations We agree that an empirical evaluation of our algorithm would be exciting future work. That said, we believe the theoretical results we provide here are more interesting than an empirical evaluation for two reasons.

First, since RoPE attention appears more complicated than regular attention, we previously believed it would be possible to prove a greater lower bound for RoPE attention. Our algorithm shows this intuition is wrong, and an improved lower bound is not possible, since the same approximation can be achieved quickly for RoPE attention. This main message, that RoPE attention is no harder than attention in theory, was very surprising to us.

Second, many different works have explored different ways to make use of positional information in attention and attention variants. One could empirically evaluate a particular version of RoPE attention, and a new variant would likely be introduced the very next day. However, our new principled way to combine the polynomial method and FFT should work for any such variant, including new ones that researchers develop in the future, and can hopefully guide future approaches to incorporating positional information.

questions

• Yes you're right, thanks for noting the typo.
• See the paragraph called "Bounded entries in practice" on page 2. In short, a number of prior works have used techniques to bound the entries or observed that entries are already bounded in order to speed up algorithms.
• The papers discussed in the paragraph on "Bounded entries in practice" on page 2 give algorithmic techniques for imposing that the bound on $B$ must hold. In those papers, they find that imposing this bound only slightly hurts the quality of the outputs of the language models.

Thank you for the review. The simplification you suggest to claim 4.8 does not work, and touches on the core reason why RoPE attention is more difficult to apply the polynomial method to than regular attention.

If it were true that $A^{\ell_1,\ell_2}$ were an outer product of vectors as you suggest, then it would be a rank 1 matrix, and A, the entire attention matrix before exponentiation, would have low rank at most $|S| \leq O(\log n)$. In fact, the main difficulty with RoPE attention, and the reason this type of approach cannot work, is that $A$ can have full rank even when $d=1$. See the section on page 4 called "Technique Overview: Limitation of Prior Techniques" where we discuss this in more detail.

You're right that $\ell_1, \ell_2$ are the indices into $W_{i-j}$, which is a $d \times d$ matrix. However, in definition 4.7 we are defining an $n \times n$ matrix $C^{\ell_1, \ell_2}$ whose $(i,j)$ entry is equal to the $(\ell_1, \ell_2)$ entry of $W_{i-j}$. Thus, this is a circulant matrix (since its $(i,j)$ entry depends only on $i-j$) but it is an $n \times n$ matrix, and not a scalar as you suggest. On line 412 we relate the $(i,j)$ entry of $A^{\ell_1, \ell_2}$ to the $(i,j)$ entry of $C^{\ell_1, \ell_2}$, from which the matrix relation we give on line 415 follows. Since $C^{\ell_1, \ell_2}$ is an $n \times n$ circulant matrix, this does not show that $A^{\ell_1, \ell_2}$ has rank 1 as you are suggesting.

Finally, to respond to the specific bullet points you raise in the weaknesses section:

• $W_{i-j}$ is the matrix given as input to the problem in Definition 1.1; this refers to the $(\ell_1, \ell_2)$ entry of that matrix.
• As above, such a simplification is not possible.
• Thanks for noticing this, we will define $B$ earlier.
• In Claim 4.8 we wanted to remind the reader the definition of $C^{\ell_1, \ell_2}$ to emphasize that we are about to use it.
• We do use this notation, for instance on line 450 on page 9.

Thank you for the review. We strongly believe that all the claims and assumptions in our paper are correct.

1. The Theorem statements are correct as written. The relationship between $n \log n$ and $n^{1 + o(1)}$ is the opposite of what you claim. Since, for every $\varepsilon>0$, we have $n \log n < n^{1 + \varepsilon}$ for big enough $n$, it follows that $n \log n \leq n^{1 + o(1)}$. See, for instance, https://cs.stackexchange.com/a/152315
2. We're not sure what claim you're saying is invalidated here. This is actually the main point of our paper: that although RoPE attention appears more complicated than regular attention, we are actually able to give a new algorithm to approximately compute it in the same running time.
3. We believe you're referring to Theorems 1.4 and 6.2, rather than Theorems 1.5 and 5.2? At any rate, these theorems are also correct as written. We observe that the ARAttC problem is hard even in the special case when all the W matrices are the identity matrix, and thus it is not possible to design a faster algorithm for ARAttC.
4. Our claim is correct. Entry-wise exponentiating a matrix does not maintain its rank as you suggest, and this is why the usual attention cannot be computed exactly in almost linear time. As a simple example, the following $2 \times 2$ matrix has rank 1, but when you entry-wise exponentiate it, it has rank 2: $\begin{bmatrix} 1 & 2 \\\\ 2 & 4 \end{bmatrix}$ https://www.wolframalpha.com/input?i=rank+of+%7B%7Be%5E1%2C+e%5E2%7D%2C+%7Be%5E2%2C+e%5E4%7D%7D
5. We're not sure what algorithm you're referring to here, but the straightforward algorithm takes $\Theta(n^2 d)$ time since it constructs the $n \times n$ attention matrix, which cannot be done in subquadratic time in $n$.
6. You observe that $C^{1,1}$ has the property that "Its elements depend only on difference of positions (i-j)". This is exactly the definition of a circulant matrix! Thus, your subsequent claim that "it's a symmetric matrix with an additional structure which differs from circulant matrix structure" is incorrect, and this is not in any way a counterexample.
7. When $d=1$, each $R_{j-i}$ is a scalar value, and $M$ is the matrix whose $(i,j)$ entry is the value $R_{j-i}$. You can thus pick the scalars so that $M$ is any circulant matrix you would like.

1. Unfortunately, I cannot agree that computational complexity $n \log n \leq n^{1+o(1)}$. The formal definition of o(1) function g(n) is $\lim_{n \to \infty} \frac{g(n)}{1} = \lim_{n \to \infty} g(n) = 0$. The same is stated in the answers of the StackExchange questions you provided (https://cs.stackexchange.com/a/152315). Please also note that, while $n \log n \leq n^{1+\epsilon}$ is a true statement, any constant positive function $g(n)=\epsilon$ is not in o(1) class. So, $\lim_{n \to \infty} n^{1+o(1)} = \lim_{n \to \infty} n^{1+0} = \lim_{n \to \infty} n \le \lim_{n \to \infty} n \log n$.

2. The issue from point 1 invalidates of the main claims of the paper (lines 155-159) that the algorithm is in the same runtime complexity class ($n^{1+o(1)}$ to be specific) as in paper [1], which you build upon, and [2]. The paper [1] establishes a truly $n^{1+o(1)}$ algorithm. If you compare your algorithm with plain exact attention calculation , which has $O(n^2)$ runtime complexity, and not with the approximate algorithm from [1], please state it explicitly. As of now, the phrasing "we are able to achieve the same running time for this more expressive version" in conjunction with claimed $n^{1+o(1)}$ complexity indicates comparison with [1].

3. Thank you for the correction, I referred to Theorems 1.4 and 6.2, indeed. The special case when all the W matrices are identity matrix is the case of the regular attention, which is proved in [1]. But the general case of RoPE attention in your work when the W matrices are not identity matrix is not addressed. In order to prove the Theorem, you should show that all other cases can be reduced to this special case.

4. In the paper you talk about pre-exponentiated attention matrix (M in your notation from lines 182-188), to quote: "the underlying matrix which exp is applied to no longer needs to have low rank" and "since M itself is not low-rank."

5. I refer to the standard RoPE calculation algorithm which runs in O(nd) time (see, for example, https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py#L203). You suggest using $O(nd^2)$ algorithm for the same operation, which is impractical, at least in case of standard RoPEs.

6. Consider a matrix $C^{1, 1}$ in case of standard RoPE and sequence length n=3: $\begin{bmatrix} a & b & c \\\\ b & a & b \\\\ c & b & a \end{bmatrix}$, where $a = \cos 0\theta$, $b = \cos 1\theta$, and $c = \cos 2\theta$. This is clearly not a circulant matrix.

7. Unfortunately, this claim does not make sense for me.

I thank reviewer sE6r for the clarifications and would like to respond in details both to you and the authors.

1. I agree that $n \log(n) = n^{1 + o(1)}$. Thank you for the proof. However, I believe notation $n^{1 + o(1)}$ is confusing and unintuitive and can lead to misinterpretations. For example, take two functions: $f_1(n) = \log(\log n) / \log(n)$ and $f_1(n) = \frac{1}{n}$ which tend to zero and, thus, $o(1)$. Then, formally, both $n^{f_1(n)}$ and $n^{f_2(n)}$ are in $n^{o(1)}$ class. But $\lim_{n \to \infty} n^{f_1(n)} = \lim_{n \to \infty} \log n = \infty$ and $\lim_{n \to \infty} n^{f_2(n)} = \lim_{n \to \infty} n^{\frac{1}{n}} = 1$ which are quite different cases.

I think notation $O(n \log(n))$ is more established and universally used in the literature, especially if the algorithm is based on FFT (see e. g. [1-3]).

3. I agree. Perhaps, the authors could explicitly recite the statement of the corresponding theorem in [4] to make their argument clearer.

4. Your construction for this specific case is correct. I think the authors should explicitly include it in the relevant section of the paper. Nevertheless, I don't believe their stronger claim (point 7 in my original review) is true. More on that later.

6. Let me politely disagree. The matrix I provided is a symmetric matrix but not a circulant matrix. To see it, note that by definition of a circulant matrix, all rows in it are identical up to a circular shift by some offset. That's not the case for this matrix. More generally, the definition of $C^{l_1, l_2}$ doesn't imply it is circulant as the elements of $(W_{i-j})_{l_1, l_2} can be arbitrary, unless specifically chosen to induce a circular matrix.

7. After having read the latest author comments and the construction in 4, I'm still not convinced pre-exponentiation matrix $M$ could be made circular always by varying $W_{i-j}$ (or $R_{i-j}$, using the paper's notation from another page). It works by specific construction of matrices $Q$, $K$, $W_{i-j}$ with carefully chosen values. For every fixed set $\\{R_{i-j}\\}$, there exist such values $Q$ and $K$ which do not lead to a circular $M$. Example: for $Q=\begin{bmatrix} 1 \\\\ 2 \end{bmatrix}$ and $K=\begin{bmatrix} 3 \\\\ 4 \end{bmatrix}$, $R_0 \neq 0$ , the result will be $\begin{bmatrix} 3R_0 & 4R_1 \\\\ 6R_{-1} & 8R_0 \end{bmatrix}$ which is not a circular matrix. In practice, $Q$ and $K$ are not fixed and data-dependent while $R_{i-j}$ are constant.

References

[1] Fu et al., FlashFFTConv: Efficient Convolutions for Long Sequences with Tensor Cores. ICLR 2024

[2] Gu et al., Efficiently Modeling Long Sequences with Structured State Spaces. ICLR 2022

[3] Lee-Thorp et al., "FNet: Mixing Tokens with Fourier Transforms." NAACL 2022

[4] Josh Alman and Zhao Song. Fast attention requires bounded entries. In NeurIPS, 2023.