NVIDIA Cosmos Predict 2.5 を LoRA/DoRA でファインチューニング、ロボット動画生成
NVIDIA の大規模世界モデル「Cosmos Predict 2.5」を対象に、ロボティクス分野への適応を目的とした効率的なファインチューニング手法が公開されました。本モデルはテキストや画像、動画クリップを条件として物理的に妥当な動画生成が可能ですが、特定ドメインへの適合には依然としてターゲット学習が必要です。実機からのデータ収集は時間とコストがかかるため、ファインチューニングされたモデルによる合成データの生成が有効な代替案となります。しかし、20 億パラメータのモデル全体を学習させるのは高コストであり、既存の一般知識が失われるリスク(カタストロフィック・フォッゲッティング)があります。 これを解決するため、本研究では LoRA(Low-Rank Adaptation)および DoRA(Weight-Decomposed Low-Rank Adaptation)を用いたパラメータ効率的なファインチューニングアプローチを提案しています。これらは基盤モデルの重みを固定したまま、小さなアダプターモジュールを追加する方式であり、メモリ要件を大幅に削減しつつ、異なるドメイン間でのアダプターの柔軟な切り替えを可能にします。実装は PyTorch の加速ライブラリや Diffusers を使用し、単一 GPU もしくは複数 GPU 環境でのトレーニングがサポートされています。 トレーニングプロセスでは、VideoDataset を用いてデータセットを読み込み、VideoProcessor による前処理後に DiT(Diffusion Transformer)の注意力機構や前全結合層のみに LoRA アダプターを挿入します。勾配計算の安定性を保つため、訓練可能なパラメータのみを浮動小数点 32 倍精度で処理します。学習ロスは、ノイズから清浄データへ直線的に移動させる速度を予測する「リフレクトフロー」方式に基づき、条件付きフレーム以外のフレームに対してのみ計算されます。オプティマイザには AdamW、スケジューラには線形ウォームアップ・デケイを用い、アダプター重みは定期的にチェックポイントとして保存されます。 実験結果によると、100 エポック(8 枚の H100 GPU で約 2.5 時間)の学習で、サンプソン誤差(幾何学的一貫性)、物理法則の遵守度、指示従順性の 3 つの指標が大幅に改善されました。LoRA と DoRA は同程度の性能を発揮し、特にランク 32 の設定はロボットの動作や物体との相互作用といった指示への従順性を向上させることが確認されました。一方、幾何学的および物理的な事前知識は基盤モデルに既に含まれているため、アダプターはドメイン固有の見た目のシフトだけで済み、ランク 8 でも十分な成果が得られると結論づけられています。メモリ制約が厳しい場合はランク 8 の LoRA を、学習の安定性が課題となる場合はランク 32 の DoRA を推奨します。
