Flash Attention とは？

基本のアテンションでは、トークン同士の関係を表す大きな中間行列ができるよ。それを保存する実装では、行列の要素数はトークン数の2乗になる。1024から2048へ倍にすると、その行列の要素数は4倍。モデル全体のメモリが必ず4倍になるという意味ではないんだ。

入力を小さなブロックに分け、GPU上の小さく速いSRAMで計算するよ。大容量のHBMへ中間行列を丸ごと書き出さず、メモリ間の読み書きを減らす。HBM自体も高速だけど、SRAMに比べてアクセスの負担が大きいんだ。

基本の密なアテンションは、関係を間引く近似ではなく、同じ数式を計算するよ。ただし浮動小数点では計算順序や精度による丸め差があるので、ビット単位の完全一致を保証するわけではない。公式実装も数値の許容誤差を設けて確認しているんだ。

タイリングでブロックごとの結果を正しく合成し、学習の逆伝播では必要な中間値を再計算するんだ。入力と出力はHBMにあり、ブロックを繰り返し読み、結果も書き戻す。「一度読むだけ」ではないよ。中間行列を保存しないことで追加メモリを長さに対して線形に抑えるけど、基本の計算量は2乗のままなんだ。

速度の倍率はGPUや入力の長さ、比較対象、実装によって変わるよ。アテンション単体の高速化とモデル全体の学習速度も別。公式実装にはGPUやデータ精度などの対応条件があり、非公開モデルも含めて全部が使っているとは断定できないんだ。