Fast-LeWM : La prédiction parallèle de préfixes d'actions réduit les coûts de planification des modèles de monde latent

Les architectures prédictives à encodage conjoint (JEPA), y compris le récent LeWorldModel (LeWM), ont été saluées comme une voie prometteuse pour les modèles de monde visuel sans reconstruction. Mais lorsqu'il s'agit de planification visuelle, elles se sont heurtées à un mur. Le problème : LeWM évalue les séquences d'actions candidates en appliquant de manière répétée un modèle de transition latente locale en une étape de manière autorégressive. Ce déroulement est coûteux, et à mesure que l'horizon s'allonge, les erreurs s'accumulent rapidement.

Dans un nouvel article intitulé Fast LeWorldModel, les chercheurs décrivent une solution. Au lieu de dérouler les prédictions étape par étape, Fast-LeWM prédit le résultat de préfixes d'actions entiers en une seule fois. L'idée : encoder les préfixes d'une séquence d'actions candidates et prédire tous les états latents futurs atteints après l'exécution de ces préfixes en parallèle.

Comment fonctionne Fast-LeWM

Étant donné l'état latent actuel et une séquence d'actions candidates, Fast-LeWM encode ses préfixes, les premières actions, et prédit les états latents futurs qui seraient atteints après l'exécution de ces préfixes. En faisant des préfixes d'actions l'unité de prédiction de base, le modèle apprend directement les effets cumulés des actions sur plusieurs horizons, plutôt que de ne modéliser que des transitions d'état en une étape.

Pendant la planification, le prédicteur peut utiliser le dernier jeton de préfixe de la séquence d'actions encodée pour évaluer l'état latent futur correspondant sans passer explicitement par chaque état intermédiaire imaginé. Cette conception élimine le goulot d'étranglement autorégressif et réduit l'exposition à l'accumulation d'erreurs qui affecte les prédictions à long horizon.

Gains de performance

Sur plusieurs tâches, Fast-LeWM améliore les taux de réussite moyens par rapport à LeWM tout en réduisant considérablement le temps de planification. La méthode atteint une perte latente en boucle ouverte plus faible, dont la croissance devient significativement plus lente à mesure que l'horizon de déroulement augmente. Cela suggère que la supervision au niveau du préfixe force le modèle à apprendre comment les états évoluent continuellement sous différents préfixes d'actions, plutôt que de seulement modéliser des transitions en une étape.

Les auteurs rapportent que Fast-LeWM conserve les avantages sans reconstruction des JEPA tout en surmontant leur principale limitation pour la planification : le coût de calcul des déroulements autorégressifs. En parallélisant la prédiction, la méthode ouvre la voie à la planification visuelle en temps réel en robotique, navigation autonome et autres domaines nécessitant une prise de décision rapide sur de longs horizons.

Implications pour les modèles de monde visuel

Ce travail représente une avancée pratique pour rendre les modèles de monde latent plus efficaces pour la planification. Bien que les JEPA aient montré de bonnes performances en apprentissage de représentations et en modélisation sans reconstruction, leur utilisation en planification a été limitée par la nature séquentielle des modèles de transition latente. Fast-LeWM comble cette lacune avec un changement architectural simple mais efficace.

Les travaux futurs pourraient explorer l'extension de l'approche aux espaces d'actions continus ou son intégration avec des algorithmes d'apprentissage par renforcement basés sur un modèle. L'article suggère également que le paradigme de prédiction par préfixe pourrait être appliqué à d'autres composants autorégressifs dans les modèles de monde au-delà de la planification visuelle.