[论文解读] Attention as an RNN

Leo Feng, Frederick Tung+4 more|arXiv (Cornell University)|May 22, 2024

Educational and Psychological AssessmentsPsychology被引用 3

一句话总结

本文提出 Aaren，一种新颖的注意力机制，将自注意力重新表述为可微分的循环神经网络（RNN），从而在保持 Transformer 模型并行训练能力的同时，实现高效、恒定内存的推理，且支持逐 token 的增量更新。Aaren 在 38 项序列建模任务中达到与 Transformer 相当的性能，同时将计算和内存复杂度从二次方级降低为线性级。

ABSTRACT

The advent of Transformers marked a significant breakthrough in sequence modelling, providing a highly performant architecture capable of leveraging GPU parallelism. However, Transformers are computationally expensive at inference time, limiting their applications, particularly in low-resource settings (e.g., mobile and embedded devices). Addressing this, we (1) begin by showing that attention can be viewed as a special Recurrent Neural Network (RNN) with the ability to compute its extit{many-to-one} RNN output efficiently. We then (2) show that popular attention-based models such as Transformers can be viewed as RNN variants. However, unlike traditional RNNs (e.g., LSTMs), these models cannot be updated efficiently with new tokens, an important property in sequence modelling. Tackling this, we (3) introduce a new efficient method of computing attention's extit{many-to-many} RNN output based on the parallel prefix scan algorithm. Building on the new attention formulation, we (4) introduce extbf{Aaren}, an attention-based module that can not only (i) be trained in parallel (like Transformers) but also (ii) be updated efficiently with new tokens, requiring only constant memory for inferences (like traditional RNNs). Empirically, we show Aarens achieve comparable performance to Transformers on $38$ datasets spread across four popular sequential problem settings: reinforcement learning, event forecasting, time series classification, and time series forecasting tasks while being more time and memory-efficient.

研究动机与目标

• 为解决 Transformer 在移动设备和嵌入式设备等低资源环境下的高推理成本问题。
• 弥合 Transformer 的并行性与 RNN 的高效、增量式推理之间的差距。
• 开发一种新型注意力形式，同时支持并行训练与恒定内存、流式推理。
• 提出一种方法，利用并行前缀扫描高效计算注意力的多对多 RNN 输出。
• 证明所提出的 Aaren 模块在保持 Transformer 性能的同时，显著降低时间与内存复杂度。

提出的方法

• 将标准自注意力重新表述为一种多对一 RNN 的特例，从而为注意力计算提供递归解释。
• 提出一种基于新型 RNN 的注意力形式，利用并行前缀扫描算法高效计算多对多 RNN 输出。
• 推导出基于 softmax 的注意力的可微分、精确重实现，作为 RNN，避免了先前近似方法中出现的指数时间戳偏差。
• 设计 Aaren 为可学习的 RNN 模块，保持恒定的隐藏状态大小，并支持新 token 的增量更新。
• 采用可学习的初始隐藏状态向量 $ q $，使 RNN 能够捕获长距离依赖关系，而无需存储过去的激活值。
• 利用高效的并行前缀和算法（如 Hillis-Steele）在 $ O(N \log N) $ 时间内计算注意力权重，实现可扩展的推理。

实验结果

研究问题

• RQ1 标准自注意力能否被正式重新解释为循环神经网络？
• RQ2 为何基于注意力的模型（如 Transformer）尽管具备并行训练能力，却无法支持高效的增量推理？
• RQ3 是否可能设计一种注意力机制，同时结合 Transformer 的并行性与 RNN 的恒定内存效率？
• RQ4 基于并行前缀扫描的新注意力形式能否在降低计算复杂度的同时，实现与 Transformer 竞争的性能？
• RQ5 在多种序列任务中，所提出的 Aaren 模块与 Transformer 相比，在推理效率和模型性能方面表现如何？

主要发现

• Aaren 在 38 项涵盖强化学习、事件预测、时间序列分类与时间序列预测的多样化序列建模数据集上，性能与 Transformer 相当。
• Aaren 的累积推理时间与内存使用随序列长度呈线性增长，而采用 KV 缓存的 Transformer 则呈二次方增长。
• 与等效的 Transformer 相比，Aaren 仅增加 0.016% 的参数量（3,152,896 对比 3,152,384），使效率提升具有成本效益。
• 该模型在推理过程中保持恒定的内存使用，适用于移动设备和嵌入式系统等低资源设备的部署。
• 基于并行前缀扫描的注意力计算实现了高效、可微分的 RNN 式推理，同时保留了完整注意力机制的表达能力。
• 实证结果证实，Aaren 在所有评估任务中均优于或匹配 Transformer 的速度与内存效率。