[论文解读] Simple Hardware-Efficient Long Convolutions for Sequence Modeling

Daniel Y. Fu, Elliot L. Epstein+6 more|arXiv (Cornell University)|Feb 13, 2023

Topic Modeling被引用 6

一句话总结

该论文提出对长卷积核进行直接参数化以用于序列建模，表明通过简单的正则化方法（如压缩核权重）即可获得平滑的核函数，从而实现与状态空间模型（SSM）相当的性能。借助一种面向输入输出（I/O）的FlashButterfly算法，该方法利用Butterfly分解，在长卷积中实现了相较于cuFFT 2.2倍的加速，并在64K长度序列上实现7.2倍的训练速度提升，优于SSM和Transformer在多个基准测试中的表现，包括Path256和脑部fMRI数据。

ABSTRACT

State space models (SSMs) have high performance on long sequence modeling but require sophisticated initialization techniques and specialized implementations for high quality and runtime performance. We study whether a simple alternative can match SSMs in performance and efficiency: directly learning long convolutions over the sequence. We find that a key requirement to achieving high performance is keeping the convolution kernels smooth. We find that simple interventions--such as squashing the kernel weights--result in smooth kernels and recover SSM performance on a range of tasks including the long range arena, image classification, language modeling, and brain data modeling. Next, we develop FlashButterfly, an IO-aware algorithm to improve the runtime performance of long convolutions. FlashButterfly appeals to classic Butterfly decompositions of the convolution to reduce GPU memory IO and increase FLOP utilization. FlashButterfly speeds up convolutions by 2.2$ imes$, and allows us to train on Path256, a challenging task with sequence length 64K, where we set state-of-the-art by 29.1 points while training 7.2$ imes$ faster than prior work. Lastly, we introduce an extension to FlashButterfly that learns the coefficients of the Butterfly decomposition, increasing expressivity without increasing runtime. Using this extension, we outperform a Transformer on WikiText103 by 0.2 PPL with 30% fewer parameters.

研究动机与目标

• 探究是否可通过直接参数化长卷积在无需复杂初始化或专用实现的情况下，达到状态空间模型（SSM）在序列建模中的性能表现。
• 解决现代硬件（尤其是GPU内存I/O瓶颈和未充分利用的算术强度）下长卷积的运行时效率低下问题。
• 探索结构化矩阵分解（如Butterfly分解）是否能同时提升长卷积的效率与表达能力。
• 评估所提方法在多种模态（包括自然语言、视觉和神经影像）中的泛化能力。
• 证明通过简单正则化与高效算法设计，可在保持标准卷积性能的同时，缩小与SSM的性能差距，并实现长序列上的更快训练。

提出的方法

• 对核权重应用Squash操作，以强制实现稀疏性并促进频域中的平滑性，从而提升模型质量。
• 在时域中引入Smooth操作，进一步正则化核权重，增强频域平滑性。
• 开发FlashButterfly，一种面向I/O的算法，通过使用分块稀疏Butterfly矩阵分解重写基于FFT的卷积，以减少内存I/O并提升算术强度利用率。
• 通过Butterfly分解将卷积重构为一系列稀疏矩阵乘法，以更高效地利用GPU矩阵乘法单元。
• 引入可学习的Butterfly扩展，使分解系数可训练，从而在不增加推理成本的前提下提升表达能力。
• 将所得架构作为S4和H3等模型中SSM层的即插即用替代品，并在LRA、CIFAR、WikiText103和fMRI等多样化基准上进行评估。

实验结果

研究问题

• RQ1 在无需复杂初始化或专用核函数的情况下，直接参数化的长卷积是否能在长序列建模任务中达到与SSM相当的性能？
• RQ2 哪些简单的正则化技术可通过促进核函数的平滑行为，使长卷积与SSM具有竞争力？
• RQ3 面向I/O的算法（如FlashButterfly）是否能显著提升现代GPU上长卷积的运行时效率？
• RQ4 使用结构化Butterfly分解是否能同时实现加速与更高的表达能力？
• RQ5 在fMRI序列建模等具有挑战性的现实应用场景中，长卷积是否能超越Transformer和SSM？

主要发现

• 简单的正则化（特别是压缩和光滑化核权重）可生成平滑的频域核函数，并在长程任务基准（LRA）上恢复SSM级别的性能，即使在随机初始化下也与S4性能相当。
• 在图像分类任务中，长卷积在顺序CIFAR上比S4高出0.3%的准确率，并在2D CIFAR上与S4ND-ISO相差仅0.8%。
• 在语言建模任务中，长卷积在PILE数据集上的表现与H3相当（PPL相差仅0.3），在OpenWebText上也仅相差0.3 PPL。
• 在脑部fMRI建模中，长卷积分别以0.14和0.16的MAE点数优势超越Transformer和SSM。
• FlashButterfly使长卷积在cuFFT基础上实现2.2倍加速，并支持在Path256（64K序列长度）上训练，实现SOTA性能，准确率比之前工作高出29.1分，训练速度提升7.2倍。
• 可学习的Butterfly扩展在不增加运行时成本的前提下提升了表达能力，使长卷积在WikiText103上以30%更少的参数量，实现0.2 PPL的PPL优势，超越Transformer。