[論文レビュー] Were RNNs All We Needed?

Leo Feng, Frederick Tung+3 more|arXiv (Cornell University)|Oct 2, 2024

Nursing Education, Practice, and Leadership被引用数 5

ひとこと要約

本論文は、隠れ状態への依存を取り除くことで並列トレーニングを可能にするLSTM/GRUの再検討を行い、minLSTMとminGRUを導入し、それらがさまざまなタスクで最近のシーケンスモデルと同等以上で、学習速度を大幅に改善することを示す。

ABSTRACT

The introduction of Transformers in 2017 reshaped the landscape of deep learning. Originally proposed for sequence modelling, Transformers have since achieved widespread success across various domains. However, the scalability limitations of Transformers - particularly with respect to sequence length - have sparked renewed interest in novel recurrent models that are parallelizable during training, offer comparable performance, and scale more effectively. In this work, we revisit sequence modelling from a historical perspective, focusing on Recurrent Neural Networks (RNNs), which dominated the field for two decades before the rise of Transformers. Specifically, we examine LSTMs (1997) and GRUs (2014). We demonstrate that by simplifying these models, we can derive minimal versions (minLSTMs and minGRUs) that (1) use fewer parameters than their traditional counterparts, (2) are fully parallelizable during training, and (3) achieve surprisingly competitive performance on a range of tasks, rivalling recent models including Transformers.

研究の動機と目的

• 並列可能なトレーニングの文脈において、従来の RNN（LSTM/GRU）の再評価。
• 隠れ状態への依存を取り除く最小限でパラメータ効率の良い variants を開発。
• minLSTMとminGRUが大幅な速度向上を伴い並列トレーニング可能で、タスクを横断して現代のシーケンスモデルと同等の性能を発揮することを示す。

提案手法

• LSTM/GRUのゲートを、parallel scan 形式に適合させるために h_{t-1} への依存を落とすように書き換える。
• 出力/状態の範囲制約を tanh ベースから取り除き、出力を安定化・正規化する。
• 並列プリフィックススキャンアルゴリズムによる、はるかに少ないパラメータで並列トレーニングが可能な minGRU および minLSTM の導出。
• 複数のタスク（合成、RL、言語モデリング）を横断して、minGRU/minLSTMをGRU/LSTMおよび最近のモデル（Mamba など）と経験的に比較する。

実験結果

リサーチクエスチョン

• RQ1 古典的な LSTM/GRU アーキテクチャは、時間をさかのぼる誤差伝播なしで並列トレーニングを可能に再定式化できるか？
• RQ2 最小化された variant（minGRU/minLSTM）は、パラメータ数を大幅に減らした状態で、並列実行とともに Transformer や最先端の再帰モデルと競争力のある性能を達成するか？
• RQ3 隠れ状態依存と出力範囲制約を取り除く際の速度、メモリ、および安定性のトレードオフは何か？
• RQ4 最小限のRNNは、現代のシーケンスモデルを評価するタスク（Selective Copying、D4RL RLタスク、言語モデリング）にスケールしますか？

主な発見

• minGRU と minLSTM は parallel scan アルゴリズムを用いて並列トレーニングを可能にし、報告された実験でシーケンス長 512 に対して従来のRNNに比べ約175×、LSTM（長さ512）に対して約235×の速度向上を達成。
• 最小モデルはパラメータ数が大幅に少なく（例：minGRU ~ O(2 d_h d_x) 対 GRU ~ O(3 d_h (d_x + d_h)); minLSTM ~ O(3 d_h d_x) 対 LSTM ~ O(4 d_h (d_x + d_h))）。
• トレーニング/ランタイム比較では、minGRU/minLSTM は Mamba と同等の実行時間を示し、従来のRNNより劇的に速い。シーケンス長512での実行時間は：2.97 ms (minLSTM)、2.72 ms (minGRU)、2.71 ms (Mamba)。
• 長いシーケンス長4096の場合も、minGRUとminLSTMは大きな速度向上を維持（GRU/LSTM対比、各自の最小モデルと比べて1324×および1361×高速化）。
• Selective CopyingとD4RL RL ベンチマークでは、minGRU/minLSTMはS4、Hyena、Transformerベースのベースラインと競合または上回る水準でタスクを解き、いくつかのデータセットで平均性能がDecision S4を上回る。
• Shakespeareの言語モデリングでは、minGRU/minLSTMはテストロスがMambaおよびTransformersに近づき、同等の性能を達成するにはTransformersがかなり多くの学習ステップを要する（約2.5倍のステップ）