WERSA让长序列注意力计算变快10倍,还更准
Scaling Attention to Very Long Sequences in Linear Time with Wavelet-Enhanced Random Spectral Attention (WERSA)
- 用小波和随机谱特征混合,实现线性时间注意力
- 在论文分类任务上准确率提升1.2%,训练快81%
- 适合低资源设备运行,特别适合超长序列处理
Transformer 模型在长序列上计算成本高,因标准注意力具有二次方 $O(n^2)$ 时间复杂度。我们提出一种新型机制 WERSA,具备线性 $O(n)$ 时间复杂度,可实现无性能损失的长序列处理。WERSA 融合内容自适应的随机谱特征、多分辨率哈尔小波与可学习参数,能选择性关注数据中的有效尺度,同时保持线性效率。大规模单卡对比实验涵盖视觉、NLP 和分层推理等多种基准及多种注意力机制(如 Multiheaded Attention、Flash-Attention-2、FNet、Linformer、Performer、Waveformer),结果均显示 WERSA 统一领先。在 ArXiv 分类任务中,其准确率相比原生注意力提升 1.2%(86.2% vs 85.0%),训练时间减少 81%(296s vs 1554s),FLOPS 降低 73.4%(26.2G vs 98.4G)。尤为关键的是,在极端长序列的 ArXiv-128k 数据上,其他方法因显存溢出失败,而 WERSA 达到最佳准确率(79.1%)和 AUC(0.979),且速度为次优方法 Waveformer 的两倍。通过显著降低计算负载而不牺牲精度,WERSA 使低成本硬件上的长上下文模型成为可能,推动更可持续、可扩展的 AI 发展。
原文摘要 · Abstract (English)
Transformer models are computationally costly on long sequences since regular attention has quadratic $O(n^2)$ time complexity. We introduce Wavelet-Enhanced Random Spectral Attention (WERSA), a novel mechanism of linear $O(n)$ time complexity that is pivotal to enable successful long-sequence processing without the performance trade-off. WERSA merges content-adaptive random spectral features together with multi-resolution Haar wavelets and learnable parameters to selectively attend to informative scales of data while preserving linear efficiency. Large-scale comparisons \textbf{on single GPU} and across various benchmarks (vision, NLP, hierarchical reasoning) and various attention mechanisms (like Multiheaded Attention, Flash-Attention-2, FNet, Linformer, Performer, Waveformer), reveal uniform advantages of WERSA. It achieves best accuracy in all tests. On ArXiv classification, WERSA improves accuracy over vanilla attention by 1.2\% (86.2\% vs 85.0\%) while cutting training time by 81\% (296s vs 1554s) and FLOPS by 73.4\% (26.2G vs 98.4G). Significantly, WERSA excels where vanilla and FlashAttention-2 fail: on ArXiv-128k's extremely lengthy sequences, it achieves best accuracy (79.1\%) and AUC (0.979) among viable methods, operating on data that gives Out-Of-Memory errors to quadratic methods while being \textbf{twice as fast} as Waveformer, its next-best competitor. By significantly reducing computational loads without compromising accuracy, WERSA makes possible more practical, more affordable, long-context models, in particular on low-resource hardware, for more sustainable and more scalable AI development.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。