技術要約: SparsePR – ビデオ生成およびワールドモデルのためのトレーニングフリーな疎結合アテンション

1. 問題提起

ビデオ生成モデル（例：DiT）やビデオ・ワールドモデルは、長大な時空間トークンシーケンスを処理するため、二次関数的な自己アテンションが主要な推論のボトルネックとなります。トレーニングフリーのブロック疎結合アテンションは加速への道を提供しますが、既存の手法は、疎結合を行単位の特性として扱う際に、以下の2つの根本的な限界に直面しています。

1. 共有ルートのサポートにおける曖昧性（Shared-Route Support Ambiguity）: 行単位のアテンション集中は、複数のクエリが単一のルーティング決定を共有する場合に必要となる実行可能なサポート（support）を特定していません。グループ内のクエリは、キーのサポートが十分に重なっていない場合があり、その結果、プールされたサポート（グループが必要とするキーの和集合）は、個々のサポートの合計よりも大幅に大きくなることがあります。
2. ソフトマックス後の誤差推定ミス（Post-Softmax Error Misestimation）: 保持されたアテンションの質量（mass）だけでは、スキップされた相互作用によって導入される誤差を決定できません。保持された質量が同程度であっても、除外されたサポートと保持されたサポートによって誘発される出力の差によって、2つのルートが生成する誤差は大きく異なる場合があります。

現在の手法は、これらの依存関係を考慮できていないことが多く、その結果、過剰な密度（加速の低下）または重大な品質低下（再構成誤差の増大）を招いています。

2. 手法: SparsePR

著者らは、**応答結合型パーティショニング（Response-Coupled Partitioning）とプローブ適合残差再構成（Probe-Fitted Residual Reconstruction）**を組み合わせた、SparsePRを提案します。この手法は、オフラインのトレーニングを必要とせず、完全にオンラインで動作します。

A. 応答結合型パーティショニング（Response-Coupled Partitioning）

クエリとキーを独立してクラスタリングしたり、交互に洗練を行ったりする従来の手法とは異なり、SparsePRは、現在のاتンテーション・コールから得られる応答幾何学（response geometry）に基づき、単一のアシンメトリックなパスで実行可能なパーティションを構築します。

• キー・応答座標（Key-Response Coordinates）: サンプリングされたクエリ（$Q_s$）は、指標 $M_K = Q_s^\top Q_s / M_s$ を定義します。キーは、これらのサンプリングされたクエリ下での応答プロファイルに基づいてクラスタリングされます（$M_K$ の固有ベクトルに対する $k$-means を使用）。これにより、ペアとなるK/Vグループが形成されます。
• クエリ・応答座標（Query-Response Coordinates）: 得られたK/Vグループの重心は、クエリのための新しい指標空間を定義します。クエリは、これらのK/Vグループに対する事前ソフトマックス・ロジット・プロファイルの類似性に基づいてクラスタリングされます。
• 結果: この結合により、ルートを共有するクエリが同様のサポート要件を持つようになり、セマンティックなクラスタリングで見られる「プールされたサポート」の拡大を最小限に抑えます。

B. プローブ適合残差再構成（Probe-Fitted Residual Reconstruction）

スキップされた相互作用による誤差に対処するため、SparsePRはブロックレベルの近似や時間的再利用には依存しません。

• プローブ選択（Probe Selection）: 少数のクエリ行（$M \ll N_q$）を選択し（クエリグループ間で層化）、それらを厳密に評価することで、高密度な残差（$R = O_{dense} - O_{sparse}$）を計算します。
• アフィン校正（Affine Calibration）: 手法は、残差を疎な出力のアフィン関数（$R \approx b + xB$）としてモデル化します。リッジ回帰によって正則化されたプローブ残差を用いて、コール固有のアフィン写像を適合させます。
• 部分空間制限（Subspace Restriction）: 未観測の出力方向への外挿を防ぐため、補正はプローブ残差の主要な特異ベクトルによって張られる部分空間に投影されます。
• 補正（Correction）: 適合された写像は、未プローブの行に対する残差を推定し、それを疎な出力に加算することで、高密度近似を復元します。

C. 疎な実行（Sparse Execution）

フレームワークは、トークンを連続したグループ・メジャーなレイアウトに並べ替え、FlashInferの可変ブロック疎結合アテンション・プリミティブを使用します。総実行ペア密度（$\rho_{exec}$）には、ルーティングされた疎なセルと厳密なプローブ行の両方が含まれます。

3. 主な貢献

1. 構造的分析: 本論文は、実行可能なブロック疎結合アテンションにおける3つの重要な依存関係を特定しました。
   • クエリごとの疎性は、共有ルートのサポートを決定しない（プールされたサポートは3〜4倍高くなる可能性がある）。
   • 保持されたアテンションの質量は、ソフトマックス後の誤差を保証しない。
   • パーティションの幾何学が、残差の「アフィン説明可能（affine-explainable）」な割合に大きく影響する。
2. SparsePRフレームワーク: 応答結合型パーティショニング（サポート拡大を最小化するため）と、プローブ適合残差再構成（スキップされた相互作用を補正するため）を組み合わせた、トレーニングフリーの手法を導入。
3. 包括的な評価: 4つのヘテロジニアスなモデル（HunyuanVideo, Wan2.2, Cosmos-Predict2.5, Cosmos3-Nano）に対して、厳密なペア計数およびエンドツーエンドのタイミングを用いて評価を行い、一貫した誤差減少と高速化を実証。

4. 実験結果

著者らは、テキスト・トゥ・ビデオ、イメージ・トゥ・ビデオ、および物理ワールド予測のタスクにおいてSparsePRを評価しました。

• 品質保持: SparsePRは高密度ベンチマークの品質に密接に一致します。例えば、HunyuanVideo-13Bにおいて、PSNR 31.844（高密度ベースラインに対して）およびLPIPS 0.087を達成し、SVG2やSpargeAttnといった従来の疎結合手法を上回りました。
• 効率性: 本手法は、1.48倍から2.61倍のエンドツーエンドの高速化を実現しながら、実現された実行ペア密度をわずか 22.0%–26.0% に維持しています。
• アブレーションによる洞察:
  • プローブ適合（Probe Fitting）: 誤差減少の大部分を占めています。
  • 応答結合型パーティショニング: 直交する残差エネルギーを低減することで、「ハードドロップ（hard-drop）」誤差を大幅に下げ、有限のプローブ予算下での再構成を改善します。
  • オーバーヘッド: プローブ適合による再構成は、無視できる程度のオーバーヘッド（総レイテンシの約1.1%）しか追加しません。これは、主なコスト削減が疎なルーティングから得られるためです。

5. 重要性と主張

本論文は、ビデオモデルにおける効果的な疎な実行は、しばしば別々に扱われる2つの要因、すなわちパーティション幾何学（共有ルーティングに必要なサポートを決定する）とスキップされた相互作用の再構成（出力誤差を決定する）に依存すると主張しています。

SparsePRは、パーティショニングのプロセスを現在のコールの応答幾何学に結合し、少数の厳密なプローブを用いて残差再構成を校正することで、大幅な計算節約を実現しながら、高密度と同等の忠実度を得られることを示しています。著者らは、このアプローチがトレーニングフリーであることを強調しており、再学習や蒸留を必要とせずに、既存の学習済みビデオDiTやワールドモデルに適用可能です。結果は、単純な行単位の疎性はビデオ生成には不十分であり、高品質な加速には、構造化され、応答を考慮したパーティショニングと、精密な残差補正の組み合わせが必要であることを示唆しています。