FlashAttention

FlashAttention とは、Transformer の中心であるアテンションの計算を、GPU のメモリ階層をうまく使って速く、省メモリで行う実装手法です。2022年に発表された研究が元になっており、アテンションの計算結果そのものは変えずに、計算のやり方だけを改善している点が大きな特徴です。

なぜ必要か

アテンションは、入力の全トークンどうしの関連度を計算します。入力が長くなると、その組み合わせは文の長さの2乗で増えます。素朴な実装では、巨大な関連度の表を GPU の大きなメモリ（HBM）に書き出し、また読み戻す作業が発生します。実は、時間がかかる主な原因は計算そのものより、このメモリの読み書きです。

GPU には、容量は小さいものの非常に高速なメモリ（SRAM）もあります。大きなメモリとの往復を減らして、小さく速いメモリの中で処理を終わらせれば、速くなります。

しくみ

• 入力を小さなブロックに分け、ブロックごとに計算して高速なメモリの中で完結させる（タイリング）
• ソフトマックスという正規化の計算を、全体を一度見なくても、ブロックごとに逐次更新できる形にする
• 巨大な関連度の表を、大きなメモリに保存しない

結果として、メモリの使用量が文の長さに比例する程度に抑えられ、速度も向上します。これは近似ではなく、計算を並べ替えているだけなので、結果は通常の計算と（数値の丸め誤差を除いて）一致します。

具体例

モデルの学習でも、推論でも使われます。長い文書や長い対話を扱う文脈長の大きいモデルは、こうした仕組みがあることで、現実的な時間とメモリで動かせます。その後も改良版が登場しており、対応する GPU の世代に合わせた最適化が進んでいます。

使うときの注意

• 利用者が直接触る場面は少なく、多くの推論ライブラリや学習フレームワークの内部で、設定により有効になります。
• 対応する GPU の種類やソフトウェアの版に条件があります。
• 計算量そのものが減るわけではなく、文脈が非常に長い場合には、2乗に増える計算の重さは残ります。
• 高速化できるのは主にアテンション部分で、モデル全体の速度は他の処理にも左右されます。

関連する考え方

文脈長が伸びた背景には、モデルの工夫だけでなく、こうした計算の実装技術の進歩があります。アテンションそのものの計算量を減らす別のアプローチとして、状態空間モデルのような構造の違うモデルも研究されています。