[論文レビュー] Gated Linear Attention Transformers with Hardware-Efficient Training

Yang Song-lin, Bailin Wang+3 more|arXiv (Cornell University)|Dec 11, 2023

Neural Networks and Applications被引用数 5

ひとこと要約

本稿では、標準的なソフトマックスアテンションをデータに依存するゲート付き線形アテンション機構に置き換えることで、ハードウェアに最適化された変種であるゲート付き線形アテンション（GLA）を提案する。メモリアクセスのパターンを最適化し、I/Oに配慮した新規アルゴリズムであるFLASHLINEARATTENTIONを導入することで、短いシーケンス（例：1K）であってもFlashAttention-2を上回る高速な学習を達成する一方で、線形時間の推論を維持する。GLA-Transformerは、言語モデリングにおいてLLaMA、RetNet、Mambaといった強力なベースラインと同等またはそれを上回る性能を発揮し、特に長さ一般化能力とリ콜中心のタスクで優れた結果を示す。

ABSTRACT

Transformers with linear attention allow for efficient parallel training but can simultaneously be formulated as an RNN with 2D (matrix-valued) hidden states, thus enjoying linear-time inference complexity. However, linear attention generally underperforms ordinary softmax attention. Moreover, current implementations of linear attention lack I/O-awareness and are thus slower than highly optimized implementations of softmax attention. This work describes a hardware-efficient algorithm for linear attention that trades off memory movement against parallelizability. The resulting implementation, dubbed FLASHLINEARATTENTION, is faster than FLASHATTENTION-2 (Dao, 2023) as a standalone layer even on short sequence lengths (e.g., 1K). We then generalize this algorithm to a more expressive variant of linear attention with data-dependent gates. When used as a replacement for the standard attention layer in Transformers, the resulting gated linear attention (GLA) Transformer is found to perform competitively against the LLaMA-architecture Transformer (Touvron et al., 2023) as well recent linear-time-inference baselines such as RetNet (Sun et al., 2023a) and Mamba (Gu & Dao, 2023) on moderate-scale language modeling experiments. GLA Transformer is especially effective at length generalization, enabling a model trained on 2K to generalize to sequences longer than 20K without significant perplexity degradations. For training speed, the GLA Transformer has higher throughput than a similarly-sized Mamba model.

研究の動機と目的

• トランスフォーマーにおける線形アテンションと標準的なソフトマックスアテンションの性能差を是正すること。
• I/Oに配慮したハードウェア最適化アルゴリズムを設計することで、線形アテンションの学習効率を向上させること。
• データに依存するゲーティングを導入することで線形アテンションの表現力を高め、長文脈およびリコール中心のタスクでの性能を向上させること。
• 線形時間の推論複雑度を維持しつつ、標準的なトランスフォーマーと同等の性能を達成すること。
• 訓練長さを超えたより長いシーケンスへの一般化性能を強化し、性能の低下を防ぐこと。

提案手法

• 現代のGPU向けにメモリアクセスと並列性を最適化した、ハードウェアに最適化された線形アテンションのためのアルゴリズムであるFLASHLINEARATTENTIONを提案する。
• I/O効率と学習速度のバランスを取るために、チャンク間の再帰とチャンク内での並列計算を組み合わせたチャンク別並列学習形式を導入する。
• 隠れ状態の更新がデータに依存するゲートによって調整されるゲート付き線形アテンション機構を設計し、モデルの表現力の向上を図る。
• 累積和のテクニックを用いて学習可能なパラメータαとβの閉形式勾配を導出することで、高帯域幅メモリに中間状態を保存する必要を回避する。
• FLASHLINEARATTENTIONアルゴリズムをゲート付きバージョンに一般化し、GLA-Transformerの効率的学習を可能にする。
• RNNに類似した再帰構造を維持しながら、チャンク化による並列学習を可能にする修正されたアテンション計算を採用する。

実験結果

リサーチクエスチョン

• RQ1 ハードウェアに最適化された線形アテンションの実装が、短いシーケンス（例：1Kトークン）においても、高度にチューニングされたソフトマックスアテンション実装（例：FlashAttention-2）を上回ることができるか？
• RQ2 線形アテンションにデータに依存するゲーティングを導入することで、固定減衰係数や通常の線形アテンションと比較して性能が顕著に向上するか？
• RQ3 GLA-Transformerが標準的なトランスフォーマー（例：LLaMA）と同等の性能を発揮しつつ、線形時間の推論複雑度を維持できるか？
• RQ4 GLA-Transformerは、訓練長さよりも長いシーケンスに一般化できるか。特にリコール中心のタスクにおいてはいかがな性能を示すか？
• RQ5 GLA-Transformerの学習スループットは、Mamba や RetNet といった最先端の線形時間モデルと比較してどの程度か？

主な発見

• FLASHLINEARATTENTIONは、短いシーケンス（例：1Kトークン）においても、FlashAttention-2を上回る速度を示し、優れたI/O効率を実証した。
• GLA-Transformerは、言語モデリングベンチマークで競争力のある性能を発揮し、LLaMAアーキテクチャのモデルや最近の線形時間モデル（例：RetNetやMamba）を上回る結果を得た。
• 340Mおよび1.3Bパラメータのモデルにおいて、11のゼロショットタスク全体で平均正解率48.0%および55.5%を達成し、いくつかのベンチマークでMamba や RetNet を上回った。
• GLA-Transformerは20Kトークンを超える長さのシーケンスに対しても効果的に一般化でき、2Kシーケンスで訓練した場合でもパープレキシティの低下が最小限に抑えられた。
• 同サイズのMambaモデルと比較して、GLA-Transformerはより高い学習スループットを達成しており、大規模事前学習におけるスケーラビリティの優位性を示した。
• BoolQ や ARC といったリコール中心のタスクにおいても、優れた性能を示しており、記憶保持力および長距離依存関係のモデリング能力の向上が示唆された。