摘要

Transformer的出现标志着序列建模领域的一次重大突破，其提供了一种高性能的架构，能够充分利用GPU的并行计算能力。然而，Transformer在推理阶段计算开销较大，限制了其在资源受限场景（如移动设备和嵌入式设备）中的应用。针对这一问题，本文（1）首先表明，注意力机制可被视作一种特殊的循环神经网络（RNN），能够高效计算其“多对一”的RNN输出。随后，（2）我们进一步证明，诸如Transformer等主流基于注意力的模型本质上可被视为RNN的变体。然而，与传统RNN（如LSTM）不同，这些模型在引入新输入标记时难以实现高效的更新，而这一特性在序列建模中至关重要。为解决该问题，（3）我们提出一种基于并行前缀扫描（parallel prefix scan）算法的新方法，用于高效计算注意力机制的“多对多”RNN输出。在此新注意力形式的基础上，（4）我们提出了Aaren——一种基于注意力机制的模块，该模块不仅（i）支持并行训练（如同Transformer），还能够（ii）以高效方式对新输入标记进行更新，且在推理过程中仅需常数级内存（如同传统RNN）。实验结果表明，Aaren在涵盖四个主流序列任务设置（强化学习、事件预测、时间序列分类与时间序列预测）的共38个数据集上，性能可与Transformer相媲美，同时在时间和内存效率方面表现更优。