LoRA/DoRAを用いたNVIDIA Cosmos Predict 2.5のファインチューニングによるロボット動画生成

NVIDIA Cosmos Predict 2.5は、テキスト、画像、またはビデオクリップから物理的に妥当なビデオを生成できる大規模世界モデルです。これをロボット操作や特定のカメラ視点といった特定のドメインに適応させるには、ターゲットを絞ったファインチューニングが必要です。しかし、20億パラメータモデルのフルファインチューニングは高コストであり、汎用知識の破滅的忘却リスクがあります。LoRA（低ランク適応）とDoRA（方向性低ランク適応）は、凍結されたベースモデルに小型の学習可能アダプターモジュールを注入することで、メモリ要件を削減し、アダプターファイルを小型でポータブルに保ちます。これにより、単一GPUでのファインチューニングが実用的になり、推論時に異なるドメイン用のアダプターを柔軟に交換できます。

本ガイドでは、diffusersおよびaccelerateライブラリを使用し、単一GPUおよび複数GPUトレーニングをサポートする、LoRAおよびDoRAを用いたCosmos Predict 2.5のパラメータ効率的なファインチューニングを説明します。その後、ファインチューニングされたモデルを使用して、ダウンストリームのロボット学習タスク用の合成ロボット軌跡を生成する方法を示します。

まず、環境要件としてPython 3.10+、PyTorch 2.5+、diffusers、accelerateなどをインストールします。データ準備では、GR00T Dreams後訓練レシピと同じデータセットを使用します。訓練データセットは92個のロボット操作ビデオとテキストプロンプト、テストデータセットは50個の（プロンプト、画像）ペアです。

訓練実装では、VideoDatasetクラスが各サンプルを（キャプション、ビデオ）ペアとして読み込み、長いビデオからランダムな連続ウィンドウをサンプリングします。アダプター初期化では、VAE、テキストエンコーダ、DiTの全重みを凍結し、DiTの注意射影層とフィードフォワード層にLoRAアダプターを注入します。損失関数は整流流（rectified flow）を採用し、条件付けされていないフレームに対してのみMSE損失を計算します。最適化にはAdamWを使用し、線形学習率スケジューラでウォームアップと減衰を行います。チェックポイントはLoRA重みをdiffusers形式で保存します。訓練コマンド例では、lora_rank=32（約5000万訓練可能パラメータ）をデフォルトとし、accelerateで分散訓練をサポートします。100エポックで単一H100上約17時間、8台のH100で2.5時間で妥当な結果が得られます。

推論時には、ファインチューニングされたLoRA重みをパイプラインにロードし、融合することで推論オーバーヘッドを排除します。arch_invariant_rand関数で初期潜在ノイズを生成し、再現性を確保します。評価指標としてSampson誤差を用い、生成ビデオの幾何学的整合性を測定します。