Sub-JEPA: Regularização Gaussiana de Subespaço para Modelos de Mundo Estáveis de Ponta a Ponta

Resumo

As Arquiteturas Preditivas de Incorporação Conjunta (JEPAs) fornecem uma estrutura simples para aprender modelos do mundo prevendo representações latentes futuras. No entanto, o treinamento de JEPAs está sujeito a um tradeoff entre viés e variância. Sem restrições estruturais suficientes, a variância representacional excessiva faz com que o modelo colapse em soluções triviais. O recente LeWorldModel (LeWM) mostra que esse problema pode ser aliviado simplesmente restringindo as incorporações latentes com um prior Gaussiano isotrópico. No entanto, as representações latentes inerentemente estão em variedades de baixa dimensão dentro de um espaço ambiente de alta dimensão, e impor um prior Gaussiano isotrópico diretamente nesse espaço ambiente introduz um viés excessivamente forte. Neste trabalho, propomos ame, que busca um ponto operacional favorável na fronteira viés-variância aplicando restrições Gaussianas em múltiplos subespaços aleatórios, em vez de no espaço de incorporação original. Esse projeto relaxa a restrição global enquanto preserva seu efeito anti-colapso, levando a um melhor equilíbrio entre estabilidade de treinamento e flexibilidade representacional. Extensos experimentos em quatro ambientes de controle contínuo demonstram que ame supera consistentemente o LeWM com margens muito claras. Nosso método é simples mas eficaz, e serve como uma forte baseline para futuras pesquisas em modelos do mundo baseados em JEPA. fdefinedeeemode O código está disponível em 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.