分散表現学習のためのFlashAttention
著者：Asher Bond (asher.bond@distillative.ai)
要旨
FlashAttention [1] は、メモリトラフィックを大幅に削減しつつ厳密な自己注意 [2] を計算する、IO を意識したアルゴリズムである。これが削減するのは注意機構のメモリおよびメモリ帯域幅のコストであって、演算(FLOP)量ではない。結果は厳密であり、実時間(wall-clock)の高速化は、注意行列の全体を高帯域幅メモリ上に一度も具現化しないことから生じる。本稿は、これが分散表現の学習に何をもたらすのかを明らかにし、カーネル融合(厳密)とブロックスパース版(近似)との境界を——正確に——引く。
解説
標準的な自己注意は、長さ n の系列に対して n × n のスコア行列を形成し、これは時間・メモリのいずれにおいても二次のオーダーとなる [2]。長い系列では、実務上支配的なコストは生の演算そのものではなく、その巨大な中間行列を GPU のメモリ階層間で移動させることであることが多い。FlashAttention [1] が対処するのは、まさにこの点である。計算をタイル化し、スコア・ソフトマックス・値の集約という各段階を融合することで、n × n の行列全体を高帯域幅メモリに一切書き出さず、メモリ使用量を n について二次から線形へと削減しつつ、通常の注意と同一の結果を計算する。同論文は、実測された実時間の高速化と、素朴な実装では不可能な、より長い文脈での学習を報告している。
分散表現学習にとって、その帰結は直接的である。固定のメモリ予算のもとでより長い文脈が現実的な負担で扱えるようになり、注意パターンを近似することなく、より広い範囲——文書全体、長時間の音声——から表現を学習できる。これこそが主張すべき利点であり、その根拠は [1] にある。
この論を支える区別が二つあり、それらは正確に述べておく価値がある。
- 厳密であり、FLOP において劣二次ではない。 基本の FlashAttention は、厳密な注意の漸近的な演算コストを保つ。低減するのはメモリトラフィックであり、その結果として実時間である [1]。劣二次の挙動が生じるのは、注意パターンを変更する別個のブロックスパース版に限られる——これは近似であり、近似として明示されている。
- カーネル融合とブロックスパース性は別物である。 融合(行列の具現化を回避すること)は厳密性を保つ。ブロックスパース性は、さらなる削減と引き換えに厳密性を手放す。両者は異なる設計点であり、混同して扱ってはならない。
根拠と適用範囲
本稿における定量的な言明はすべて、Transformer をバックボーンとする [2] FlashAttention 論文 [1] の実測結果である。すなわち、厳密な FlashAttention は注意の二次の FLOP 量を保持するメモリ/IO の最適化であり、その高速化はメモリトラフィックの低減に由来する実時間の高速化である。劣二次の計算量そのものが要件となる場面では、設計点は近似注意の手法——スパース系あるいは低ランク系——へと移り、これらはスケールと引き換えに厳密性を手放す。本稿は、その地形の地図であり、両者の境界に沿って描かれている。
参考文献
- Tri Dao et al. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. Advances in Neural Information Processing Systems (NeurIPS). arXiv:2205.14135.
- Ashish Vaswani et al. (2017). Attention Is All You Need. Advances in Neural Information Processing Systems (NeurIPS). arXiv:1706.03762.