【ITニュース解説】FlashAttention by hand

自己注意機構は、今日の多くのAIモデル、特にTransformerにおいて中心的な役割を果たす。この仕組みは、入力された情報同士の関係性を数値化し、どの情報が他のどの情報と強く関連しているかを学習することで、文脈を理解する能力を高める。具体的には、入力された各情報からクエリ（Query）、キー（Key）、バリュー（Value）という3種類のベクトルを生成し、これらを使って計算を進める。

従来の自己注意機構の計算手順は、大きく分けて三つの段階を踏む。まず、クエリとキーの類似度を計算し、ログットと呼ばれるスコアの行列を生成する。次に、このログット行列に対してソフトマックス関数を適用し、すべてを合計すると1になるようなアテンションスコアの行列を算出する。最後に、このアテンションスコア行列をバリュー行列と掛け合わせることで、最終的な出力ベクトルを得る。

この従来の計算方式には、特に長い入力系列を処理しようとすると課題があった。それは、途中で生成されるアテンションスコア行列が、入力情報の長さの二乗に比例して巨大になるためである。この巨大な行列を計算し、GPUのメインメモリであるグローバルメモリ（DRAM）に書き出し、再度読み込むという処理が頻繁に発生すると、データの転送速度が計算速度のボトルネックとなり、全体的な処理が遅くなるだけでなく、大量のメモリを消費してしまう。

FlashAttentionは、このメモリ転送のボトルネックを解消するために開発された革新的なアルゴリズムである。その主要な目的は、自己注意機構の計算において、途中で発生する巨大なアテンションスコア行列をグローバルメモリに書き出すことを完全に回避し、最終的な出力を直接計算することにある。これにより、メモリのアクセス回数を劇的に減らし、計算速度とメモリ効率を大幅に向上させる。

FlashAttentionがこの目標を達成する鍵は、「融合されたタイル化されたアテンション」というアプローチにある。これは、計算を細かく分割し、GPUの高速なオンチップメモリであるSRAM（スタティックRAM）を最大限に活用するという考え方だ。具体的には、クエリ、キー、バリューの各行列を小さなブロック（「タイル」と呼ぶ）に分割する。そして、出力行列の各行を計算する際、対応するキーとバリューのタイルをSRAMに順に読み込み、限られたデータだけをSRAM内で処理していく。

この処理の中心にあるのは、いわゆる「オンラインソフトマックス」の原理である。従来のソフトマックス計算では、すべてのログット値を一度に見てから最大値を特定し、指数化と正規化を行う必要があった。しかし、オンラインソフトマックスは、データを小さなブロックに分けて逐次的に処理しながら、それまでの最大値とソフトマックスの分母をリアルタイムで更新していく。FlashAttentionは、このオンラインソフトマックスの考え方を応用し、バリューとの積の計算も統合する。

処理の流れとして、まず最終的な出力ベクトルを保持するための初期値、これまでの最大ログット値、ソフトマックスの分母をSRAM内に用意する。次に、入力となるキーとバリューの各タイルをグローバルメモリからSRAMへ読み込む。SRAMにロードされたデータに対して、ログット値の計算、そのタイル内での最大値の特定、そしてソフトマックスの計算の一部を行う。ここで得られた「ローカルな」統計値を、SRAM内に保持している「グローバルな」統計値に統合し、更新する。

この更新の際に重要なのが、新しいタイルにこれまでの最大値を上回るログット値が含まれている場合、ソフトマックスの計算基準が変わるため、それまでの出力ベクトルと分母も新しい基準に合わせて調整する「再スケーリング」である。FlashAttentionはこの再スケーリングを効率的に行い、各タイルの処理が終わるたびに、現在のところ処理されたすべての情報に基づいた出力ベクトルと分母を維持する。

この一連の処理を、すべてのタイルに対して繰り返す。次のキーとバリューのタイルをSRAMにロードし、前回のタイルで計算されたグローバルな統計値を使って同様の計算と更新を行う。全てのタイル処理が完了した後、最終的に得られた出力ベクトルを最終的な分母で一度だけ割ることで、完全に正規化された自己注意機構の出力が得られる。

この方法の最大の利点は、グローバルメモリへのアクセスが大幅に削減されることにある。FlashAttentionでは、中間のアテンションスコア行列自体はグローバルメモリには一度も書き出されず、SRAM内で一時的に計算されてすぐにバリューとの積に利用されるため、メモリの消費も抑えられる。

FlashAttentionのアルゴリズムは、入力埋め込みからクエリ、キー、バリューの活性化行列を生成する前段階の計算には手を加えていない。これらは学習可能な重み行列と入力埋め込みの積によって生成されるものであり、FlashAttentionが最適化するのは、これらの活性化行列が与えられた後の、ソフトマックスとバリューの積の計算フェーズである。また、記事内の手計算例では、説明を簡潔にするため、最終的な正規化を一度に行う方式を採用しているが、これは逐次的に正規化を行う論文の方式と数学的に等価であり、最終結果は変わらない。

まとめると、FlashAttentionは、計算の融合、データのタイル化、そしてオンラインソフトマックスの応用によって、従来の自己注意機構が抱えていたメモリのボトルネックと計算効率の課題を解決する画期的な手法である。これにより、特に大規模なTransformerモデルや長い系列データを扱う際の、学習速度と推論速度の大幅な向上が実現され、AI技術のさらなる発展に貢献している。