[論文レビュー] Block-Recurrent Transformers

DeLesley Hutchins, Imanol Schlag+3 more|arXiv (Cornell University)|Mar 11, 2022

Topic Modeling被引用数 27

ひとこと要約

The Block-Recurrent Transformer processes long sequences by applying a transformer layer recurrently over blocks, achieving linear complexity in sequence length and substantially better perplexity on long documents with similar or lower compute than non-recurrent baselines.

ABSTRACT

We introduce the Block-Recurrent Transformer, which applies a transformer layer in a recurrent fashion along a sequence, and has linear complexity with respect to sequence length. Our recurrent cell operates on blocks of tokens rather than single tokens during training, and leverages parallel computation within a block in order to make efficient use of accelerator hardware. The cell itself is strikingly simple. It is merely a transformer layer: it uses self-attention and cross-attention to efficiently compute a recurrent function over a large set of state vectors and tokens. Our design was inspired in part by LSTM cells, and it uses LSTM-style gates, but it scales the typical LSTM cell up by several orders of magnitude. Our implementation of recurrence has the same cost in both computation time and parameter count as a conventional transformer layer, but offers dramatically improved perplexity in language modeling tasks over very long sequences. Our model out-performs a long-range Transformer XL baseline by a wide margin, while running twice as fast. We demonstrate its effectiveness on PG19 (books), arXiv papers, and GitHub source code. Our code has been released as open source.

研究の動機と目的

• 長距離メモリを二次的なアテンションコストなしに保持するトランスフォーマーアーキテクチャを動機付け、開発する。
• トークンと状態ベクトルのブロック上で動作するリカレントセルを導入し、長い文脈モデル化を可能にする。
• 標準的なトランスフォーマーと計算パラティを維持しつつ、長いシーケンスの言語モデリングタスクで顕著な perplexity の改善を達成する。

提案手法

• スライディングウィンドウ（因果的）アテンションとブロック単位のリカレンスを使用し、長さ N のセグメントを W のブロックで処理する。
• 垂直方向（自己注意 + リカレント状態へのクロスアテンション）と水平方向（状態から状態への自己注意と入力へのクロスアテンション）を組み合わせたリカレントセルを実装し、情報を忘却/保持するゲートを導入する。
• 垂直方向と水平方向のキーと値を共有しつつ、クエリは別々のものを保持する。
• 学習を長いシーケンスで安定させるための学習可能な状態IDと相対位置バイアスを導入する。
• ゲート構成（固定ゲート、LSTM風ゲート）を試し、安定性、スケーリング、アブレーションを評価する。

実験結果

リサーチクエスチョン

• RQ1 ブロックごとのリカレンスを持つトランスフォーマーは、Transformer-XL ベースラインと同等以下の計算量で、非常に長いコンテキスト（例: 60k トークン）の記憶を得られるか。
• RQ2 リカレンスゲートの型と状態サイズが perplexity、学習安定性、計算コストに与える影響はどのようか。
• RQ3 単一のリカレント層と複数の非リカレント層を長い文書データセットで比較した場合の性能とコストはどうなるか。

主な発見

Model	segment length	window length	step time	PG19 (bytes/tokens)	arXiv (tokens)	GitHub (tokens)
XL:512	512	512	0.88	1.01	3.62±0.01	1.45	1.21
XL:1024	1024	1024	1.20	0.997	3.59±0.01	1.37	1.08
XL:2048	2048	2048	2.11	0.990	3.58±0.01	1.31	1.01
Slide:12L	4096	512	0.93	0.989	3.60	1.43	1.19
Slide:13L	-	-	1.00	0.989	3.58±0.01	1.42	1.17
Rec:lstm:dual	4096	512	1.06	0.985	3.54±0.01	1.26	1.01
Rec:lstm:single	-	-	1.05	0.962	3.54±0.01	1.29	1.03
Rec:lstm:skip	-	-	1.00	0.969	3.56±0.01	1.31	1.10
Rec:fixed:dual	-	-	1.01	0.957	3.52±0.01	1.27	0.991
Rec:fixed:single	-	-	1.02	0.966	3.58±0.01	1.25	1.00
Rec:fixed:skip	-	-	0.99	0.952	3.53±0.01	1.24	0.976
Feedback:lstm:single	4096	512	1.40	0.977	3.50	1.22	-
Feedback:fixed:skip	-	-	1.35	0.935	3.49	1.24	-
Memorizing Trans. 64k	512	512	1.94	0.950	3.53	1.22	-

• 単一のレイヤーによるリカレンスは、同等の計算量で非リカレントレイヤを追加するよりも perplexity の改善が大きい。
• PG19、arXiv、GitHub において、Block-Recurrent バリアントは Transformer-XL ベースラインを上回る、または同程度のステップ時間で、より高速な広範囲コンテキストのベースラインよりも速く動作することがある。
• 最良の設定（Rec:fixed:skip）は他の設定を上回ることが多く、XL:2048 を凌ぐ一方で2倍以上の速さで実行可能。さらに1.3BパラメータのモデルでPG19の語彙レベル perplexity（26.50）という最先端を達成。
• スケーリング結果は、モデルサイズに関与する利益を提供し、パラメータの倍増を超えて利益が拡大する。
• 1層のリカレンスでほとんどの利点を得るのが一般的には十分であり、2層以上のリカレント層は限られた改善かデータセット依存。
• ブロック再帰モデルは、専用 CUDA カーネルなしでベースラインの性能に匹敵・上回ることができ、オープンソース公開にも適している。