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

2025-05-17

AI Summary

This work addresses the quadratic attention complexity $O(n^2)$ in RoPE-enhanced Transformers. We propose the first near-linear-time approximation algorithm for RoPE attention under a bounded-input assumption. Our core method synergistically integrates polynomial interpolation with the Fast Fourier Transform (FFT), overcoming the fundamental limitation that conventional fast attention techniques fail under RoPE’s rotational positional encoding structure. By constructing a low-rank approximation of RoPE embeddings in the frequency domain, our approach reduces per-layer attention computation to $O(n log n)$—achieving theoretical optimality and matching the known lower bound. Extensive experiments demonstrate that our method significantly accelerates long-sequence inference while preserving model accuracy, establishing a new paradigm for efficient deployment of RoPE-based architectures.

Technology Category

• Machine Learning: Mixture of Experts (MoE)
• Search and Optimization: Non-convex Optimization
• Planning, Routing, and Scheduling: Planning with Language Models

Application Category

• Graph Algorithms and Modeling for the Web: Algorithms and analysis for heterogeneous, signed, attributed, multi-relational, temporal, higher-order, and annotated Web-related graphs
• Search and Retrieval-Augmented AI: Efficiency and scalability of Web search engines
• Responsible Web: Human-perceived consequences of algorithmic deployment on the web

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.

Problem

• Overcoming RoPE's complexity in attention computation
• Enabling almost linear time RoPE attention algorithms
• Combining polynomial method and FFT for efficiency

Innovation

• Combines polynomial method with Fast Fourier Transform
• Enables almost linear time RoPE attention computation
• Works under bounded entry regime constraints