中文

SPECTRE:基于快速傱里叶变换的自注意力长上下文高效替代方案

机器学习 2025-05-20 v7

摘要

长上下文变压器因自注意力的二次计算成本,面临显著的效率挑战。然而,许多现代应用——从多轮对话到高分辨率视觉任务——都需要处理数万甚至数十万token的上下文。我们提出SPECTRE方法,通过用快速实数傱里叶变换、内容自适应谱门和逆傱里叶变换替换每个注意力头,将单层复杂度从 O(L2)\mathcal{O}(L^{2}) 降至 O(LlogL)O(L\log L),同时保持周围架构不变。我们通过 Prefix-FFT 缓存将这种效率优势扩展到自回归生成过程中,并通过可选的小波模块增强局部特征表示,计算开销可忽略不计。我们的实验表明,SPECTRE 在 128k token 上比 FlashAttention-2 快达 7 倍,在 PG-19 语言建模和 ImageNet-1k 分类任务上表现出与基线相当甚至更好的效果。通过仅增加不足 6% 的参数,SPECTRE 使百万级 token 上下文处理在无需专用硬件的常规 GPU 上变得可行。

关键词

引用

@article{arxiv.2502.18394,
  title  = {SPECTRE: An FFT-Based Efficient Drop-In Replacement to Self-Attention for Long Contexts},
  author = {Jacob Fein-Ashley and Neelesh Gupta and Rajgopal Kannan and Viktor Prasanna},
  journal= {arXiv preprint arXiv:2502.18394},
  year   = {2025}
}