SPECTRE:基于快速傱里叶变换的自注意力长上下文高效替代方案
机器学习
2025-05-20 v7
摘要
长上下文变压器因自注意力的二次计算成本,面临显著的效率挑战。然而,许多现代应用——从多轮对话到高分辨率视觉任务——都需要处理数万甚至数十万token的上下文。我们提出SPECTRE方法,通过用快速实数傱里叶变换、内容自适应谱门和逆傱里叶变换替换每个注意力头,将单层复杂度从 降至 ,同时保持周围架构不变。我们通过 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}
}