FlashFFTConv: Efficient Convolutions for Long Sequences with Tensor Cores

卷积模型使用长滤波器在许多长序列任务中展示了最先进的推理能力，但在墙钟时间上落后于最优化的Transformer。主要瓶颈是快速傅里叶变换（FFT）——它允许长卷积在序列长度N的$O(N logN)$时间内运行，但硬件利用率较差。在本文中，我们研究如何优化FFT卷积。我们发现两个关键瓶颈：FFT没有有效地使用专用的矩阵乘法单元，并且在内存层次结构的层之间产生昂贵的I/O。为此，我们提出了FlashFFTConv。FlashFFTConv使用矩阵分解计算FFT，使用矩阵乘法单元并启用长序列的内核融合，从而减少I/O。我们还提出了两个稀疏卷积算法——1）部分卷积和2）频率稀疏卷积——可以通过跳过矩阵分解中的块来简单实现，从而提供了进一步的内存和计算节省机会。FlashFFTConv将精确FFT卷积的速度提高了最多7.93倍，超过了PyTorch，并实现了最多4.4倍的端到端加速。在相同的计算预算下，FlashFFTConv使Hyena-GPT-s在PILE上实现了比固定参数模型更好的2.3个困惑度，并使M2-BERT-base在GLUE得分上实现了3.3个更高的得分——与参数数量翻倍的模型相匹配。FlashFFTConv还在Path-512上实现了96.1%的准确率，这是一个高分辨率的视觉任务，之前没有任何模型可以实现比50%更好的结果。此外，部分卷积使得更长序列的模型成为可能——产生了第一个可以处理最长人类基因（2.3M个碱基对）的DNA模型——而频率稀疏卷积则可以在保持或提高模型质量的同时加速预训练模型。

优化FFT卷积算法的硬件利用率和内存层次结构之间的I/O开销，以提高卷积模型的计算效率和性能。

FlashFFTConv使用矩阵分解实现FFT卷积，使得FFT卷积能够更有效地使用专用矩阵乘法单元，并减少内存层次结构之间的I/O开销。此外，论文还提出了两种稀疏卷积算法，即部分卷积和频率稀疏卷积，进一步减少内存和计算开销。

FlashFFTConv在精确FFT卷积方面比PyTorch快了多达7.93倍，并实现了高达4.4倍的端到端加速。在相同的计算预算下，FlashFFTConv使Hyena-GPT-s在PILE数据集上的perplexity提高了2.3个点，M2-BERT-base在GLUE数据集上的得分提高了3.3个点，与参数数量翻倍的模型相匹配。FlashFFTConv还在Path-512数据集上实现了96.1%的准确率，这是一个高分辨率视觉任务，之前没有任何模型能够达到50%以上的准确率。此外，部分卷积使得模型能够处理更长的序列，实现了第一个能够处理最长人类基因（2.3M个碱基对）的DNA模型，而频率稀疏卷积则在保持或提高模型质量的同时加速了预训练模型。

相关论文：1）Efficient Transformers: A Survey；2）Longformer: The Long-Document Transformer；3）Big Bird: Transformers for Longer Sequences；4）Sparse Transformers；5）Reformer: The Efficient Transformer；6）Linformer: Self-Attention with Linear Complexity；7）Performer: Fast and Scalable Attention；8）Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention；作者和机构：1）Zihang Dai, et al.，Google Research；2）Iz Beltagy, et al.，AI2；3）Manzil Zaheer, et al.，Google Research；4）Dense Sparse Dense Training for Deep Neural Networks；5）Nikita Kitaev, et al.，Google Research；6）Sinong Wang, et al.，Tsinghua University；7）Krishnamurthy Dvijotham, et al.，Google Research；8）Hanxiao Liu, et al.，MIT CSAIL。