[2022]FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness

[takumi7110]

【背景】
• Transformerモデルは、長いシーケンスに対して時間とメモリの複雑さが高く、遅くなりメモリを多く消費する。
• 近似的な注意メカニズムは、計算の複雑さを減らすためにモデルの品質を犠牲にしてきたが、実際の速度向上を達成していない。
• 注意アルゴリズムがIOに対応していないことが問題であり、GPUメモリの読み書きを考慮する必要がある。
【目的】
• FlashAttentionというIOに対応した正確な注意アルゴリズムを提案する。
• FlashAttentionは、GPUの高帯域幅メモリ（HBM）とGPU内部のSRAMの間のメモリ読み書き回数を減らすためにタイリングを使用する。
• FlashAttentionは、既存の近似的な注意手法よりも高速な注意アルゴリズムを提案する。
【手法】
• FlashAttentionは、タイリングを使用してHBMへのメモリ読み書き回数を減らす。
• FlashAttentionは、SRAMにブロックをロードし、各ブロックで注意計算を行い、結果をHBMに書き込む。
• FlashAttentionは、逆伝播時に中間の注意行列をHBMから読み込む代わりに、SRAM上で注意を再計算することで高速化する。
【実験方法】
• FlashAttentionを実装し、BERT-large、GPT-2、および長いシーケンスのモデルトレーニングを行う。
• モデルのトレーニング時間、メモリ使用量、品質の改善などを評価する。
【実験結果】
• FlashAttentionは、BERT-largeで15%、GPT-2で3倍、長いシーケンスで2.4倍のトレーニング速度向上を実現する。
• FlashAttentionは、GPT-2のパープレキシティで0.7の改善と、長いドキュメント分類で6.4ポイントの向上を実現する。
• ブロックスパースなFlashAttentionは、既存の近似的な注意手法よりも高速であり、長いシーケンスでの性能向上を実現する。
【考察】
• FlashAttentionは、注意アルゴリズムのIOに対応することで、モデルのトレーニング時間と品質を向上させることができる。
• FlashAttentionは、他の注意操作にも応用可能であり、さらなる拡張が可能である。
• FlashAttentionの実装はオープンソースで提供されている。

[takumi7110]

PDF:
https://arxiv.org/pdf/2205.14135.pdf