[論文レビュー] Transformer Quality in Linear Time

Weizhe Hua, Zihang Dai+2 more|arXiv (Cornell University)|Feb 21, 2022

Topic Modeling被引用数 45

ひとこと要約

tldr: FLASH は Gated Attention Unit および mixed chunk attention を導入し、Transformer に類似した品質で線形時間のアテンションを実現し、長文脈言語モデリングタスクで大幅な学習スピードアップを達成する。

ABSTRACT

We revisit the design choices in Transformers, and propose methods to address their weaknesses in handling long sequences. First, we propose a simple layer named gated attention unit, which allows the use of a weaker single-head attention with minimal quality loss. We then propose a linear approximation method complementary to this new layer, which is accelerator-friendly and highly competitive in quality. The resulting model, named FLASH, matches the perplexity of improved Transformers over both short (512) and long (8K) context lengths, achieving training speedups of up to 4.9$ imes$ on Wiki-40B and 12.1$ imes$ on PG-19 for auto-regressive language modeling, and 4.8$ imes$ on C4 for masked language modeling.

研究の動機と目的

• Transformers の長い文脈の制限を克服し、スケーラブルで効率的なアテンションを実現する。
• 高品質を維持しつつ弱いアテンションで動作する新しいレイヤー（Gated Attention Unit）を開発する。
• 局所的な二次の要素とグローバルな線形成分を統合する線形時間アテンション近似（mixed chunk attention）を提案する。
• FLASH が Transformer++ 品質に匹敵または接近しつつ、大規模データセット（Wiki-40B、PG-19、C4）での学習スピードアップを実証する。
• 最先端の線形および二次アテンションモデルに対する extensive ablations と比較を提供する。

提案手法

• O = (U ⊙ ĤV)W_o および ĤV = AV を用いた Gated Attention Unit（GAU）を導入し、A はシンプルな単一ヘッドのアテンションを使用できるゲーティング機構から計算される。
• GAU は retrievedV via attention によって U をゲーティングすることで、弱いアテンションでも Transformer 品質に合わせられることを示す。
• 混合チャンクアテンションを開発する：シーケンスを非重複のチャンクに分割し、各チャンク内で局所的な二次アテンションを適用し、チャンク間でグローバルな線形アテンション（因果的または非因果的）を適用し、結果をゲーティングと組み合わせる。
• GAU を FLASH 内に表現し、GAU ブロックと mixed chunk attention を組み合わせて、品質の損失を最小化しつつ線形計算量を達成する。
• 擬似コードを提供し、メモリ再フォーマットを最小化し、実務的に線形スケーリングを維持することに焦点を当てたアクセラレータ対応の実装を議論する。
• C4、Wiki-40B、PG-19 における双方向および自己回帰言語モデリングの大規模実験を、モデルサイズ約 110M–500M パラメータで実施する。

実験結果

リサーチクエスチョン

• RQ1 GAU は弱い単一ヘッドアテンションを用いて Transformer レベルの品質を達成できるか？
• RQ2 混合チャンクアテンションは長いシーケンスで高品質を保ちながら線形時間計算を実現できるか？
• RQ3 FLASH-Quad（二次）と FLASH（線形）は、 MLM および LM タスクにおいて Transformer++ や他の線形アテンション手法と、文脈長の変動に対してどのように比較されるか？
• RQ4 大規模データセットでの文脈長を 512 から 8K にスケールした際のスピードアップと品質のトレードオフはどうなるか？
• RQ5 下流の長文脈タスク（例：TriviaQA）へのファインチューニングで、GAU ベースの FLASH は Transformer++ やベースラインと比較してどう機能するか？

主な発見

• GAU は小さな単一ヘッドアテンションだけで Transformer レベルの品質に匹敵する。
• FLASH は線形スケーラビリティと大幅な学習スピードアップを達成：自己回帰 LM で wiki-40B で最大 4.9×、PG-19 で 12.1×、MLM で C4 で 4.8×。
• FLASH-Quad と FLASH は、Wiki-40B の文脈長 512–8192 に対して低遅延と競争力のあるパープレキシティを提供し、長い文脈でより大きな利得を示す（例：PG-19）。
• Transformer++ と比較して、FLASH は PG-19 で学習コストを最大 12.1×削減し、タスク間で有利な品質と速度のトレードオフを実現する。
• アブレーション研究は、局所的な二次成分とグローバルな線形成分の両方が重要であることを示しており、GAU は混合チャンク変種の MC-TFM++ を大幅に上回る。
• チャンクサイズと設計選択（例：First-to-all トークンオプション）はファインチューニングの性能と速度に影響を与え、一般に長い文脈に対しては大きなチャンクが有利である。