注意力即RNN
机器学习
2024-05-29 v2
摘要
Transformer的出现标志着序列建模的重大突破，提供了一种能够利用GPU并行性的高性能架构。然而，Transformer在推理时计算成本高昂，限制了其应用，尤其是在低资源环境（如移动和嵌入式设备）中。针对这一问题，我们（1）首先证明注意力机制可以被视为一种特殊的循环神经网络（RNN），能够高效计算其\textit{多对一}RNN输出。然后（2）表明流行的基于注意力的模型（如Transformer）可被视为RNN变体。然而，与传统的RNN（如LSTM）不同，这些模型无法高效地使用新令牌进行更新，而这是序列建模中的一个重要特性。为解决这一问题，我们（3）引入了一种基于并行前缀扫描算法的高效计算注意力\textit{多对多}RNN输出的新方法。基于新的注意力公式，我们（4）引入了\textbf{Aaren}，一个基于注意力的模块，它不仅可以（i）像Transformer一样并行训练，而且可以（ii）高效地使用新令牌进行更新，推理时仅需恒定内存（如传统RNN）。实验上，我们展示了Aaren在涵盖强化学习、事件预测、时间序列分类和时间序列预测任务四个流行序列问题设置的38个数据集上取得了与Transformer相当的性能，同时更加节省时间和内存。
引用
@article{arxiv.2405.13956,
title = {Attention as an RNN},
author = {Leo Feng and Frederick Tung and Hossein Hajimirsadeghi and Mohamed Osama Ahmed and Yoshua Bengio and Greg Mori},
journal= {arXiv preprint arXiv:2405.13956},
year = {2024}
}