用于序列建模的简单的硬件高效长卷积

Simple Hardware-Efficient Long Convolutions for Sequence Modeling

13 February, 2023

https://arxiv.org/abs/2302.06646

https://github.com/HazyResearch/safari

Daniel Y. Fu, Elliot L. Epstein, Eric Nguyen, Armin W. Thomas, Michael Zhang, Tri Dao, Atri Rudra, Christopher Ré

状态空间模型（SSM）在长序列建模上具有高性能，但需要复杂的初始化技术和专门的实现来实现高质量和运行时性能。我们研究了一种简单的替代方案能否在性能和效率上与SSM相匹配：直接学习序列上的长卷积。我们发现，实现高性能的关键要求是保持卷积核的平滑。我们发现，简单的干预措施（如挤压核权重）可以使内核平滑，并恢复SSM在一系列任务中的性能，包括long range arena、图像分类、语言建模和大脑数据建模。接下来，我们开发FlashButterfly，一种IO感知算法，以提高长卷积的运行时性能。FlashButterfly采用经典的蝶形卷积分解来减少GPU内存IO并提高FLOP利用率。FlashButterfly将卷积速度提高了2.2倍，并允许我们在Path256上进行训练，这是一项具有挑战性的任务，序列长度为64K，我们将最新技术设置为29.1分，训练速度比之前的工作快7.2倍。最后，我们介绍了FlashButterfly的一个扩展，它学习Butterfl分解的系数，在不增加运行时的情况下提高了表现力。使用此扩展，我们比WikiText103上的Transformer性能提高了0.2 PPL，参数减少了30%