技術概要：Veda - 蒸留された疎アテンションによるスケーラブルな動画拡散

1. 問題定義

高解像度かつ長時間の動画生成における拡散トランスフォーマ（DiT）のスケーリングは、現在、時空間トークン系列長に対する自己アテンション機構の二次的な計算量およびメモリ複雑度（$O(N^2)$）によってボトルネックとなっています。疎アテンションはこれを緩和する理論的な道筋を提供しますが、既存の手法は以下の重要なトレードオフに直面しています：

• 静的パターン： 事前定義された疎パターン（例：スライディングウィンドウ）を使用する手法は、DiT によって学習される動的かつヘッド固有のアテンション構造への適応性が欠如しています。
• 動的選択： 動的マスクを学習する手法（例：VSA、VMOBA）は、しばしば不正確な推定に悩まされます。これらは拡散目的関数による暗黙の教師あり学習と平均プーリングなどの粗い統計量に依存しており、完全アテンションの真の幾何学構造との不一致を招きます。
• 核心的な課題： 実証的証拠は、高い疎性レベル（≥90%）において、これらの既存アプローチが生成品質を劣化させ、空間的な歪み、「水紋」パターン、時間的なフリッカーなどの構造的アーティファクトを生成することを示しています。本論文は、この劣化は疎性そのものによるものではなく、疎マスクが完全アテンションのタイルごとの幾何学構造と整合しないことに起因すると仮説立てています。

2. 手法：Veda フレームワーク

Veda は、極端な疎性下でも完全アテンションのタイルごとの構造とランキングを保持するように設計された、蒸留された疎アテンションフレームワークです。これは、タイル選択を拡散目的関数の暗黙的な副産物ではなく、明示的な再構成問題として扱います。

2.1. 蒸留されたタイルスコアリング

拡散損失に依存して疎性を形成するのではなく、Veda は軽量な推定器を用いて、完全アテンションバックボーンからタイルレベルのアテンションスコアを明示的に学習します。

• ターゲット構築： 参照用の「オラクル」マスクは、クエリ - キータイル領域にわたって完全アテンション行列に対して最大値プーリングを適用することで構築されます。平均化ではなく最大値プーリングが選択されたのは、背景ノイズによって希薄化されがちな顕著な高周波信号のピークを保持するためです。
• 統計量認識型推定器： これらのスコアを再構成するために、Veda は各タイルに対して平均、最大、最小の統計量を連結するTripPool記述子を使用します。これにより、単純な平均プーリングを超えたタイル表現が強化されます。
• ヘッド固有の投影： 推定器は、異質な依存パターンを捉えるために、各アテンションヘッドに対して個別の MLP 投影を使用します。
• 最適化： モデルは、予測されたタイルスコア分布を完全アテンション参照と整合させるために、行ごとのKL 発散蒸留目的関数（$L_{distill}$）を用いて訓練されます。重要なのは、推定器に供給されるバックボーン特徴に対して勾配停止操作が適用されることです。これにより、マスク学習と特徴学習が分離され、事前学習された生成多様体の破壊が防止されます。

2.2. ヘッド認識型タイル化

アテンションヘッドが時空間依存性において顕著な異質性を示す（一部は局所的な空間相互作用を捉え、他方は長距離の時間依存性を捉える）ことを認識し、Veda は均一なタイル化戦略を放棄します。

• 構成探索： 各レイヤーとヘッドに対して、Veda はハードウェアのタイルサイズを因数分解する最適なタイル化構成 $\pi_{l,h} = (p_t, p_h, p_w)$ を探索します。
• オフライン選択： 最適な構成は、完全アテンション出力と疎アテンション出力間の近似誤差を最小化するように、較正セット上でオフラインで選択され、タイル化戦略が各ヘッドの特定の構造的ニーズに合致することを保証します。

2.3. ハードウェア効率的な実装

理論的な FLOP 削減を実際のウォールクロック速度向上に変換するために、Veda は ThunderKittens DSL を使用したカスタムタイルスキップ疎アテンションカーネルを実装します。

• 非同期実行： このカーネルは、NVIDIA Hopper の Tensor Memory Access（TMA）とワープ特化を活用します。プロデューサー - コンシューマーパラダイムを用いてデータ移動と計算を分離し、プロデューサーワープが非連続なキー/バリュータイルを共有メモリにフェッチしている間に、コンシューマーワープがテンソルコア操作を実行します。
• 効率性： この設計は、密行列計算の背後でメモリ遅延を隠蔽し、FlashAttention-3 のメモリ FLOPs 利用率（MFU）の約**80%**を達成します。
• 訓練効率： 専用の TileLang カーネルが 2 パスでグランドトラウスヒートマップを生成し、パフォーマンスを損なうことなく訓練オーバーヘッドを削減するための疎な教師（クエリタイルのランダムな部分集合）を可能にします。

3. 主要な貢献

1. 実証的洞察： 本論文は、疎動画拡散における生成品質は、疎性比率そのものではなく、疎マスクと完全アテンションのタイルごとの幾何学構造との整合性によって決定されることを実証しています。
2. 蒸留された疎アテンション： Veda は、タイル選択を明示的な再構成問題として定式化するフレームワークを導入し、統計量認識型推定器とヘッド固有の投影を用いて推定誤差を最小化します。
3. ヘッド認識型タイル化： 異なるアテンションヘッドに異なる時空間タイル因数分解を割り当てる新規戦略であり、Video DiT における依存パターンの異質性に対処します。
4. ハードウェア最適化カーネル： 系列長に対してほぼ線形にスケーリングし、高いハードウェア利用率を達成するカスタムタイルスキップカーネルにより、アルゴリズム的な疎性を具体的なレイテンシ向上に変換します。

4. 実験結果

大規模な動画拡散モデル（Waver-T2V-12BおよびWan2.1-T2V-14Bを含む）を用いた実験は、最大 720P の解像度および最大 241 フレームの系列長で行われました。

• 速度向上： Waver-T2V-12B（720P、241 フレーム）において、Veda は5.1 倍のエンドツーエンド速度向上（サンプリング時間を 19.4 分から 3.8 分に短縮）と10.5 倍の自己アテンション速度向上を達成しました。アテンションのオーバーヘッドは 92% から 50% に削減されました。
• スケーラビリティ： 速度向上は系列長とともに増加します。245K トークンにおいて、Veda は完全アテンションの二次的な成長（1576.5ms）と比較して、ほぼ線形にスケーリング（309.1ms）し、5.1 倍の速度向上をもたらします。
• 品質保持：
  • 人間評価： Waver-bench 1.0 において、90% 疎性の Veda は完全アテンションと知覚的な同等性を達成します。95% 疎性においても、VSA 法（より低い疎性で動作する場合でも）を大幅に上回ります。
  • アーティファクト削減： 高い疎性で水紋パターンや時間的フリッカーに悩まされるベースラインとは異なり、Veda は高い視覚忠実度と時間的一貫性を維持します。
  • 定量的指標： VBench 評価により、Veda は 95% 疎性であっても、完全アテンションのベースラインと同等の被写体の一貫性と滑らかな動きを維持することが示されました。

5. 意義と主張

本論文は、Veda が構造的完全性を損なうことなく、高解像度動画生成のスケーラビリティボトルネックを成功裡に解決したと主張しています。疎アテンションをヒューリスティックな探索や暗黙の学習タスクではなく、蒸留問題として再定義することで、Veda は以前は深刻な品質劣化なしには達成不可能だった積極的な疎性（最大 95%）を可能にします。

著者らは、これらの成果が単に理論的なものではなく、カスタムカーネルが疎性を実用的なウォールクロック加速に変換することを強調しています。このフレームワークは、時空間解像度に対して有利にスケーリングするモジュール式かつハードウェアを認識したソリューションとして提示されており、現在のハードウェア制約下でも高忠実度かつ長時間の動画生成を可能にします。今後の改善は、より密なカーネル融合、タイムステップ全体での適応的疎性戦略、およびタイルスコアの時間的キャッシングに焦点を当てることが示唆されています。