[論文レビュー] Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention

Angelos Katharopoulos, Apoorv Vyas+2 more|arXiv (Cornell University)|Jun 29, 2020

Neural Networks and Applications参考文献 35被引用数 338

ひとこと要約

本論文は softmax アテンションをカーネルベースの線形アテンションに置換する線形トランスフォーマを導入し、O(N) 時間/メモリと自回帰推論を数千倍高速化しつつ、従来のトランスフォーマと同等の性能を維持する。

ABSTRACT

Transformers achieve remarkable performance in several tasks but due to their quadratic complexity, with respect to the input's length, they are prohibitively slow for very long sequences. To address this limitation, we express the self-attention as a linear dot-product of kernel feature maps and make use of the associativity property of matrix products to reduce the complexity from $\\mathcal{O}\\left(N^2\ ight)$ to $\\mathcal{O}\\left(N\ ight)$, where $N$ is the sequence length. We show that this formulation permits an iterative implementation that dramatically accelerates autoregressive transformers and reveals their relationship to recurrent neural networks. Our linear transformers achieve similar performance to vanilla transformers and they are up to 4000x faster on autoregressive prediction of very long sequences.

研究の動機と目的

• 長いシーケンスに対するトランスフォーマの自己注意の二次計算コストを動機づけ、解決する。
• 複雑さを削減するためにカーネル特徴写像を用いた線形アテンションの定式化を提案する。
• 線形アテンションを用いた因果マスキングを示し、効率的な自回帰推論のためのRNN様の再帰を導出する。
• ベースラインと比較して、画像生成と音声認識で経験的に検証する。

提案手法

• 自己注意を phi という特徴写像を用いたカーネル化ドット積として再定式化し、V' = (phi(Q) (phi(K)^T V)) / (phi(Q)^T sum_j phi(K_j)) のように表す。
• 結合律を用いて全てのクエリに対して和 S = sum_j phi(K_j) V_j^T と Z = sum_j phi(K_j) を計算し、フォワードパスを O(N) 時間にする。
• 自己回帰生成のために S_i および Z_i の増分更新を導入して因果マスキングを実現し、線形時間と一定メモリを達成する。
• 因果マスキングを備えたトランスフォーマ層のRNN様ビューを導出し、注意メモリ s と正規化メモリ z、そして2段階更新を特徴とする。

実験結果

リサーチクエスチョン

• RQ1 トランスフォーマの自己注意を、性能を犠牲にせずに線形（O(N)）の時間とメモリで再定式化できるか。
• RQ2 正の類似性関数を用いたカーネルベースの線形アテンションは、softmaxアテンションと同等の自回帰推論を実現できるか。
• RQ3 因果マスキングを備えたトランスフォーマ層はRNNとして解釈でき、効率的な逐次生成を可能にするか。
• RQ4 線形アテンションを用いた長いシーケンスのタスク（画像生成、音声認識）における精度と速度の経験的トレードオフは、標準のトランスフォーマーおよび Reformer と比べてどうか。
• RQ5 自回帰タスクにおける訓練と推論の条件で、線形アテンションの性能はどうなるか。

主な発見

• 線形トランスフォーマは、長いシーケンスの自回帰予測で最大4000倍高速でありつつ、ベースのトランスフォーマと同等の性能を達成する。
• 線形アテンションはメモリと時間がシーケンス長に対して線形にスケーリングし、softmaxアテンションの二乗スケーリングとは異なる。
• MNIST、CIFAR-10、WSJ の実験で、線形トランスフォーマは競争力のある bits/dim または PER を達成しつつ、はるかに高いスループットと低いメモリを示す。
• 線形アテンションによる自回帰推論は、内部状態（S_i, Z_i）を保持・更新することで、1ステップあたりのコストを一定にできる。
• 本研究は、因果マスキングを備えたトランスフォーマとRNNの間の理論的な関連を、注意メモリのRNN様再帰によって確立する。