用快速傅里叶变换替代自注意力,让长文本处理快7倍且不丢性能。
SPECTRE: An FFT-Based Efficient Drop-In Replacement to Self-Attention for Long Contexts
- 用FFT+谱门控替换注意力头,复杂度从O(L²)降到O(L log L)
- 128k token上下文下速度比FlashAttention-2快7倍,任务性能相当
- 参数增加少于6%,能在普通显卡跑超长文本任务
长序列Transformer因自注意力的二次复杂度面临显著效率瓶颈。现代应用如多轮对话和高分辨率视觉需处理数万标记的上下文。我们提出SPECTRE,将每个注意力头替换为快速实数FFT、内容自适应谱门控和逆FFT,使每层复杂度从$/mathcal{O}(L^{2})$降至$O(L ext{log}L)$,同时保持原有架构。通过前缀FFT缓存实现自回归生成的效率扩展,并引入可选的小波模块增强局部特征表示,计算开销极低。实验表明,SPECTRE在128k标记上下文上运行速度比FlashAttention-2快达7倍,在PG-19语言建模和ImageNet-1k分类任务中性能持平或更优。其额外参数不足基础模型的6%,使百万级标记上下文处理可在消费级GPU上实现,无需专用硬件。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。