技術要約: Flash-WAM

問題提起

World Action Models (WAMs) は、反復的な拡散プロセスを通じて将来のビデオフレームとロボットの動作を共同で生成する、エンボディドAI（身体性AI）における有望な方向性を示している。これらのモデルは操作ベンチマークにおいて強力な性能を発揮するが、ビデオとアクションの両方のストリームに対して数十回のデノイジングステップを必要とすることが、計算コストの増大を招いている。例えば、最先端のLingBot-VAモデルは、NVIDIA L40S GPU上で1チャンクあたり8.1秒を要し、リアルタイムのクローズドループ制御を不可能にしている。

ステップ蒸留（step distillation）は、拡散モデルを加速させるための自然な解決策であるが、既存の「オフザシェルフ」な手法は、ビデオとアクションの共同設定に適用した場合、機能しない。核心的な問題は、WAMsに採用されている非対称なノイズ・スケジュールにある：

• ビデオ・ストリームは、高SNR（信号対雑音比）シフトを利用し、学習の質量を高ノイズレベル（$\sigma$）に集中させている。
• アクション・ストリームは、より低いSNRシフトを利用し、学習の質量を全範囲に分散させ、低$\sigma$領域に大きな密度を持たせている。

標準的な一貫性蒸留法（LCMなど）は、単一の一貫性関数に依存している。著者らは、これらの関数が $\sigma \to 0$ において二次的な勾配消失（$|b(\sigma)| \propto \sigma^2$）を示すことを実証した。その結果、これを共同モデルに一様に適用すると、アクション・ストリーム（低$\sigma$で学習される）に対する勾配信号が消失し、ビデオの再構成が維持されていても、アクションの精度とタスク成功率が崩壊（実験では90%超から約24%へ低下）するという事態を招く。

手法: Flash-WAM

著者らは、この構造的な不適合性を解決するために設計された、モダリティ認識型のステップ蒸留フレームワークであるFlash-WAMを提案する。両方のモダリティに単一の一貫性関数を適用するのではなく、Flash-WAMは各ストリームの特定のノイズ・レジームに合わせた、一貫性関数ファミリーの異なるメンバーを選択する。

1. 勾配スケーリングの構造解析

本論文は、一貫性関数が $f(x_\sigma, \sigma) = a(\sigma)x_\sigma + b(\sigma)v_\theta$ と定義されるとき、一貫性蒸留における勾配信号の強さは $|b(\sigma)|$ に比例することを確立している。

• 命題1: 有効な一貫性関数の場合、$\sigma=0$ 付近での最適な勾配スケーリングは線形（$O(\sigma)$）であり、これは $b'(0) \neq 0$ の場合にのみ達成可能である。
• 標準的なLCMの失敗: 標準的なLCMのパラメータ化は $b'(0)=0$ を満たすため、二次的なスケーリング（$O(\sigma^2)$）となる。これは高$\sigma$のビデオ学習には十分であるが、低$\sigma$のアクション学習には致命的である。

2. モダリティ認識型パラメータ化

Flash-WAMは、異なるパラメータ化を割り当てることでこれに対処する：

• アクション・ストリーム（低$\sigma$レジーム）: 著者らは、$a(\sigma) = 1$ および $b(\sigma) = -\sigma$ となる線形勾配スケーリング・パラメータ化を導出した。これにより $|b(\sigma)| = \sigma$ が保証され、アクション・ストリームの低ノイズ学習分布全体にわたって強力な線形勾配信号が提供される。
• ビデオ・ストリーム（高$\sigma$レジーム）: 高$\sigma$に集中するビデオ・ストリームに対しては、標準的な分散保存型Karrasパラメータ化（LCMで使用されるもの）を維持する。この選択は、高ノイズレベルにおける分散の保存と出力範囲の限定の必要性に由来しており、これらは低次元で有界なアクション・ターゲットにはそれほど重要ではない特性である。

3. 共同学習目的関数

生徒モデルは、以下の結合損失を用いて、両方のモダリティに対して一貫性の特性を同時に満たすように訓練される：
$$\mathcal{L} = \mathcal{L}_v + \lambda_a \mathcal{L}_a$$
ここで、$\mathcal{L}_v$ と $\mathcal{L}_a$ はそれぞれのモダリティ固有の一貫性関数を利用する。アーキテクチャは共有トランスフォーマー・バックボーンを維持し、モダリティ認識型のパラメータ化は損失ヘッドにのみ影響を与え、教師モデルからのステップあたりの計算コストは変わらない。

主な結果

本フレームワークはLingBot-VAに実装され、両方のモダリティに対して1ステップ（1v/1a）、または1つのビデオステップと2つのアクションステップ（1v/2a）へと推論を圧縮した。

• 高速化: Flash-WAMは23倍の高速化を達成し、NVIDIA L40Sにおける1チャンクあたりのレイテンシを8.1秒から348ミリ秒に短縮し、リアルタイム推論（500ミリ秒の予算以内）を可能にした。
• シミュレーション・ベンチマーク:
  • RoboTwin 2.0: 1v/2aにおいて、Flash-WAMは85.5%の成功率（教師モデルの91.2%に対し）を達成した。これは、ナイーブな共同LCMの23.9%と比較して大幅に高い。1v/1aにおいても81.4%の成功率を維持している。
  • LIBERO: 1v/2aにおいて、ビデオのみの蒸留ベースラインを上回る95.7%の平均成功率（教師モデルの98.6%に対し）を達成した。
• 実世界への展開: Unitree G1 ヒューマノイドロボットを用い、3つの操作タスクで評価を行った。
  • 加速なしの教師モデル（3v/10a）は66.7%を達成。
  • 蒸留なしでのステップ削減（1v/2a）は、性能を40.0%まで低下させた。
  • ナイーブなビデオのみのLCM（1v/2a）は43.3%であった。
  • Flash-WAM (1v/2a) は**60.0%**を達成し、加速なしの性能を大幅に回復させた。
• 定性的分析: 可視化により、ナイーブな蒸留が物体の消失やぼやけを引き起こす一方で、Flash-WAMはシーンの構造と物体の同一性を保持していることが示された。

重要性と主張

本論文は、Flash-WAMが、非対称なマルチモーダル・システムに単一モダリティの蒸留を適用した際に生じる構造的な失敗モードに対処することで、共同ビデオ・アクション拡散モデルのリアルタイム推論を初めて可能にしたフレームワークであると主張している。

• 診断: 本研究は、標準的な一貫性蒸留と共同拡散の間の不適合性が、ハイパーパラメータの調整の問題ではなく、非対称なノイズ・スケジュールによって引き起こされる勾配信号の可用性の根本的なミスマッチであることを特定した。
• 解決策: 勾配スケーリングの構造的分析に基づき一貫性関数を選択することで、Flash-WAMは、各モダリティのデノイジング・ステップを1ステップへと圧縮しながら、タスク成功率を維持する原理的な手法を提供している。
• 影響: 結果は、WAMが持つ優れた生成能力を犠牲にすることなく、汎用的なエッジ・ハードウェア上でWAMを用いたリアルタイム・ロボット制御が可能であることを示している。著者らは、ビデオ生成自体を削除する手法とは異なり、自身のアプローチは元のWAMの推論構造を維持している点を強調している。