- FlashAttentionは、GPU内のメモリの読み書きを減らし、注意計算を速くする方法を提案した。
- 原論文は厳密な注意計算を保つ方式と、別の近似的な疎な拡張方式を区別して評価した。
- 特定の機器と系列長での高速化を、すべてのAI処理の同じ倍率の改善と読み替えない。
何を確かめた研究か
Transformerが長い文章を読むとき、各位置がほかの位置を参照する「注意計算」は、文章の長さが増えるほど重くなる。計算式の演算数だけを減らしても、GPU内の遅いメモリと速いメモリの間を何度も往復すれば、実際の時間は短くならない。FlashAttentionの原論文は、メモリ移動を減らす設計で注意計算を速くできるかを問うた。
ここでいう「厳密」は、近似で一部の注意関係を捨てず、元の注意計算と同じ結果を得ることを指す。論文には、ブロック単位で疎にする近似的な拡張も登場する。両者を同じ方法として扱わない。
どう検証したか
FlashAttentionは、注意行列全体を遅い高帯域メモリに置いて何度も読み書きする代わりに、計算を小さな区画に分け、GPUチップ上の速いメモリを使う。必要な中間値は一部を再計算し、大きな中間行列の保存を避ける。これは計算を単純に減らすというより、読み書きの費用と再計算を交換する設計だ。
論文は、注意計算単体の時間・メモリ使用量と、モデル全体を学習する時間を分けて調べた。BERT-large、GPT-2、長い系列を扱う公開課題などで、当時の標準的な実装や近似手法と比較する。比較対象、GPU、系列長、学習設定が異なれば速度比も変わる。
論文は演算の回数だけでなく、高帯域メモリへの読み書き回数をモデル化している。速いメモリに収まる区画の大きさを決め、前向き計算と勾配計算の両方で大きな行列を保存しない構成を検討した。
何が分かったか
要約では、Hugging Face実装との比較で系列長1024のGPT-2学習が約3倍、系列長512のBERT-large学習で約15%の総所要時間短縮を報告する。後者の比較対象はMLPerf 1.1の当時の学習速度記録である。同じ方法でも、モデル全体で注意計算が占める割合が違えば、総時間の改善は異なる。
注意計算単体では、長い系列でメモリ使用量を抑え、より長い入力を扱える条件を示した。ただし、厳密なFlashAttentionの演算量は系列長に対して依然として二次的に増える。メモリ使用量が直線的に増えることと、演算時間も常に直線的になることは別だ。
限界は何か
論文第5節は、新しい注意方式ごとに低水準のCUDA実装が必要で、開発に手間がかかると述べる。特定のGPU向けの実装が別の構成にそのまま移るとも限らない。複数GPU間の通信を含む最適化も、同論文の中心的な実証範囲の外にある。
高速化の数字は、当時の機器、ソフトウェア、比較実装に依存する。また、長い入力を処理できるようになっても、モデルが離れた情報を正しく利用できることまで保証しない。長い文書の理解や実務上の正確さは、別途評価する必要がある。
仕事との接点
長文処理の費用を比べるなら、入力長ごとの総所要時間、GPUメモリ、出力の正しさを同じ条件で測る。実装の互換性と保守の費用も確認したい。これはTRAIL編集部の提案であり、原論文が現在のすべての製品や環境での速度改善を保証したわけではない。