SPECTRE: An FFT-Based Efficient Drop-In Replacement to Self-Attention for Long Contexts

TL;DR

SPECTRE introduces an FFT-based self-attention replacement that significantly improves efficiency for long-context transformers, enabling faster processing of tens of thousands of tokens with minimal parameter overhead.

Contribution

It proposes a novel FFT-based attention mechanism that reduces complexity from quadratic to logarithmic, facilitating long-context processing without specialized hardware.

Findings

Operates up to 7× faster than FlashAttention-2 on 128k-token contexts

Matches or exceeds baseline performance on language and vision tasks

Adds fewer than 6% parameters to the base model

Abstract

Long-context transformers face significant efficiency challenges due to the quadratic cost of self-attention. However, many modern applications-from multi-turn dialogue to high-resolution vision-require contexts spanning tens of thousands of tokens. We introduce SPECTRE, a method that replaces each attention head with a fast real FFT, a content-adaptive spectral gate, and an inverse FFT, reducing per-layer complexity from $\mathcal{O}(L^{2})$ to $O(L\log L)$ while preserving the surrounding architecture. We extend this efficiency to autoregressive generation through our Prefix-FFT cache and enhance local feature representation with an optional wavelet module that adds negligible computational overhead. Our experiments demonstrate that SPECTRE operates up to 7$\times$ faster than FlashAttention-2 on 128k-token contexts while matching or exceeding baseline performance on PG-19 language modeling and ImageNet-1k classification tasks. SPECTRE achieves these improvements by adding fewer than 6\% parameters to the base model, making hundred-kilotoken context processing feasible on commodity GPUs without specialized hardware.