摘要
自注意力机制近年来被广泛应用于各类序列建模任务中。尽管其表现优异，但自注意力机制在序列长度上的计算和内存开销呈二次方增长，限制了其在长序列场景下的应用。此前的高效化方法主要聚焦于局部滑动窗口或与内容无关的固定稀疏位置的注意力机制。本文提出一种动态稀疏注意力模式的学习方法，能够避免将计算与内存资源分配给与目标查询无关的序列内容。本研究融合了两条技术路线：一方面继承了基于内容的稀疏注意力机制在建模灵活性方面的优势，另一方面结合了局部时序稀疏注意力方法在效率上的提升。我们提出的模型——路由Transformer（Routing Transformer），引入基于在线k-means的稀疏路由模块，将自注意力机制的整体复杂度从传统的 $O(n^2d)$ 降低至 $O(n^{1.5}d)$，其中 $n$ 为序列长度，$d$ 为隐藏维度。实验结果表明，在 Wikitext-103 语言建模任务中，我们的模型以 18.3 的困惑度优于对比模型的 15.8；在 ImageNet-64 图像生成任务中，达到 3.43 bits/dim 的性能，略优于对比模型的 3.44 bits/dim，同时使用更少的自注意力层。此外，在新发布的 PG-19 数据集上，我们取得了新的最先进水平，采用 22 层的路由Transformer模型，训练序列长度为 8192，测试困惑度达到 33.2。