以线性时间将注意力扩展到极长序列:小波增强随机谱注意力(WERSA)

Transformer 模型在长序列上计算开销巨大,因为常规注意力具有 $O(n^2)$ 的二次时间复杂度.我们引入小波增强随机谱注意力(WERSA),一种具有线性 $O(n)$ 时间复杂度的全新机制,对于在不牺牲性能的前提下实现成功的长序列处理至关重要.WERSA 将内容自适应的随机谱特征与多分辨率 Haar 小波以及可学习参数相融合,能够在保持线性效率的同时选择性地关注数据中信息丰富的尺度.在单 GPU 上的大规模比较以及跨多种基准(视觉、NLP、层次推理)和多种注意力机制(如多头注意力、Flash-Attention-2、FNet、Linformer、Performer、Waveformer)的实验表明,WERSA 具有一致的优势.它在所有测试中均取得了最佳准确率.在 ArXiv 分类任务上,WERSA 相比 vanilla 注意力将准确率提升 1.2%(86.2% 对 85.0%),同时将训练时间削减 81%(296 秒对 1554 秒),FLOPS 削减 73.4%(26.2G 对 98.4G).值得注意的是,在 vanilla 和 FlashAttention-2 失效的场景下 WERSA 表现卓越:在 ArXiv-128k 的极长序列上,它在可行方法中取得了最佳准确率(79.1%)和 AUC(0.979),能够处理使二次方法出现显存溢出错误的数据,同时速度是其次优竞争者 Waveformer 的两倍.通过在不牺牲准确率的前提下显著降低计算负载,WERSA 使得更实用、更经济的长上下文模型成为可能,尤其是在低资源硬件上,为更可持续、更具可扩展性的 AI 发展提供了支持.