Résumé technique : Sub-JEPA

Énoncé du problème

Les architectures prédictives à encodage conjoint (JEPAs) offrent un cadre pour apprendre des modèles du monde en prédisant des représentations latentes futures sans reconstruire les pixels bruts. Cependant, l'entraînement de ces modèles de bout en bout fait face à un compromis biais-variance critique. Sans contraintes suffisantes, l'encodeur s'effondre vers des solutions triviales (mappant toutes les entrées sur des embeddings identiques) ; inversement, des contraintes excessives suppriment la richesse des représentations apprises.

Des travaux récents, LeWorldModel (LeWM), ont résolu ce problème en appliquant un prior gaussien isotrope aux embeddings latents, permettant un entraînement stable de bout en bout. Cependant, les auteurs identifient une limitation fondamentale dans cette approche : les tâches de contrôle naturelles résident typiquement sur des variétés de basse dimension au sein d'espaces ambiants de haute dimension. Imposer un prior gaussien isotrope global directement dans l'espace ambiant de haute dimension introduit un biais excessivement fort qui ne correspond pas à la géométrie intrinsèque des dynamiques sous-jacentes, risquant ainsi de nuire aux performances sur des tâches à faible dimensionnalité intrinsèque.

Méthodologie : Sub-JEPA

Les auteurs proposent Sub-JEPA (Régularisation gaussienne par sous-espace), une méthode qui cherche un point de fonctionnement favorable sur la frontière biais-variance en déplaçant les contraintes gaussiennes de l'espace ambiant vers plusieurs sous-espaces de basse dimension.

Mécanisme central

1. Projection sur sous-espaces orthogonaux : Au lieu de régulariser le vecteur latent complet $z$ de dimension $D$, la méthode projette $z$ sur $K$ matrices de projection aléatoires, figées et à lignes orthonormales $\{P_k\}$. Chaque projection mappe le vecteur latent vers un sous-espace de dimension inférieure de dimension $d_s$ (où $d_s \approx D/K$).
   $$z^{(k)} = P_k z$$
2. Régularisation gaussienne multi-sous-espaces : La contrainte gaussienne est appliquée indépendamment dans chacun des $K$ sous-espaces.
   • Dans chaque sous-espace $k$, les embeddings projetés sont ensuite projetés sur $M$ vecteurs unitaires aléatoires.
   • La statistique de normalité d'Epps–Pulley est calculée pour ces projections 1D afin d'imposer la gaussianité.
   • La perte de régularisation totale $L_{reg}$ est la moyenne de ces statistiques sur tous les sous-espaces et toutes les directions aléatoires.
3. Objectif d'entraînement : Le modèle est entraîné de bout en bout en utilisant une perte combinée :
   $$L_{total} = L_{pred}(\hat{z}_{t+1}, z_{t+1}) + \lambda L_{reg}(Z)$$
   où $L_{pred}$ est l'erreur de prédiction de l'état latent (MSE) et $\lambda$ est le poids de régularisation.

Choix de conception

• Projections figées : Les matrices de projection sont initialisées de manière orthogonale et maintenues figées. Cela empêche le régularisateur de s'adapter à la distribution latente évolutive, garantissant que la contrainte reste stable et indépendante.
• Orthogonalité : L'utilisation de projections orthogonales assure que les sous-espaces fournissent des vues non redondantes et équilibrées de la représentation latente, préservant la structure géométrique.

Contributions clés

1. Identification du biais : L'article identifie que le prior gaussien isotrope global dans LeWM introduit un biais excessif pour les tâches à faible dimensionnalité intrinsèque, ne correspondant pas à la géométrie sous-jacente de la tâche.
2. Cadre Sub-JEPA : Les auteurs proposent une méthode qui assouplit le prior global en appliquant une régularisation gaussienne sur plusieurs vues projetées aléatoires, de basse dimension et à lignes orthonormales. Cela préserve l'effet anti-effondrement tout en permettant à l'espace latent de mieux se conformer à la dimensionnalité intrinsèque des dynamiques.
3. Validation empirique : Des expériences extensives sur quatre environnements de contrôle continu démontrent que Sub-JEPA surpasse systématiquement LeWM. Les gains de performance sont corrélés à des réductions du rang effectif des embeddings appris, confirmant que la méthode respecte mieux la dimensionnalité intrinsèque.

Résultats expérimentaux

La méthode a été évaluée sur quatre benchmarks de contrôle continu : Two-Room (faible dimensionnalité intrinsèque), Reacher, PushT, et OGB-Cube (3D visuellement riche).

• Performance de planification : Sub-JEPA a systématiquement surpassé LeWM dans tous les environnements. L'amélioration a été la plus marquée sur Two-Room, où le prior gaussien en espace complet de LeWM était le plus inadapté à la faible dimensionnalité de la tâche. Sub-JEPA a également surpassé PLDM (une baseline multi-objectifs) et DINO-WM (une baseline à encodeur figé) sur plusieurs tâches sans recourir à des backbones visuels pré-entraînés.
• Analyse du rang effectif : Une corrélation directe a été observée entre la réduction du rang effectif (de LeWM à Sub-JEPA) et l'augmentation du taux de réussite en planification. Les tâches avec des compressions de rang plus importantes (Two-Room, OGB-Cube) ont montré les plus grands gains de performance.
• Décodabilité physique : Des expériences de sondage sur l'environnement PushT ont montré que les représentations Sub-JEPA préservaient les informations d'état physique (position de l'agent, position du bloc) aussi bien ou mieux que LeWM, avec une décodabilité non linéaire améliorée.
• Qualité des trajectoires latentes :
  • Cohérence temporelle : Les visualisations ont montré que Sub-JEPA produisait des trajectoires latentes plus cohérentes temporellement par rapport à l'évolution moins organisée dans LeWM.
  • Redressement des trajectoires : Sub-JEPA a atteint une plus grande « rectitude » dans les vecteurs de vitesse latente, indiquant des dynamiques plus lisses et évoluant plus linéairement, ce qui est crucial pour la planification.
  • Stabilité à long horizon : Dans les déroulements en boucle ouverte, Sub-JEPA a maintenu la géométrie de la scène et a montré moins de dérive spatiale sur de longs horizons par rapport à LeWM.

Importance et affirmations

L'article affirme que Sub-JEPA offre une solution simple mais efficace au compromis biais-variance dans les modèles du monde basés sur JEPA. En déplaçant la régularisation gaussienne vers des sous-espaces, la méthode :

• Assouplit le prior isotrope global sans perdre les bénéfices anti-effondrement.
• Correspond mieux à la dimensionnalité intrinsèque des dynamiques sous-jacentes, en particulier pour les tâches où l'espace d'état est de basse dimension.
• Sert de baseline solide pour les futures recherches sur les modèles du monde basés sur JEPA, démontrant que les contraintes structurelles peuvent être simplifiées et rendues plus flexibles sans sacrifier la stabilité.

Les auteurs soulignent que leur approche conserve la simplicité de la recette d'entraînement de LeWM (n'utilisant que la perte de prédiction et un régularisateur) tout en introduisant un prior structurel plus flexible qui génère des performances de planification supérieures et une meilleure qualité de représentation.