面向序列建模的简单硬件高效长卷积
机器学习
2023-02-15 v1
摘要
状态空间模型（SSMs）在长序列建模上性能优异，但需要复杂的初始化技术与专用实现以保证质量与运行效率。我们研究一种简单替代方案能否在性能与效率上匹敌SSMs：直接学习序列上的长卷积。我们发现实现高性能的一个关键要求是保持卷积核平滑。我们发现简单干预——例如压缩核权重——可产生平滑核，并在包括长程竞技场、图像分类、语言建模与脑数据建模的一系列任务上恢复SSM性能。接下来，我们开发了FlashButterfly，一种IO感知算法以改善长卷积的运行效率。FlashButterfly借助经典的蝴蝶分解来减少GPU内存IO并提高FLOP利用率。FlashButterfly将卷积加速2.2，并使我们能在Path256这一序列长度64K的极具挑战性任务上训练，我们以领先先前工作29.1个点刷新最优结果，同时训练快7.2。最后，我们引入FlashButterfly的扩展以学习蝴蝶分解的系数，在不增加运行时间的前提下提升表达能力。使用该扩展，我们在WikiText103上以少30%的参数超越Transformer 0.2 PPL。
引用
@article{arxiv.2302.06646,
title = {Simple Hardware-Efficient Long Convolutions for Sequence Modeling},
author = {Daniel Y. Fu and Elliot L. Epstein and Eric Nguyen and Armin W. Thomas and Michael Zhang and Tri Dao and Atri Rudra and Christopher Ré},
journal= {arXiv preprint arXiv:2302.06646},
year = {2023}
}