[論文レビュー] Attention as an RNN

Leo Feng, Frederick Tung+4 more|arXiv (Cornell University)|May 22, 2024

Educational and Psychological AssessmentsPsychology被引用数 3

ひとこと要約

本稿では、自己注意機構を微分可能なRNNに再定式化することで、変換器と同様の並列学習と、段階的トークン更新による効率的で定数メモリの推論を可能にする、Aarenと呼ばれる新しい注目メカニズムを提案する。Aarenは38の順序付きモデリングタスクにおいて変換器水準の性能を達成するとともに、計算およびメモリの複雑さを二次的から一次的に削減する。

ABSTRACT

The advent of Transformers marked a significant breakthrough in sequence modelling, providing a highly performant architecture capable of leveraging GPU parallelism. However, Transformers are computationally expensive at inference time, limiting their applications, particularly in low-resource settings (e.g., mobile and embedded devices). Addressing this, we (1) begin by showing that attention can be viewed as a special Recurrent Neural Network (RNN) with the ability to compute its extit{many-to-one} RNN output efficiently. We then (2) show that popular attention-based models such as Transformers can be viewed as RNN variants. However, unlike traditional RNNs (e.g., LSTMs), these models cannot be updated efficiently with new tokens, an important property in sequence modelling. Tackling this, we (3) introduce a new efficient method of computing attention's extit{many-to-many} RNN output based on the parallel prefix scan algorithm. Building on the new attention formulation, we (4) introduce extbf{Aaren}, an attention-based module that can not only (i) be trained in parallel (like Transformers) but also (ii) be updated efficiently with new tokens, requiring only constant memory for inferences (like traditional RNNs). Empirically, we show Aarens achieve comparable performance to Transformers on $38$ datasets spread across four popular sequential problem settings: reinforcement learning, event forecasting, time series classification, and time series forecasting tasks while being more time and memory-efficient.

研究の動機と目的

• モバイルや組み込みデバイスなどのリソースが限られた環境における変換器の高い推論コストを解消すること。
• 変換器の並列性とRNNの効率的で段階的な推論の間のギャップを埋めること。
• 並列学習と定数メモリ、ストリーミング推論を両立できる新しい注目形式を開発すること。
• 並列プレフィックススキャンを用いて、多くの入力から多くの出力へのRNN出力を効率的に計算する手法を導入すること。
• 得られたAarenモジュールが変換器の性能を再現する一方で、時間的およびメモリ的複雑さを顕著に低減できることを示すこと。

提案手法

• 標準的な自己注意機構を、多くの入力から1つの出力への特別なRNNのケースとして再定式化し、注目計算に再帰的解釈を可能にする。
• 多くの入力から多くの出力へのRNN出力を効率的に計算するため、並列プレフィックススキャンアルゴリズムを用いた新しいRNNベースの注目形式を導入する。
• ソフトマックスベースの注目機構を微分可能かつ正確に再実装するRNNとして定式化し、先行研究の近似手法で見られる指数的タイムスタンプバイアスを回避する。
• 一定の隠れ状態サイズを維持し、新しいトークンの段階的更新をサポートする学習可能なRNNモジュールとしてAarenを設計する。
• 過去の活性化を保存せずに長距離依存関係を捉えるために、学習可能な初期隠れ状態ベクトル $ q $ を採用する。
• 効率的な並列プレフィックス和アルゴリズム（例：Hillis-Steele）を活用し、$ O(N \log N) $ 時間で注目重みを計算することで、スケーラブルな推論を実現する。

実験結果

リサーチクエスチョン

• RQ1 標準的な自己注意機構を形式的に再解釈して再帰的ニューラルネットワークとして捉えることは可能か？
• RQ2 変換器のような注目ベースのモデルは、並列学習が可能であるにもかかわらず、なぜ効率的な段階的推論をサポートできないのか？
• RQ3 変換器の並列性とRNNの定数メモリ効率性を両立できる注目機構を設計することは可能か？
• RQ4 並列プレフィックススキャンに基づく新しい注目形式は、変換器と同等の性能を達成しながら、計算複雑さを低減できるか？
• RQ5 提案されたAarenモジュールは、多様な順序付きタスクにおいて、変換器と比較して推論効率およびモデル性能の点でどのように差をつけるか？

主な発見

• Aarenは強化学習、イベント予測、時系列分類、時系列予測を含む、38の多様な順序付きモデリングデータセットにおいて、変換器と同等の性能を達成する。
• Aarenの累積的推論時間およびメモリ使用量は、系列長に比例して線形に増加するが、KVキャッシュを用いた変換器は二次的に増加する。
• 同等の変換器と比較して、パラメータ数はわずか0.016%増加（3,152,896 対 3,152,384）にとどまるため、効率性の向上はコスト効率が良い。
• 推論中に定数メモリ使用量を維持するため、モバイルや組み込みシステムなどのリソースが限られたデバイスへのデプロイが可能である。
• 並列プレフィックススキャンに基づく注目計算により、微分可能で効率的なRNNスタイルの推論が可能になり、同時に完全な注目機構の表現力も保持される。
• 実験的結果から、Aarenはすべての評価タスクにおいて変換器を上回るか同等の速度およびメモリ効率性を示している。