AB-Attention、動画拡散Transformerの注意機構を高精度に疎化

動画拡散Transformerは高品質な映像を生成できますが、フレーム数や解像度が増えるほど推論コストが急激に膨らみます。特に3D全注意機構は、多数の動画トークン間で相互作用を計算するため、二次的な計算量が大きな制約になります。動的なブロックスパース注意機構はこのコストを抑えられる一方、疎化を強くすると画質や数値精度が低下しやすいという課題がありました。

AB-Attentionの狙いは、単に計算ブロックを減らすことではなく、密な注意機構をより忠実に近似することです。従来の手法では、粗い要約値からキー・バリューブロックを評価したり、全トークンを対象にした高コストなクラスタリングを利用したりします。また、タイル単位で保持・削除を判断する場合、タイル内の一部のキーだけが重要でも、タイル全体を残すか捨てるかの二択になりがちです。削除されたブロックはsoftmaxにも参加しないため、注意分布が密な計算からずれる可能性があります。

第1の要素であるIntra-Block Clustering Selector（IBCS）は、各クエリブロックの内部だけでクエリをクラスタリングします。トークンの並べ替えを行わないため、グローバルなクラスタリング構造を作るコストを避けながら、クエリの局所的な違いを利用できます。これにより、単一の粗いブロックスコアよりも、どのキー・バリューブロックがどのクエリ群に関係するかを細かく推定します。

第2の要素はSampling with Inclusion-Probability Weighting（SIP）です。重要度の低いブロックを決定的に切り捨てるのではなく、推定された重要度に応じてブロックをサンプリングし、包含確率の逆数で結果を再重み付けします。つまり、選ばれなかったブロックを常に無関係とみなすのではなく、1回の疎な計算から密な注意機構の推定値を得る設計です。高い疎化率で誤った削除の影響が大きくなる場面に適しています。

さらに、AB-AttentionはGPU実行を意識したカーネルも用意します。標準タイルより小さいキー・バリューブロックを集め、完全なTensor Coreタイルとして処理することで、Tensor Coreの実行形態を保ちながら、より細かな選択を行います。論文によれば、Tensor Coreの計算コストを増やさずに選択粒度を高められる点が特徴です。

Wan2.1とHunyuanVideoを用いた実験では、AB-Attentionは注意計算の89〜92%をスキップしました。比較対象は1.8〜3.6倍多くの注意計算を実行しており、AB-Attentionは最も強いベースラインに対してPSNRを最大4.4 dB改善し、エンドツーエンドでは最大2.23倍の高速化を報告しています。

この研究の意義は、動画拡散の疎な注意機構を、粗い切り捨てから確率的に補正された近似計算へ広げた点にあります。ただし、示された結果はWan2.1とHunyuanVideo、および特定の実験条件に基づくものです。実際の効果はGPU、系列長、目標疎化率、生成タスクによって変わる可能性があり、より多くのモデルでの検証が今後の課題になります。

出典：Hugging Face Daily Papers