A Generalization Theory for JEPA-Based World Models

2026-06-25

AI Summary

This work addresses the lack of theoretical understanding regarding the generalization capabilities of JEPA-based world models. By formulating JEPA pretraining as a conditional spectral graph learning problem, we characterize its learning objective through low-rank decomposition of an action-conditioned co-occurrence matrix and establish, for the first time, a theoretical link between pretraining error and downstream planning regret. We derive finite-sample generalization bounds for JEPA world models, revealing an intrinsic trade-off governed by the latent dimension between approximation error and sampling error. Our analysis provides the first theoretical foundation elucidating both the advantages and limitations of latent-space predictive models.

Abstract

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.

Problem

• JEPA
• world models
• generalization theory
• latent predictive models
• spectral graph learning

Innovation

• JEPA
• generalization theory
• world models
• latent predictive modeling
• spectral graph learning