FlashAttentionはデータ移動を減らすことで正確なアテンションを高速化した

長い系列を扱うTransformerは、標準的な自己アテンションが大きな中間行列を実体化し、GPUの高帯域メモリと演算ユニットの間で繰り返し移動させていたため、低速でメモリを大量に消費していた。

魅力的な修正案は近似アテンションだった。モデルの数学を変えて計算量を減らすのである。FlashAttentionは逆の道を取った。正確なアテンションを保ちながら、ブロックをSRAMで処理し、実行中のソフトマックス統計を再計算することで、完全なアテンション行列を保存する必要をなくした。

論文は、MLPerf 1.1の学習記録と比較してBERT-largeでエンドツーエンドの実時間で15%の高速化、系列長1KのGPT-2で3倍の高速化、Long Range Arenaのタスクで2.4倍の高速化を報告した。

最適化すべきは数学ではなくメモリトラフィック

核心的な制約は浮動小数点演算だけでなく、HBMとSRAMの間のIOだった。

タイル処理により、データがすでにチップ上にある間に各ブロックが有用な仕事を行えた。

アルゴリズムを正確なまま保つことで、多くの近似手法を魅力のないものにしていた品質の劣化を避けられた。

同じプリミティブが、Transformerの抽象化を変えることなく、より長い文脈と新しい系列タスクを可能にした。

ハードウェア上のデータ移動がボトルネックである場合、より優れたアルゴリズムとは、同じ答えを計算しながらメモリ階層をより少ない回数しか横断しないものでありうる。

FlashAttentionは、効率的なTransformerの学習と推論のための標準的な構成要素となった。そのより大きな教訓はアテンションの枠を超えて広がった。現代のアクセラレータ向けアルゴリズムは、新しいモデルの数学を発明する前に、メモリ階層を中心に設計することで勝つことが多い。