Sub-JEPA: Subruimte Gaussiaanse Regularisatie voor Stabiele End-to-End Wereldmodellen

Samenvatting

Joint-Embedding Predictive Architectures (JEPA's) bieden een eenvoudig raamwerk voor het leren van wereldmodellen door toekomstige latente representaties te voorspellen. JEPA-training is echter onderhevig aan een bias-variantie-afweging. Zonder voldoende structurele beperkingen leidt overmatige representatievariantie ertoe dat het model instort tot triviale oplossingen. Het recente LeWorldModel (LeWM) laat zien dat dit probleem kan worden verholpen door eenvoudigweg latente embeddings te beperken met een isotrope Gaussiaanse prior. Latente representaties liggen echter inherent op laagdimensionale manifolds binnen een hoogdimensionale omgevingsruimte, en het rechtstreeks opleggen van een isotrope Gaussiaanse prior in deze omgevingsruimte introduceert een te sterke bias. In dit werk stellen we Sub-JEPA voor, dat een gunstig werkpunt op de bias-variantiegrens zoekt door Gaussiaanse beperkingen toe te passen in meerdere willekeurige deelruimten in plaats van in de oorspronkelijke embeddingruimte. Dit ontwerp versoepelt de globale beperking terwijl het anti-instorteffect behouden blijft, wat leidt tot een betere balans tussen trainingsstabiliteit en representatieflexibiliteit. Uitgebreide experimenten in vier continue besturingsomgevingen tonen aan dat Sub-JEPA consequent en met zeer duidelijke marges beter presteert dan LeWM. Onze methode is eenvoudig maar effectief en dient als een sterke basislijn voor toekomstig JEPA-gebaseerd wereldmodelonderzoek. De code is beschikbaar op 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.