Sub-JEPA : Régularisation Gaussienne de Sous-espace pour des Modèles de Monde Stables de Bout en Bout

Résumé

Les architectures prédictives par embedding conjoint (JEPAs) fournissent un cadre simple pour apprendre des modèles du monde en prédisant des représentations latentes futures. Cependant, l'entraînement des JEPAs est soumis à un compromis biais-variance. Sans contraintes structurelles suffisantes, une variance représentationnelle excessive conduit le modèle à s'effondrer vers des solutions triviales. Le récent LeWorldModel (LeWM) montre que ce problème peut être atténué en contraignant simplement les embeddings latents par un a priori gaussien isotrope. Cependant, les représentations latentes résident naturellement sur des variétés de faible dimension dans un espace ambiant de haute dimension, et imposer un a priori gaussien isotrope directement dans cet espace ambiant introduit un biais trop fort. Dans ce travail, nous proposons ame, qui cherche un point de fonctionnement favorable sur la frontière biais-variance en appliquant des contraintes gaussiennes dans de multiples sous-espaces aléatoires plutôt que dans l'espace d'embedding original. Cette conception assouplit la contrainte globale tout en préservant son effet anti-effondrement, conduisant à un meilleur équilibre entre stabilité de l'entraînement et flexibilité de la représentation. Des expériences approfondies menées dans quatre environnements de contrôle continu montrent que ame surpasse systématiquement LeWM avec des marges très nettes. Notre méthode est simple mais efficace, et sert de baseline solide pour les futures recherches sur les modèles du monde basés sur JEPA.fdefinedeeemodeLe code est disponible à l'adresse https://github.com/intcomp/Sub-JEPA.

English

Joint-Embedding Predictive Architectures (JEPAs) provide a simpleframework for learning world models by predicting future latent representations.However, JEPA training is subject to a bias-variance tradeoff.Without sufficient structural constraints, excessive representationalvariance causes the model to collapse to trivial solutions.The recent LeWorldModel (LeWM) shows that this issue can be alleviated bysimply constraining latent embeddings with an isotropic Gaussian prior.However, latent representations inherently lie on low-dimensional manifoldswithin a high-dimensional ambient space, and enforcing an isotropic Gaussianprior directly in this ambient space introduces an overly strong bias.In this work, we propose ame, which seeks a favorable operatingpoint on the bias-variance frontier by applying Gaussian constraints inmultiple random subspaces rather than in the originalembedding space.This design relaxes the global constraint while preserving itsanti-collapse effect, leading to a better balance between trainingstability and representation flexibility.Extensive experiments across fourcontinuous-control environments demonstrate that consistentlyoutperforms LeWM with very clear margins.Our method is simple yet effective, and serves as a strong baseline for future JEPA-based world model research.fdefinedeeemodeThe code is available at https://github.com/intcomp/Sub-JEPA.