[Paper Review] Gated recurrent neural networks discover attention

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

Neural Networks and ApplicationsComputer Science3 citations

TL;DR

This paper demonstrates that gated recurrent neural networks (RNNs) with multiplicative gating mechanisms can exactly implement linear self-attention through learned weight configurations, revealing that gradient descent in practice discovers attention-like computation. Trained RNNs on in-context learning tasks implicitly adopt the same attention-based algorithm as linear self-attention, suggesting RNNs may unknowingly perform attention under the hood.

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.

Motivation & Objective

• To investigate whether gated RNNs with multiplicative gating can implement linear self-attention mechanisms.
• To understand how gradient descent in RNNs leads to attention-like computation in practice.
• To compare the inductive biases of different RNN architectures—specifically deep linear RNNs, LSTMs, and GRUs—regarding attention implementation.
• To reverse-engineer trained RNNs and determine whether they learn the same algorithm as linear self-attention on in-context learning tasks.
• To assess the role of multiplicative interactions in enabling RNNs to perform attention-like computation.

Proposed method

• The authors define a class of gated diagonal linear RNNs with input and output gating using elementwise multiplication: $ h_{t+1} = \lambda \odot h_t + g^{\text{in}}(x_t) $, $ y_t = D g^{\text{out}}(h_t) $, where $ g^{\text{in/out}}(x_t) = (W_m^{\text{in/out}} x_t) \odot (W_x^{\text{in/out}} x_t) $.
• They derive a parameter configuration under which such RNNs exactly implement linear self-attention, with the key insight that the recurrent state accumulates $ \sum_t v_t k_t^\top $, analogous to the attention weight matrix.
• The method involves reverse-engineering trained RNNs on linear regression and in-context learning tasks to compare learned weights with the theoretical attention construction.
• They analyze RNNs with input-output gating and side gating, showing that only the side-gated variant successfully emulates linear self-attention under specific weight configurations.
• The authors use JAX and Flax for training, and Numpy, Scikit-learn, and Matplotlib for analysis and visualization of weight patterns and attention matrices.
• They compare the learned $ M(x,y) $ matrices from RNNs to those of linear self-attention, finding rank-1 components with a positive peak and a symmetric negative peak, which does not affect classification due to one-hot encoding.

Experimental results

Research questions

• RQ1 Can gated RNNs with multiplicative gating exactly implement linear self-attention under a specific parameter configuration?
• RQ2 Does gradient descent in trained RNNs lead to the discovery of attention-based computation, even without explicit architectural design for attention?
• RQ3 Why do deep linear RNNs succeed in implementing attention, while LSTMs and GRUs fail under the same configuration?
• RQ4 To what extent do trained RNNs on in-context learning tasks adopt the same algorithm as linear self-attention?
• RQ5 How do multiplicative interactions in RNNs enable or constrain the emergence of attention-like behavior?

Key findings

• Gated RNNs with multiplicative gating can exactly implement linear self-attention under a specific parameter configuration, with the recurrent state accumulating $ \sum_t v_t k_t^\top $.
• Trained RNNs on in-context learning tasks learn to implement the same attention-based algorithm as linear self-attention, indicating that gradient descent discovers attention implicitly.
• The side-gated RNN architecture successfully emulates linear self-attention, with learned weights matching the theoretical construction: $ W_x^{\text{in}} = W^{\text{side}} $, $ W_m^{\text{in}} = [0|B] $, $ D = B^\top $.
• The RNN's output matrix $ M(x,y) $ is of rank 1 and contains a positive peak at $ (i,j) $ and a symmetric negative peak, which does not affect classification due to one-hot encoding.
• LSTMs fail to implement the attention construction under the same configuration, suggesting weaker inductive bias toward attention compared to deep linear RNNs.
• The study reveals that multiplicative interactions are essential for RNNs to implement attention, and that such mechanisms may underlie the success of modern RNNs in sequence modeling.