摘要

我们提出了一种新颖的方法，通过在训练过程中学习选择最具信息量的标记表示来稀疏化Transformer模型中的注意力机制，从而聚焦于输入数据中与任务相关的部分。得益于一个稳健的可训练Top-$k$操作器，该方法将原本二次方的时间和内存复杂度降低至次线性级别。在一项具有挑战性的长文档摘要任务上的实验表明，即使我们的简单基线模型也已达到当前最先进（SOTA）水平，而通过引入可训练的池化机制，我们能够在保持顶级性能的同时，使训练速度提升1.8倍，推理速度提升4.5倍，并在解码器阶段实现高达13倍的计算效率提升。