A Generalization Theory for JEPA-Based World Models
Provides theoretical foundations for JEPA-based world models, which are important for representation learning and planning in AI, but the work is primarily theoretical and does not include empirical validation.
This paper develops the first generalization theory for JEPA-based world models, showing that the JEPA objective is equivalent to low-rank factorization of an action-conditioned co-occurrence matrix and establishing a finite-sample generalization bound that reveals a trade-off between approximation and sample errors with respect to latent dimension.
Joint Embedding Predictive Architectures (JEPAs) have recently emerged as a promising paradigm for world modeling by learning predictive dynamics in a latent space rather than generating future observations at the input level. Despite their empirical success, the theoretical understanding of JEPA-based world models remains limited. In this paper, we develop the first generalization theory for JEPA-based world models. We formulate JEPA pretraining as a conditional spectral graph learning problem and show that the JEPA objective is equivalent to a low-rank factorization of an action-conditioned co-occurrence matrix. Building on this characterization, we establish a connection between JEPA pretraining error and downstream planning regret, leading to a finite-sample generalization bound for JEPA-based world models. Our analysis reveals an inherent trade-off between approximation and sample errors with respect to the latent dimension, providing theoretical insights into the advantages and limitations of latent predictive models compared with input-level predictive approaches.