MC-Sparse、拡散Transformerの疎注意と密注意の差を縮める

長い系列を扱う拡散Transformerでは、注意計算が推論遅延の大きな要因になる。動画生成や高解像度3Dアセット生成ではトークン数が増えやすく、すべてのクエリとキー・バリューの組み合わせを計算する密注意は高コストだ。そこで疎注意が使われるが、単純に多くのトークンを捨てると、生成結果の品質や細部の忠実度が低下する可能性がある。

MC-Sparseの特徴は、新しい疎パターンを提示するだけでなく、密注意との差がどこから生じるかを検証した点にある。論文は制御されたoracle比較を通じて、主な原因を三つに分けた。第一に、クエリをグループ化すると、本来異なるクエリが同じ選択結果を共有させられる。第二に、重要な相互作用を予測する仕組みが誤り、必要なKVトークンを選び損なう。第三に、削除されたトークンの注意出力への寄与が、そのまま失われる。

提案手法は、次の三つの工夫でこの問題に対応する。

• KVトークンを個別に選択：粗いブロック単位だけに依存せず、正確な注意確率を用いて個々のキー・バリュートークンを選ぶ。これにより、残す相互作用をより細かく制御できる。
• GPU実行を意識したクエリ分割：細粒度の選択は不規則なメモリアクセスを生み、GPUでは非効率になり得る。そこで、類似したクエリをタイルに整列したグループへまとめ、実行効率との両立を図る。
• メタデータと残差を再利用：クエリグループ、選択したKVインデックス、密注意出力と疎注意出力の差分をキャッシュし、後続のデノイズステップで使い回す。選択判断の繰り返しを減らし、破棄された情報の影響も補う設計だ。

MC-Sparseは追加学習を必要としない。狙いは、理論上の計算削減率だけを追うのではなく、実際のGPU実行効率を確保しながら、疎注意の出力を密注意に近づけることにある。報告された実験では、動画および3D生成モデルにおいて、既存の疎注意ベースラインより密注意出力への忠実度とデノイズ高速化の両方で優れた結果を示した。密注意との比較では、Minimax-H3-Baseで1.80倍、3Dアセット生成で2.32倍のデノイズ高速化が報告され、品質低下は無視できる程度とされる。

この研究が示すのは、疎注意の効率化が単なるマスク設計の問題ではないということだ。重要な相互作用を正確に選び、GPUに適した形へ整理し、捨てた情報を補償する三つの要素が必要になる。ただし、提示された素材だけでは、すべてのモデルや疎性設定で同じ効果が得られるとは言えない。系列長、キャッシュのコスト、GPUカーネルの実装などは、実際の速度向上を左右する可能性がある。

出典：Hugging Face Daily Papers