摘要

我们提出了一种新颖的方法，通过在训练过程中学习选择最具信息量的标记表示来稀疏化Transformer模型中的注意力机制，从而专注于输入中的任务特定部分。由于引入了一个鲁棒的可训练top-$k$算子，时间复杂度和内存复杂度从二次降低到了次线性。我们在一项具有挑战性的长文档摘要任务上的实验表明，即使我们的简单基线模型也能够与当前的最先进方法（SOTA）相媲美，而在采用可训练池化后，我们不仅能够保持其顶级质量，同时在训练时速度提高了1.8倍，在推理时速度提高了4.5倍，并且在解码器中计算效率最高提升了13倍。