MALA：注意の寄与度に応じて計算量を配分するAttention

長いコンテキストを扱うAttentionでは、QKスコアを計算するコストだけが問題になるわけではない。スコア計算後も、最終的な寄与がほとんどない過去位置に対して、密なカーネルが同じ後処理を続けている。MassAlloc Attention（MALA）は、Attention自身が示す分布を利用して、どの領域に後段の計算を配分するかを決める。

仕組みのポイント

MALAは、スコア計算の前に因果的な接続を固定的に削除する方式ではない。まず、合法な因果範囲にあるQK相互作用へのアクセスを維持する。そのうえで、softmaxによって正規化された寄与を手がかりに、低寄与領域のpost-score計算を減らす。

• 前向き計算：処理中に更新されるオンラインsoftmaxの正規化値を使い、タイルごとの寄与を推定する。
• 逆伝播：最終的な正規化値を再利用し、標準的なAttentionの状態だけから保持領域を導出する。
• 共通の許容誤差：学習と推論で同じ基準を使い、必要な計算量を現在の分布に応じて変化させる。
• 融合プリミティブ：動的な選択処理をカーネルに組み込み、低寄与部分の削減を実際の速度向上につなげることを狙う。

実験から分かること

8Kコンテキストで総post-score作業量を一致させた実験では、MALAが取りこぼした注意量の平均は0.0188%だった。これは、各インスタンスの最終的な分布を知っている理想的な参照質量オラクルの0.0182%に近い。計算量をそろえた比較であるため、単純に多く計算した結果ではなく、分布に応じた配分の精度を検証している。

1Kから32Kまでのコンテキスト長では、同じ許容誤差によってFullAttnに対する出力誤差と勾配誤差を低く保った。関連付け想起の比較では、8K時点のMALAの正解率は89.67%、FullAttnは89.97%だった。少なくとも報告された評価では、後段計算の削減による能力低下は限定的だった。

システム面では、8基のH100とTP=8を使った128KトークンのAttention演算子ベンチマークで、FullAttnに対して学習時の前向きは2.2倍、逆伝播は3.0倍、推論時のデコードは1.6倍高速になった。0.6Bから14Bまでのモデルを用いたスケーリング実験も行われ、14B・32Kコンテキストの学習では総FLOPsを23.1%削減しながら、評価された能力ではFullAttnと同程度の性能を維持した。

意義と注意点

MALAの特徴は、固定された疎なパターンを仮定しないことだ。クエリや入力ごとに異なるAttention分布を見て、重要度に応じて計算量を変える。そのため、長い履歴の大部分が現在の出力にほとんど寄与しない場面では、メモリアクセスだけでなく、スコア後の積和計算や学習時の逆伝播も削減できる可能性がある。

一方、動的な保持判定や制御のオーバーヘッドが大きければ、削減した計算が相殺される。効果はモデル、タスク、コンテキスト長、Attention分布の集中度にも左右されるだろう。MALAは新しいAttention形式というより、既存のAttentionを分布適応型に実行するためのカーネル設計として捉えるのが適切だ。

出典：Hugging Face Daily Papers