Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Josh Alman, Zhao Song

TL;DR

The paper tackles accelerating RoPE-based attention in transformers under a bounded-entry regime by introducing a novel algorithm that combines the polynomial method with FFT to achieve almost linear time, specifically $n^{1+o(1)}$, for $d=O(rac{rac{ rac{rac{n}}}})$ and $B=o(rac{rac{rac{ rac{rac{log n}}{}}})$, with additive error $rac{1}{rac{poly}(n)}$. The core idea is to decompose the RoPE attention into a sum of rescaled Toeplitz matrices, each amenable to fast multiplication via FFT, and to approximate the exponential with a low-degree polynomial as in prior work. A matching $ extsf{SETH}$-based lower bound demonstrates the necessity of the $B = o(rac{rac{ rac{rac{log n}}{}}})$ condition, confirming the tightness of the approach. The work thus provides a principled route to speed RoPE-enabled attention beyond conventional linear-attention methods, potentially enabling faster training and inference for RoPE-based LLMs with bounded inputs.

Abstract

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.

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

TL;DR

The paper tackles accelerating RoPE-based attention in transformers under a bounded-entry regime by introducing a novel algorithm that combines the polynomial method with FFT to achieve almost linear time, specifically , for and , with additive error . The core idea is to decompose the RoPE attention into a sum of rescaled Toeplitz matrices, each amenable to fast multiplication via FFT, and to approximate the exponential with a low-degree polynomial as in prior work. A matching -based lower bound demonstrates the necessity of the condition, confirming the tightness of the approach. The work thus provides a principled route to speed RoPE-enabled attention beyond conventional linear-attention methods, potentially enabling faster training and inference for RoPE-based LLMs with bounded inputs.

Abstract

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.
Paper Structure (21 sections, 9 theorems, 33 equations)

This paper contains 21 sections, 9 theorems, 33 equations.

Table of Contents

  1. Introduction
  2. Related Work
  3. Preliminaries
  4. Notation
  5. Polynomial Approximation of Exponential
  6. Fast Circulant Transform
  7. Toeplitz Matrix
  8. Rescaled Toeplitz Matrix
  9. How to Compute the Linear Attention under RoPE
  10. Conclusion
  11. Roadmap.
  12. Preliminaries
  13. Background on Hardness and Complexity
  14. Low bound Existence
  15. Hardness
  16. ...and 6 more sections

Key Result

Theorem 1.3

Suppose $\epsilon = 1/\mathop{\mathrm{poly}}\nolimits(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.

Theorems & Definitions (37)

  • Definition 1.1: A General Approximate RoPE Attention Computation, $\mathsf{ARAttC}$
  • Remark 1.2
  • Theorem 1.3: main result, upper bound
  • Theorem 1.4: main result, lower bound
  • Lemma 3.1: aa22
  • Definition 3.2: Circulant matrix
  • Definition 3.4: Toeplitz matrix
  • Remark 3.6
  • Definition 3.7: Rescaled Toeplitz Matrix
  • proof
  • ...and 27 more