Flash Attention（ふらっしゅ あてんしょん）

概要

Flash Attentionは、2022年にStanford大学の研究者によって提案されたAttention機構の最適化手法です。Transformerモデルは、自然言語処理や画像認識など、様々な分野で高い性能を発揮していますが、Attention機構の計算コストがボトルネックとなる場合があります。Flash Attentionは、この問題を解決するために、Attention行列をブロックに分割し、GPUの高速なSRAMメモリを活用することで、計算効率を大幅に向上させます。

技術的な詳細

従来のAttention機構では、Attention行列全体をGPUの低速なHBMメモリに保存する必要がありました。Flash Attentionでは、Attention行列を小さなブロックに分割し、SRAMメモリにキャッシュすることで、メモリアクセス回数を削減します。また、ブロックごとの計算を並列化することで、計算速度を向上させます。具体的には、以下の技術が用いられています。

• Tiling: Attention行列をブロックに分割する。
• Recomputation: 順伝播時に必要な中間結果を再計算することで、メモリ使用量を削減する。
• Fusion: 複数の演算を融合することで、GPUの利用効率を向上させる。

性能

Flash Attentionは、従来のAttention機構と比較して、メモリ使用量を大幅に削減し、計算速度を向上させることができます。例えば、大規模言語モデルの学習において、Flash Attentionを使用することで、学習時間を数倍短縮することが可能です。また、Flash Attentionは、GPUのメモリ容量が限られている環境でも、大規模なモデルを学習することができます。

応用

Flash Attentionは、大規模言語モデルの学習だけでなく、様々な分野に応用することができます。例えば、画像認識、音声認識、機械翻訳など、Transformerモデルが使用されるあらゆるタスクにおいて、Flash Attentionを使用することで、性能を向上させることができます。

参考文献

• Dao, Tri, et al. “FlashAttention: Fast and Efficient Transformer Attention with IO-Awareness.” Advances in Neural Information Processing Systems 35 (2022).