NVIDIA Cosmos Predict 2.5 : Fine-tuning LoRA/DoRA pour la vidéo robotique
NVIDIA a publié un guide technique détaillé pour le fine-tuning de son modèle de monde Cosmos Predict 2.5, un système de génération vidéo capable de produire des séquences physiquement plausibles. L'objectif principal est d'adapter ce modèle de 2 milliards de paramètres à des domaines spécifiques, comme la manipulation de robots, sans les coûts prohibitifs et les risques d'oubli catastrophique associés au fine-tuning complet. Pour résoudre ces problèmes, l'approche proposée utilise des techniques d'adaptation paramétrique efficace, spécifiquement LoRA (Low-Rank Adaptation) et DoRA (Weight-Decomposed Low-Rank Adaptation). Ces méthodes consistent à geler les poids du modèle de base et à injecter de petits modules d'adaptation entraînés. Cette stratégie réduit considérablement les besoins en mémoire, permettant d'effectuer l'entraînement sur une seule carte graphique tout en conservant des fichiers d'adaptation légers et portables. Le guide explique comment utiliser les bibliothèques diffusers et accelerate pour configurer l'entraînement, qu'il soit mono ou multi-GPU. La préparation des données implique le téléchargement et le prétraitement de jeux de données similaires à ceux utilisés pour les post-entraînements GR00T Dreams, organisant les vidéos et les légendes textuelles pour l'apprentissage. Pendant la phase d'entraînement, les paramètres du encodeur de texte, du VAE et du DiT (Diffusion Transformer) restent figés. Les adaptations LoRA ou DoRA sont appliquées uniquement sur les projections d'attention et les couches feed-forward du DiT. Le processus utilise une perte d'erreur quadratique moyenne (MSE) basée sur le flux rectifié, où le modèle apprend à prédire la vitesse pour interpoler entre du bruit et les données propres. Les images des deux premières séquences servent de conditionnement, tandis que la perte est calculée sur les autres trames. L'optimisation repose sur l'algorithme AdamW avec un planificateur de taux d'apprentissage linéaire, et les poids de l'adaptateur sont sauvegardés sous forme de fichiers .safetensors. Une fois l'entraînement terminé, généralement après cent époques, le modèle adapté peut être utilisé pour générer des trajectoires robotiques synthétiques. La procédure d'inférence charge le modèle de base et applique les poids de l'adaptateur, optionnellement fusionnés pour éliminer tout délai d'inférence. Pour l'évaluation, les chercheurs utilisent trois métriques clés : l'erreur de Sampson pour mesurer la cohérence géométrique des mouvements, un jugement par IA (LLM) pour vérifier le respect de la physique, et un autre jugement pour évaluer la conformité aux instructions textuelles. Les résultats expérimentaux montrent que le fine-tuning corrige efficacement les défaillances du modèle de base, telles que la génération de mains humaines inappropriées, les mouvements erratiques et le non-respect des instructions. Les modèles adaptés, qu'ils utilisent LoRA ou DoRA, surpassent significativement le modèle non ajusté sur tous les critères. Le rank 32 améliore la précision de l'exécution des tâches par rapport au rank 8, bien que la cohérence géométrique et la plausibilité physique soient bien capturées même par un rank plus faible. DoRA offre une stabilité supplémentaire à faible rank mais n'est pas strictement nécessaire dans ce contexte, rendant LoRA r=8 un choix optimal pour les contraintes de mémoire. Cette approche ouvre la voie à la génération scalable de données synthétiques pour l'apprentissage des politiques robotiques.
