【DL輪読会】Cosmos World Foundation Model Platform for Physical AI

DEEP LEARNING JP [DL Papers] Cosmos World Foundation Model Platform for Physical AI Makoto Kawano (@mkt_kwn), Matsuo Lab. http://deeplearning.jp/

• http://deeplearning.jp/
• http://twitter.com/mkt_kwn

書誌情報 • Cosmos World Foundation Model Platform for Physical AI ▪ NVIDIAがCES2025で発表した世界基盤モデル ▪ 2000万時間の動画を学習した２種類のモデルを提供 Custom Dataset Post-trained WFM ✓ 拡散モデルベース Custom Dataset Post-trained WFM ✓ 自己回帰モデルベース Pre-trained WFM Custom Dataset Post-trained WFM ▪ Tokenizerや事後学習事例，Guardrail機能も同時に提供 Custom Dataset Post-trained WFM ▪ 10000個のH100を３ヶ月使って学習 ▪ プロジェクト，arXiv, github 2

• https://www.nvidia.com/en-us/ai/cosmos/
• https://arxiv.org/abs/2501.03575
• https://github.com/NVIDIA/Cosmos

アジェンダ(ほぼ論文通り) • 全体の構造 • データセット構築 • トークナイザ • 事前学習モデル • 事後学習モデル • ガードレール • まとめ 3

背景/モチベーション • Physical AI / Embodied AIは，センサとアクチュエータを持つシステム ▪ センサ：世界を観測するために利用 ▪ アクチュエータ：世界と相互作用のために利用 危険/退屈/面倒な作業を 人間の代わりにやってくれる ▪ 深層学習の発展に比べるとだいぶ遅い -> 観測とアクションをいい感じで配置した大量の訓練データを用意できない ため 物理空間のデジタルツイン(=世界基盤モデル)があればこの問題が解決するはず 4

Cosmosの構造 • 基本的にこれまでの世界モデル※１と同じアーキテクチャ ▪ 1. 画像/動画，テキスト，アクション(座標軌道)をトークン化する ▪ 2. 各トークンをいい感じでTransformerベースのモデルに入力する ▪ 3. 画像/動画トークンを予測する ▪ 4. デコーダで予測したトークンをピクセル空間に戻す ※１ GAIA-1やSora, DriveDreamer-2，Vistaなど 5

データセット構築 • インターネット上にある動画データを収集し，加工して作成 ▪ なんらかのオブジェクトとアクションが含まれている動画 1. Driving (11%) 6. Nature dynamics (20%) 2. Hand motion and object manipulation (16%) 7. Dynamic camera movements (8%) 3. Human motion and activity (10%) 8. Synthetically rendered (4%) 4. Spatial awareness and navigation (16%) 9. Others (7%). 5. First person point-of-view (8%) ▪ 2000万時間分→事前学習用に1億クリップ，事後学習用に1000万クリップを作成 ▪ 加工の工程： Raw Input Video Split • Shot Detection • GPU-based Transcoding Filtering • • • • Motion Filtering Quality Filtering Overlay Text Filtering Video Type Filtering Annotation Video Clip Database Dedup Sharding • Video Description Generation 6

分割 • シーン転換の検出 ▪ 分布(PySceneDetect)や埋め込み(Panda70M)変化検出，TransNetV2/AutoShotで検証 ▪ 複数のデータセットで検証した結果，TransNetV2が一番性能が良い • 2秒未満は削除，60秒以上は分割し，mp4へエンコード ▪ データローダの際に，フォーマットを統一していた方が安定性・効率が良い ▪ NVIDIA L40SとPyNvideoCodecを駆使することで高速化が可能 7

フィルタリング • 要件を満たさないクリップの削除，事後学習用の高品質データの選択をする ▪ Motion：静止動画や手ぶれ動画を削除し，モーションのアノテーションを実施 ✓ パン，ズーム，ティルトなど ✓ 推定されたオプティカルフローを入力として分類器を用意する ▪ Visual Quality: 低品質/低aestheticな動画を削除 ✓ スコア算出可能なNNを用意して，閾値未満の動画を削除 ▪ Text: 字幕のようなテキストが埋め込まれた動画も削除 ✓ InternVideo2による埋め込みから二値分類で削除 ▪ Video Type： 抽象的な映像やアニメ，ゲーム動画も削除 ✓ InternVideo2の埋め込みから分類して削除 8

アノテーションと重複削除 • 複数のVLMを検証し，VILA-13Bを用いてクリップ毎にキャプションを付与 ▪ 8bitにするTensorRT-LLMエンジンを用いて高速化（x10) ▪ 8フレーム+“Elaborate on the visual and narrative elements of the video in detail” ▪ 平均97単語くらいのキャプションが付与 • K-meansを実施し，クラスタ内で距離を計算し，重複を削除 ▪ 学習効率の向上と丸暗記の機会を減らす ▪ 重複した場合は高画質の方を残す ▪ InternVideo2の埋め込みで実施→30%削除 9

トークナイザ • 世界モデルのタイプに合わせて，連続トークンと離散トークンの２種類を用 意 ▪ ナイーブなAutoencoder構造で，トークン側に損失関数を設けず ✓ 連続：KLダイバージェンスなし(≠VAE) ✓ 離散：commitment lossなし(≠VQVAE)(離散化はFinite-Scalar-Quantization) ▪ 損失関数は下記の４種類 ✓ L1ロス ✓ Perceptualロス ✓ オプティカルフローロス ✓ グラム行列ロス 10

トークナイザのアーキテクチャ • 3D Wavelet変換x2 + 3D Spatial-Teporal (Causal)畳み込み層 𝑥1 𝑥2 𝑥3 𝑔0 𝑥4 𝑥5 𝑥6 𝑥7 𝑥8 𝑥9 𝑥10 𝝃1 𝑥11 𝑔3 𝑔2 𝑔1 𝝃0 𝝃2 𝝃3 𝑥12 ▪ 画像空間/時間軸方向で低周波成分/高周波成分に分解 ▪ 1回の変換で，時空間方向で1/2される W avelet3D Tran sform Sp atio -Tem p oral Convolution and Self -Attention ▪ 2回実施： CxTxHxW => 64C x T/4 x H/4 x W/4 ×N Encod er Inverse H aar W avelet3 D Causal ResBlock3D Causal Up Sam p le3D Causal Sp atioTem p oralAttn Causal Sp atioTem p oralAttn Causal Dow nSam p le3D Causal ResBlock3D ▪ Causal畳み込み：過去側のみPaddingを実施 H aar W avelet3 D 𝑥0 ▪ 空間方向にまず畳み込む ▪ そのあと時間方向に畳み込む ▪ エンコーダ・デコーダ：対称的な構造 ×N Decod er 11

トークナイザの種類 • [系列の長さ]-[連続|離散] x [画像|動画]-[時間圧縮]x[空間圧縮] ▪ 系列の長さ： 0.1→短い系列， 1.0→長い系列で学習 ▪ 連続|離散 x 画像|動画： CI, CV, DI, DVの４種類 ▪ 時間圧縮 x 空間圧縮： 4xと8x, 8x8と16x16の組み合わせ 動画の連続トークン化による再構成評価 12

トークナイザの性能 • 性能と圧縮率のトレードオフ 13

Cosmos-Tokenizerの結果 14

事前学習モデル • 拡散モデルタイプ(7B/14B)と自己回帰モデルタイプ(4B/12B)の２種類を提 案 ▪ それぞれ２ステージ学習を実施 ✓ 拡散モデル： テキスト->動画 → テキスト+動画 -> 動画 ✓ 自己回帰モデル： 動画->動画 → 動画+テキスト -> 動画 連続トークン 離散トークン 15

拡散モデルベース • EDMxDiffusionTransformer トークンをTransformerに 入力するため,パッチ化する テキストで条件付けするため， Cross-Attentionを利用 任意の解像度や系列長の動画生成の 計算量削減のため， ため，FPS-awareなRoPEや Low-Rank Adaptationを利用 学習可能な位置埋め込み 16

学習戦略 • Progressive Training：少しずつ解像度・系列長を増やして学習 • Multi-aspect Training: 異なるアスペクト比を並列で学習 ▪ 最長辺にリサイズして，1:1, 3:4, 4:3, 9:16, 16:9の中で一番近いアスペクトにリサイズす る • Mixed-Precision Training: 演算時BF16，パラメータ更新時FP32で計算 • スケールさせるために，様々な工夫を導入 ▪ Fully Sharded Data Parallelism:パラメータを分割して，分散環境に配置する ▪ Context Parallelism：コンテキストの長さに合わせてQ/KVを分割する 17

Prompt Upsampler • 学習データのテキストと推論時のユーザ入力のテキストは全然違う ▪ LLMを使って，ユーザ入力テキスト -> 学習データ（VLM)スタイルに変換する ▪ 要件として： ✓ 元のプロンプトに対して忠実に変換 ✓ 学習データに似たプロンプトに変換する ✓ より画像に忠実なプロンプトにする • ２種類のUpsamplerを用意する ▪ Text2World： できるだけ短いプロンプトに変更する ▪ Video2World: Pixtral-12B→ユーザプロンプトと動画内容を踏まえて生成 18

自己回帰モデルベース • LLMと同じnext-token predictionで学習する ▪ Z-lossℒ = 𝜆 ⋅ σ𝑖 𝑧𝑖2 を使って学習を安定化 テキストで条件付けするため， Cross-Attentionを利用 語彙数はLSQにより， 8x8x8x5x5x5=64000 YaRNを使って可変長な CoolingDownで最終的な コンテキスト窓に対応させる 重みを保存 19

推論改善 • Speculative Decoding(Medusa)でリアルタイム推論を可能にする ▪ 最終層にMedusa Head層を追加し，同時に複数個のトークンを生成 ▪ Medusa Headの結果と元のHeadの結果を比較して一致した場合に採択する 20

Diffusion Decoder • トークナイザで圧縮させすぎると生成結果にボケと人工物が発生してしまう ▪ 拡散モデルを再利用して生成結果を改善する Input Video Encoder of Cosmos-1.0TokenizerCV8x8x8 + Continuous Tokens Encoder of Cosmos-1.0TokenizerDV8x16x16 Gaussian Noise Vocabulary Embedding Discrete Tokens Denoiser of Cosmos-1.0Diffusion-7B Decoder of Cosmos-1.0TokenizerCV8x8x8 Recon. Video Conditional Input ↑ 2x Upsample 21

Diffusion Decoder • トークナイザで圧縮させすぎると生成結果にボケと人工物が発生してしまう ▪ 拡散モデルを再利用して生成結果を改善する Denoiser of Cosmos-1.0Diffusion-7B Gaussian Noise Cosmos-1.0Autoregressive Vocabulary Embedding Discrete Tokens Decoder of Cosmos-1.0TokenizerCV8x8x8 Recon. Video Conditional Input ↑ 2x Upsample 22

評価 • ２種類の一貫性で3D一貫性を評価 ✓ 幾何学一貫性：Sampson Errorとカメラのポーズ推定誤差 ✓ サンプソン誤差：ある空間における1点を２箇所のカメラ視点で見た時に一致しているか ✓ ビュー合成一貫性：生成動画から3DGaussian Splattingを学習し，テストデータで評 価 23

評価 • 物理法則に従っているかどうか ▪ PhysXとIsaacSimを使って動画を生成，1/9フレームを条件付けて生成した動画と比較 24

物理法則評価 25

事後学習モデル｜カメラ制御 • Cosmos-1.0-Diffusion-7B-Video2Worldを事後学習 ▪ 一枚の画像とカメラ軌道から動画を生成させる ▪ ノイズ埋め込みにプリュッカー座標を連結させて学習させる ✓ プリュッカー座標：直線を2点ではなく，６次元ベクトルで表現する ▪ DL3DV-10Kからクリップを作成，SfMで各フレームのカメラ位置を推定 ✓ １フレーム目を原点として残りのフレームは相対的な座標にする ✓ トークナイザは時間方向に1/8するため，8フレームのうち，４フレーム目の座標を利 用 26

結果 27

事後学習モデル｜ロボット操作 • ロボット操作を指示条件付け/アクション条件付けした時の動画生成 ▪ 拡散モデル/自己回帰モデル両方で検証 ▪ 指示条件：自前で収集したデータセット，12000エピソード（1-9秒の動画） ✓ 各エピソード1文の指示が付与，VLMでアップサンプルする ✓ 両モデル，Cross-Attentionでfine-tuning ▪ アクション条件付け：Bridge，20000エピソード ✓ グリッパーの座標7次元(Δ𝑥, Δ𝑦, Δ𝑧, Δ𝜃𝑟 , Δ𝜃𝑝 , Δ𝜃𝑦 , Δ𝐺𝑟𝑖𝑝𝑝𝑒𝑟) ✓ 自己回帰モデル：MLPでテンソル化，Cross-Attentionでfine-tuning ✓ 拡散モデル：MLPでテンソル化，タイムスタンプ埋め込みに加算してfine-tuning 28

評価｜指示条件づけ • 無作為に二つの動画をユーザに見せて，質問に対して動画を選んでもらう ▪ 指示に従っているか？ ▪ へんなものが映らずに実世界に忠実か？ ▪ 物体は映り続けているか？▪ 全体を通して道理に適っているか？ VideoLDM-Instruc on VideoLDM-Instruc on Ties Ties Cosmos-1.0-Diffusion-7B-Video2World-Sample-Instruc on 100% 90% 80% 70% 17.4% 8.7 13.0% 8.7 13.0% 100% 13.0% 90% 8.7% 80% 17.4% 26.1% 34.8% 34.8% 21.7% 17.4% 43.5% 47.8% 43.5% Instruc on Following Object Permanence Verity 30.4% 70% 60% 60% 30.4% 13.0% 50% 50% 82.6% 40% 30% Cosmos-1.0-Autoregressive-5B-Video2World-Sample-Instruc on 65.2% 73.9% 78.3% 40% 30% 20% 20% 56.5% 10% 10% 0% 0% Instruc on Following Object Permanence Verity Overall Overall 29

評価｜指示条件づけの結果 30

評価｜アクション条件付け結果 31

事後学習モデル｜自動運転 • 自動運転においては多視点動画を生成できることが好ましい ▪ 自前で収集したデータからデータセットを作成（360万x20秒の動画） ✓ 前後左右，右後ろ，左後ろの６視点動画 ✓ 天気や明るさ，自車情報，道路の種類を満遍なくなるようにデータを選択 ▪ 拡散モデルを利用して，３種類の条件付けモデルをfine-tuning ✓ Cosmos-1.0-Diffusion-7B-Text2World-Sample-MultiView ✓ Cosmos-1.0-Diffusion-7B-Text2World-Sample-MultiView-TrajectoryCond ✓ Cosmos-1.0-Diffusion-7B-Video2World-Sample-MultiView 32

多視点生成のための工夫 いまいちここがわからず • 視点に関する情報は，視点依存のCrossAttentionで入力 ▪ 位置埋め込みは，事前学習時と一緒 ▪ 視点についてはキャプションで付与されている ✓ “The video is captured from a camera mounted on a car. The camera is facing forward|left|right|backward|rear-left|rear-right.” • 軌道は64点で表現され，初期位置(0,0,0)から0.1秒間隔で区切られている ▪ （論文内では，Cosmos-1.0-Diffusion-7B-Video2Worldとあるが） Cosmos-1.0-Diffusion-7B-Text2World-MultiViewのfine-tuningをする 全体的にモデル名が混同しているっぽいのでコード公開を待つべき 33

結果 • Text2Videoでテキスト条件付けし，さらにVideo2Worldで秒数を増やす 34

ガードレール • 安全に利用できるように機能を実装 ▪ Pre-Guard: 有害なプロンプトが入力されないようにガード ✓ ハードコーディングされたキーワードマッチング ✓ Aegis：暴力や犯罪などに関係しそうなプロンプトを分類 ▪ Post-Guard: 動画コンテンツ安全性判定と顔のぼかし ✓ 事前に分類器を作成し，１フレームでも危険と判定されたら動画全体を危険とする ✓ 20x20ピクセル以上の顔に対してぼかしを入れる WFM Inputs Keyword Blocking pre-Guard Aegis World Foundation Model Video Content Safety Classifier post-Guard Face Blur WFM Outputs 35

まとめ 特段新技術はないが，大規模モデルのテクニックをふんだんに使った世界モデ ル ▪ ライセンス・コード公開など，大判振る舞い • 記載があるが，シミュレータとしてはまだまだ ▪ 物体消失や接触時のダイナミクス，指示に従い切らない ▪ 特に物理法則を遵守しないあたりはかなり致命的 ▪ 評価方法も人間が含まれているから怪しい • 拡散モデルvs自己回帰モデル ▪ 性能だけで言えば，拡散モデル ▪ LLMの重みの有効活用など自己回帰モデルでまだ探索してない可能性がある 36