技術的概要：効率的な動画生成のための不均一ステップ割当（HSA）

1. 問題提起

拡散トランスフォーマ（DiT）は動画生成において最先端の性能を確立していますが、莫大な計算コストに悩まされています。標準的な推論プロトコルは、シーケンス内のすべての時空間トークンに均一な数（$T$）のノイズ除去ステップを適用します。この単一構造のアプローチは、視覚データが本質的に非対称であることを無視しています。背景領域、静的なパッチ、粗い構造は知覚的に単純であり、詳細に富む前景オブジェクトや高運動領域よりも少ないノイズ除去ステップしか必要としません。さらに、人間の視覚は、静的なコンポーネントと比較して、非常に動的なコンポーネントにおける品質劣化に対して感度が低いです。したがって、すべてのトークンに完全なノイズ除去予算を均一に適用することは計算上の無駄であり、知覚的な優先順位と整合しません。

既存の加速手法、例えばステップレベルの機能キャッシング（TeaCache など）は、通常、グローバルな意思決定を行います。すなわち、任意の反復において、すべてのトークンが計算されるか、すべてのトークンがキャッシュされた出力を再利用するかです。これらの手法は、空間的またはトークンレベルの不均一性を活用せず、可変的な動的要件を持つトークンの集合ではなく、単一構造のブロックとしてシーケンスを扱います。

2. 手法：不均一ステップ割当（HSA）

著者は、事前学習された DiT 内でトークン固有のノイズ除去軌道を割り当てる、トレーニング不要でモデルに依存しない推論アルゴリズムである**不均一ステップ割当（HSA）**を提案します。中核的なアイデアは、トークンシーケンスを速度の動力学に基づいてグループに分割し、よりゆっくりと変化するトークンには少ないステップを割り当てながら、基準グループには完全なスケジュールを維持することです。

2.1 トークンの分割とスケジューリング

$N$個のトークンの集合は、$K$個の互いに素なグループ（$G_1, \dots, G_K$）に分割されます。各グループ$G_k$には、総ステップ数$T$の約数であるステップ予算$s_k$が割り当てられます。

• 基準グループ（$G_1$）： 完全な予算$s_1 = T$が割り当てられます。これらのトークンはすべての反復でアクティブです。
• 削減グループ： より小さな予算（例：$T/2, T/4$）が割り当てられます。これらのトークンは、グローバル反復$i$が$i \pmod{T/s_k} = 0$を満たす場合のみアクティブになります。
• 動的選択： 論文は、トークンをグループに割り当てるための 4 つの戦略を評価します。動的トークン選択戦略は、速度予測のトークンごとの L1 相対変化に基づいてトークンを割り当てるもので（背景などのゆっくりと変化する領域を特定）、一貫して最良の性能を発揮します。

2.2 KV キャッシュ同期

不均一なスケジューリングにおける主な課題は、異なるステップ数を持つトークンが、任意のウォールクロック反復において異なるノイズレベル（$\sigma$）に存在することです。自己注意（Self-attention）は、すべてのトークンが互いに注意を向けることを必要とするため、同期の不一致が生じます。HSA は、KV キャッシュ同期メカニズムを通じてこれを解決します。

1. アクティブトークン： アクティブなトークンの部分集合（$A_i$）のみが QKV 投影を介してフォワードパスを実行します。
2. キャッシュ更新： アクティブトークンの新しいキー（$K$）とバリュー（$V$）投影が、レイヤーごとのキャッシュに書き込まれ、古くなったエントリを上書きします。
3. フルコンテキスト注意： アクティブなトークンは、完全なキャッシュされた$K, V$行列（アクティブトークンの新しいエントリとスキップされたトークンの古いが有効なエントリの両方を含む）に対してクエリを計算します。
4. 複雑性の削減： これにより、自己注意の複雑さは$O(N^2)$から$O(|A_i| \cdot N)$に削減されます。ここで、$|A_i|$はステップ$i$におけるアクティブトークンの数です。スキップされたトークンは、非アクティブな反復中にクエリ計算、クロス注意、またはフィードフォワード操作を一切行いません。

2.3 キャッシュ済みオイラー更新

追加のモデル評価なしにスキップされたトークンの潜在状態を進めるために、HSA はキャッシュ済みオイラー更新を採用します。

• 各トークンは、最も最近のアクティブステップで予測された速度（$\hat{v}_n$）を保存します。
• すべてのグローバル反復において、すべてのトークン（アクティブおよびスキップされたもの）が、同じ増分オイラーステップ（$\sigma_{i+1} - \sigma_i$）だけ進められます。
• アクティブなトークンは新しく計算された速度を使用し、スキップされたトークンはキャッシュされた速度を再利用します。
• この更新は、すべての$N$個のトークンにわたる単一のテンソル操作として形式化され、分岐を回避して GPU の効率性を維持します。

2.4 位相認識キャッシングウィンドウ

初期および後期のノイズ除去段階がそれぞれグローバル構造と微細な詳細にとって重要であることを認識し、HSA はキャッシングを軌道の中間部分に制限します。最初の$m$ステップと最後の$m$ステップでは、すべてのトークンに対して完全なキャッシングなしのフォワードパスを実行し、構造的完全性と詳細の洗練が古くなった近似によって損なわれないことを保証します。

3. 主な貢献

1. HSA アルゴリズム： 速度の動力学に基づいて時空間トークングループに不均一なノイズ除去軌道を割り当てる、完全にトレーニング不要でプラグアンドプレイの推論パラダイム。
2. KV キャッシュ同期： アクティブなトークンに対して完全な$N$トークン注意コンテキストを維持しつつ、非アクティブなトークンを完全にバイパスし、反復ごとの注意コストを$O(|A_i| \cdot N)$に削減するメカニズム。
3. キャッシュ済みオイラー更新： キャッシュされた速度を使用してスキップされたトークンを進める単純なテンソル操作であり、追加のモデル評価なしに生成品質を維持します。
4. 包括的な評価： 4 つのトークングループ化戦略と 4 つの予算プリセットの分析により、異なる推論予算全体で堅牢性を示しています。

4. 実験結果

著者は、HSA をWan-2.1-1.3BおよびWan-2.2-A14B（テキストから動画、画像から動画）とLTX-2（音声から動画）で評価しました。

• 品質と実行時間のトレードオフ： HSA は、均一なフローマッチング（FM）や TeaCache などの最先端キャッシング手法と比較して、優れたパレートフロンティアを達成します。
• 積極的な加速： 利点は、積極的な加速領域（実行時間の 50% および 25%）で最も顕著です。実行時間の 25% において、HSA は VBench ベンチマークでベースラインを大幅に上回ります（総合スコア：79.87% 対、FM の 75.68%、TeaCache の 78.33%）。
• 次元の堅牢性： ベースライン手法は、厳しい予算下で「壊滅的な次元の崩壊」（運動の滑らかさや時間的一貫性などの特定の VBench 次元における深刻な劣化）に苦しむのに対し、HSA はすべての 16 の評価次元にわたって完全予算の参照を堅牢に追跡します。
• 画像から動画（I2V）： HSA は、I2V タスクにおいて画像条件付けの整合性と視覚的忠実度を成功裏に維持しますが、ステップを削減した（$T=10$）均一な FM ベースラインは、深刻な構造的崩壊を示します。
• オフラインプロファイリング不要： モデル出力の揺らぎを推定するために高価なオフラインプロファイリングを必要とする TeaCache と異なり、HSA はオフラインチューニングや微調整なしに動作します。

5. 意義と主張

本論文は、HSA が視覚データの知覚的および動的複雑さに計算リソースを整合させることで、DiT 推論における根本的な非効率性に対処すると主張しています。トークンをグローバルにではなく個別に扱うことで、HSA は、積極的な加速に伴う通常品質の劣化なしに大幅な高速化を実現します。

著者は、HSA がさまざまなトークングループ化戦略でインスタンス化可能な柔軟なフレームワークであることを強調しています。現在、動的選択が最も強力な結果をもたらしていますが、このフレームワークは、より洗練されたグループ化戦略によるさらなる改善を想定しています。この研究は、時空間データの本質的な不均一性を活用して高忠実度動画生成を民主化する、新しいクラスの推論アルゴリズムの基礎を築いています。