[论文解读] Gated recurrent neural networks discover attention

Nicolas Zucchet, Seijin Kobayashi+5 more|arXiv (Cornell University)|Sep 4, 2023

Neural Networks and ApplicationsComputer Science被引用 3

一句话总结

该论文表明，具有乘法门控机制的门控循环神经网络（RNN）可通过学习到的权重配置精确实现线性自注意力机制，揭示了梯度下降在实践中会发现类似注意力的计算。在上下文学习任务上训练的RNN隐式地采用了与线性自注意力相同的基于注意力的算法，表明RNN可能在内部不知不觉地执行注意力机制。

ABSTRACT

Recent architectural developments have enabled recurrent neural networks (RNNs) to reach and even surpass the performance of Transformers on certain sequence modeling tasks. These modern RNNs feature a prominent design pattern: linear recurrent layers interconnected by feedforward paths with multiplicative gating. Here, we show how RNNs equipped with these two design elements can exactly implement (linear) self-attention, the main building block of Transformers. By reverse-engineering a set of trained RNNs, we find that gradient descent in practice discovers our construction. In particular, we examine RNNs trained to solve simple in-context learning tasks on which Transformers are known to excel and find that gradient descent instills in our RNNs the same attention-based in-context learning algorithm used by Transformers. Our findings highlight the importance of multiplicative interactions in neural networks and suggest that certain RNNs might be unexpectedly implementing attention under the hood.

研究动机与目标

• 探究具有乘法门控的门控RNN是否能在特定参数配置下实现线性自注意力机制。
• 理解RNN中梯度下降如何在实践中导致类似注意力的计算。
• 比较不同RNN架构（特别是深度线性RNN、LSTM和GRU）在注意力实现方面的归纳偏置。
• 对训练后的RNN进行逆向工程，判断其在上下文学习任务上是否学习到与线性自注意力相同的算法。
• 评估乘法交互在使RNN实现类似注意力计算中的作用。

提出的方法

• 作者定义了一类使用逐元素乘法进行输入和输出门控的门控对角线线性RNN：$ h_{t+1} = \lambda \odot h_t + g^{\text{in}}(x_t) $，$ y_t = D g^{\text{out}}(h_t) $，其中 $ g^{\text{in/out}}(x_t) = (W_m^{\text{in/out}} x_t) \odot (W_x^{\text{in/out}} x_t) $。
• 他们推导出一种参数配置，使得此类RNN能精确实现线性自注意力，关键洞见在于递归状态累积了 $ \sum_t v_t k_t^\top $，类似于注意力权重矩阵。
• 该方法涉及对线性回归和上下文学习任务上训练的RNN进行逆向工程，以比较学习到的权重与理论注意力结构。
• 他们分析了具有输入输出门控和侧门控的RNN，表明只有侧门控变体在特定权重配置下能成功模拟线性自注意力。
• 作者使用JAX和Flax进行训练，并使用Numpy、Scikit-learn和Matplotlib进行权重模式和注意力矩阵的分析与可视化。
• 他们将RNN学习到的 $ M(x,y) $ 矩阵与线性自注意力的矩阵进行比较，发现存在一个正峰和一个对称负峰的秩-1分量，由于独热编码的存在，该分量不影响分类结果。

实验结果

研究问题

• RQ1 在特定参数配置下，具有乘法门控的门控RNN是否能精确实现线性自注意力？
• RQ2 在未显式设计为注意力架构的情况下，训练RNN中的梯度下降是否会引导其发现基于注意力的计算？
• RQ3 为何深度线性RNN能在相同配置下成功实现注意力，而LSTM和GRU却失败？
• RQ4 在上下文学习任务上训练的RNN在多大程度上采用了与线性自注意力相同的算法？
• RQ5 RNN中的乘法交互如何促进或限制类似注意力行为的出现？

主要发现

• 在特定参数配置下，具有乘法门控的门控RNN可精确实现线性自注意力，其递归状态累积了 $ \sum_t v_t k_t^\top $。
• 在上下文学习任务上训练的RNN学习到实现与线性自注意力相同的基于注意力的算法，表明梯度下降会隐式发现注意力机制。
• 侧门控RNN架构成功模拟了线性自注意力，学习到的权重与理论构造一致：$ W_x^{\text{in}} = W^{\text{side}} $，$ W_m^{\text{in}} = [0|B] $，$ D = B^\top $。
• RNN的输出矩阵 $ M(x,y) $ 为秩-1矩阵，且在 $ (i,j) $ 处存在正峰和对称负峰，由于独热编码的存在，该结构不影响分类结果。
• LSTM在相同配置下无法实现注意力结构，表明其对注意力的归纳偏置弱于深度线性RNN。
• 本研究揭示了乘法交互对RNN实现注意力机制至关重要，且此类机制可能正是现代RNN在序列建模中取得成功的原因。