[論文レビュー] 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、LSTM、GRUといった異なるRNNアーキテクチャの、注意実装に関する誘導的バイアスを比較すること。
• 訓練済み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 特定のパrameter設定下で、乗法的ゲーティングを備えたゲート付きRNNは線形自己注意を正確に実装できるか？
• RQ2 明示的な注意用アーキテクチャ設計がなくても、訓練済みRNNにおける勾配降下法が注意に基づく計算を発見するのか？
• RQ3 同じ設定下で深層線形RNNは注意を実装できるが、LSTMやGRUは失敗する理由は何か？
• RQ4 文脈内学習タスクで訓練されたRNNが、線形自己注意と同一のアルゴリズムをどの程度採用しているか？
• RQ5 RNNにおける乗法的相互作用が、注意に類似した挙動の出現をどのように促進または制限するか？

主な発見

• 特定のパrameter設定下で、乗法的ゲーティングを備えたゲート付き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が系列モデリングで成功する背後にある要因である可能性を示している。