专栏名称: 极市平台

极市平台是由深圳极视角推出的专业的视觉算法开发与分发平台，为视觉开发者提供多领域实景训练数据库等开发工具和规模化销售渠道。本公众号将会分享视觉相关的技术资讯，行业动态，在线分享信息，线下活动等。网站: http://cvmart.net/

图灵奖得主Bengio团队新作：注意力可视为RNN！新模型媲美Transformer

极市平台 · 公众号 · · 2024-05-28 22:00

正文

↑ 点击蓝字关注极市平台

作者丨机器之心

来源丨机器之心

编辑丨极市平台

极市导读

既能像 Transformer 一样并行训练，推理时内存需求又不随 token 数线性递增，长上下文又有新思路了？ >> 加入极市CV技术交流群，走在计算机视觉的最前沿

序列建模的进展具有极大的影响力，因为它们在广泛的应用中发挥着重要作用，包括强化学习（例如，机器人和自动驾驶）、时间序列分类（例如，金融欺诈检测和医学诊断）等。

在过去的几年里，Transformer 的出现标志着序列建模中的一个重大突破，这主要得益于 Transformer 提供了一种能够利用 GPU 并行处理的高性能架构。

然而，Transformer 在推理时计算开销很大，主要在于内存和计算需求呈二次扩展，从而限制了其在低资源环境中的应用（例如，移动和嵌入式设备）。尽管可以采用 KV 缓存等技术提高推理效率，但 Transformer 对于低资源领域来说仍然非常昂贵，原因在于：（1）随 token 数量线性增加的内存，以及（2）缓存所有先前的 token 到模型中。在具有长上下文（即大量 token）的环境中，这一问题对 Transformer 推理的影响更大。

为了解决这个问题，加拿大皇家银行 AI 研究所 Borealis AI、蒙特利尔大学的研究者在论文《Attention as an RNN 》中给出了解决方案。值得一提的是，我们发现图灵奖得主 Yoshua Bengio 出现在作者一栏里。

论文地址： https://arxiv.org/pdf/2405.13956

具体而言，研究者首先检查了 Transformer 中的注意力机制，这是导致 Transformer 计算复杂度呈二次增长的组件。该研究表明注意力机制可以被视为一种特殊的循环神经网络（RNN），具有高效计算的多对一（many-to-one）RNN 输出的能力。利用注意力的 RNN 公式，该研究展示了流行的基于注意力的模型（例如 Transformer 和 Perceiver）可以被视为 RNN 变体。

然而，与 LSTM、GRU 等传统 RNN 不同，Transformer 和 Perceiver 等流行的注意力模型虽然可以被视为 RNN 变体。但遗憾的是，它们无法高效地使用新 token 进行更新。

为了解决这个问题，该研究引入了一种基于并行前缀扫描（prefix scan）算法的新的注意力公式，该公式能够高效地计算注意力的多对多（many-to-many）RNN 输出，从而实现高效的更新。

在此新注意力公式的基础上，该研究提出了 Aaren（[A] ttention [a] s a [re] current neural [n] etwork），这是一种计算效率很高的模块，不仅可以像 Transformer 一样并行训练，还可以像 RNN 一样高效更新。

实验结果表明，Aaren 在 38 个数据集上的表现与 Transformer 相当，这些数据集涵盖了四种常见的序列数据设置：强化学习、事件预测、时间序列分类和时间序列预测任务，同时在时间和内存方面更加高效。

方法介绍

为了解决上述问题，作者提出了一种基于注意力的高效模块，它能够利用 GPU 并行性，同时又能高效更新。

首先，作者在第 3.1 节中表明，注意力可被视为一种 RNN，具有高效计算多对一 RNN（图 1a）输出的特殊能力。利用注意力的 RNN 形式，作者进一步说明，基于注意力的流行模型，如 Transformer（图 1b）和 Perceiver（图 1c），可以被视为 RNN。然而，与传统的 RNN 不同的是，这些模型无法根据新 token 有效地更新自身，从而限制了它们在数据以流的形式到达的序列问题中的潜力。

为了解决这个问题，作者在第 3.2 节中介绍了一种基于并行前缀扫描算法的多对多 RNN 计算注意力的高效方法。在此基础上，作者在第 3.3 节中介绍了 Aaren—— 一个计算效率高的模块，它不仅可以并行训练（就像 Transformer），还可以在推理时用新 token 高效更新，推理只需要恒定的内存（就像传统 RNN）。

将注意力视为一个多对一 RNN

查询向量的注意力可被视为一个函数, 它通过个上下文 token x_1: 的键和值将其映射到单一输出 o_N = Attention (, k_1:N , v_1:N ) 。给定 s_i = dot (q, k_i), 输出 o_N 可表述为:

其中分子为 , 分母为。将注意力视为 RNN, 可以在时, 以滚动求和的方式迭代计算和。然而, 在实践中, 这种实现方式并不稳定, 会因有限的精度表示和可能非常小或非常大的指数（即 ) 而遇到数值问题。为了缓解这一问题, 作者用累积最大值项来重写递推公式 , 计算和。值得注意的是, 最终结果是相同的的循环计算如下:

通过从 a_(k-1)、c_(k-1) 和 m_(k-1) 对 a_k、c_k 和 m_k 的循环计算进行封装，作者引入了一个 RNN 单元，它可以迭代计算注意力的输出（见图 2）。注意力的 RNN 单元以（a_(k-1), c_(k-1), m_(k-1), q）作为输入，并计算（a_k, c_k, m_k, q）。注意，查询向量 q 在 RNN 单元中被传递。注意力 RNN 的初始隐藏状态为 (a_0, c_0, m_0, q) = (0, 0, 0, q)。

计算注意力的方法：通过将注意力视为一个 RNN，可以看到计算注意力的不同方法：在 O (1) 内存中逐个 token 循环计算（即顺序计算）；或以传统方式计算（即并行计算），需要线性 O (N) 内存。由于注意力可以被看作是一个 RNN，因此计算注意力的传统方法也可以被看作是计算注意力多对一 RNN 输出的高效方法，即 RNN 的输出以多个上下文 token 为输入，但在 RNN 结束时只输出一个 token（见图 1a）。最后，也可以将注意力计算为一个逐块处理 token 的 RNN，而不是完全按顺序或完全并行计算，这需要 O (b) 内存，其中 b 是块的大小。

将现有的注意力模型视为 RNN。通过将注意力视为 RNN，现有的基于注意力的模型也可以被视为 RNN 的变体。例如，Transformer 的自注意力是 RNN（图 1b），上下文 token 是其初始隐藏状态。Perceiver 的交叉注意力是 RNN（图 1c），其初始隐藏状态是与上下文相关的潜变量。通过利用其注意力机制的 RNN 形式，这些现有模型可以高效地计算其输出存储。

然而，当将现有的基于注意力的模型（如 Transformers）视为 RNN 时，这些模型又缺乏传统 RNN（如 LSTM 和 GRU）中常见的重要属性。

值得注意的是，LSTM 和 GRU 能够仅在 O (1) 常量内存和计算中使用新 token 有效地更新自身，相比之下， Transformer 的 RNN 视图（见图 1b）会通过将一个新的 token 作为初始状态添加一个新的 RNN 来处理新 token。这个新的 RNN 处理所有先前的 token，需要 O (N) 的线性计算量。

在 Perceiver 中，由于其架构的原因，潜变量（图 1c 中的 L_i）是依赖于输入的，这意味着它们的值在接收新 token 时会发生变化。由于其 RNN 的初始隐藏状态（即潜变量）发生变化，Perceiver 因此需要从头开始重新计算其 RNN，需要 O (NL) 的线性计算量，其中 N 是 token 的数量，L 是潜变量的数量。

将注意力视为一个多对多 RNN

针对这些局限性, 作者建议开发一种基于注意力的模型, 利用 RNN 公式的能力来执行高效更新。为此, 作者首先引入了一种高效的并行化方法, 将注意力作为多对多 RNN 计算，即并行计算的方法。为此，作者利用并行前缀扫描算法（见算法 1），这是一种通过关联算子 ⊕ 从 N 个连续数据点计算 N 个前缀的并行计算方法。该算法可高效计算