A Generalization Theory for JEPA-Based World 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.

• The paper establishes that JEPA pretraining is equivalent to rank-k factorization of an action-conditioned transition matrix, enabling spectral analysis of representation quality.
• The paper proves that planning regret is bounded by the square root of pretraining risk for one-step tasks and grows linearly with the planning horizon for deterministic multi-step systems.
• The paper identifies a latent-dimension trade-off: larger representations reduce approximation error but increase sample error, helping explain why latent prediction performs better under noisy observations and long horizons.

Overview

Joint-Embedding Predictive Architectures (JEPAs) learn world models by predicting future latent representations rather than reconstructing raw observations, a paradigm instantiated in systems such as V-JEPA 2, DINO-WM, and LeWorldModel. While prior theoretical work on JEPAs has addressed static representation learning, implicit bias, density estimation under Gaussian regularization, and identifiability, no framework existed that explains how JEPA-pretrained models generalize when latent-space plans are executed in the input space. This paper fills that gap by developing the first generalization theory for JEPA-based world models. The central result connects the JEPA pretraining risk to downstream action-planning regret through a spectral characterization of the pretraining objective, yielding finite-sample bounds and an explicit approximation–sample error trade-off governed by the latent dimension.

Conditional spectral graph formulation

The paper adopts a uniformity-regularized JEPA objective: given observation $x$, action $a$, encoder $f$, and predictor $g$, the population risk combines the squared distance between predicted and target latents with a squared inner-product regularizer over independently sampled pairs, in the spirit of spectral contrastive learning (Radford et al., 2021). The key modeling device is the conditional co-occurrence matrix $M(a)$, whose entries are $w(x,x^+,a) = \mathrm{P}(x,x^+|a)$. Unlike the augmentation-graph adjacency matrices used in contrastive-learning theory, this matrix is asymmetric (transitions are directional), and unlike co-occurrence matrices in autoregressive/masked-pretraining analyses, it is conditioned on the action variable.

The first main theorem establishes an exact equivalence: for normalized embeddings,

$$\mathcal{R}_{\mathrm{JEPA}}(f,g,a) = \|\bar{M}(a) - G(F,a)^\top F\|^2 + \text{const.},$$

where $\bar{M}(a) = D^{-1/2} M(a) D_+^{-1/2}$ is the normalized co-occurrence matrix and $F$, $G(F,a)$ collect the $a$0-weighted embeddings. Consequently, minimizing the JEPA risk is exactly low-rank matrix factorization of the action-conditioned transition structure. This reduction is what makes the subsequent generalization analysis tractable, since it imports the machinery of spectral graph theory into the world-model setting.

From pretraining risk to planning regret

The evaluation protocol reflects the core asymmetry of latent world models: planning is performed via trajectory optimization in latent space (minimizing $a$1 over action sequences), but success is measured by input-space goal-reaching probability $a$2. The regret is defined against the best achievable action sequence.

Two results bound this regret by the pretraining factorization risk:

• Single-step planning: the expected regret satisfies $a$3, where $a$4 depends on the ratio $a$5. Thus any encoder–predictor pair with small pretraining risk yields near-optimal one-step actions.
• Multi-step planning: under the additional assumption that $a$6 is deterministic (each state–action pair induces a unique successor), the $a$7-step regret is bounded as $a$8.

The linear dependence on $a$9 formalizes error accumulation across rollout steps: longer horizons amplify pretraining error proportionally. The determinism assumption is defended as realistic for controlled dynamical systems, though it excludes genuinely stochastic environments from the multi-step guarantee — a substantive restriction on applicability. Both results also require the balanced-marginal condition $f$0 to be equal across actions.

Finite-sample bounds and the latent-dimension trade-off

The approximation error follows directly from the Eckart–Young–Mirsky theorem: the optimal population risk equals the tail singular-value energy $f$1 of $f$2, so increasing the latent dimension $f$3 monotonically reduces approximation error. The sample error is controlled via Rademacher complexity of the composed class $f$4 and of $f$5 itself; under boundedness assumptions ($f$6), the complexity prefactors grow polynomially in $f$7 (terms such as $f$8), so sample error increases with $f$9.

Combining these yields the finite-sample planning bound, which decomposes into an approximation error term $g$0 and a sample error term involving Rademacher complexities and a $g$1 concentration term. The trade-off has two notable consequences:

1. There exists an optimal intermediate latent dimension that balances discarded transition information against hypothesis-class complexity.
2. Input-level predictive models appear as the degenerate case $g$2 with identity encoder: zero approximation error but maximal sample error, since all observation dimensions — including nuisance and unpredictable ones — must be modeled.

This provides a formal articulation of why latent prediction should outperform reconstruction precisely when observations contain substantial task-irrelevant variation.

Empirical validation

The experiments use a synthetic 2D point-mass control environment whose observations mix task-relevant state features, nonlinear nuisance projections, and unpredictable noise dimensions. Both paradigms share a JEPA-style backbone (encoder, GRU dynamics predictor, decoder); the latent variant trains with MSE against stop-gradient target encodings plus variance regularization, while the input variant decodes predictions to observation space. Downstream evaluation uses CEM-based receding-horizon goal reaching, scored by ground-truth position success rate.

The empirical pattern matches the theory's regime analysis rather than showing universal superiority of either paradigm. For short tasks (1–5 required steps), both methods achieve near-perfect success regardless of noise. At moderate difficulty (10 steps), input-level prediction remains strong at low noise but latent prediction catches up as noise grows. In long-horizon settings (15–25 steps), the gap widens sharply: at 25 required steps, input-level prediction degrades steeply with noise whereas latent-level prediction maintains substantially higher success rates. The authors interpret this through the lens of the bounds — latent objectives are less affected by unpredictable input dimensions (smaller effective approximation error under noise), and iterative rollouts compound per-step noise, amplifying the advantage. Notably, the paper concedes that latent prediction is not universally better: when task-relevant dynamics are easily recoverable and noise is low, input reconstruction suffices or excels.

Limitations and open questions

Several restrictions qualify the results. The multi-step bound requires deterministic transition matrices, excluding stochastic dynamics common in real-world control. The balanced-marginal assumption across actions is imposed rather than derived. The analysis adopts a specific uniformity regularizer inspired by spectral contrastive loss, whereas practical JEPAs rely on stop-gradient, EMA targets, or variance/covariance regularization (e.g., VICReg-style terms); whether the equivalence and bounds transfer to those mechanisms is not established. The sample-error constants scale polynomially in $g$3 and $g$4, so the bounds are qualitative rather than tight. Finally, validation is confined to a fully controlled synthetic environment; no experiments on vision-based benchmarks (e.g., V-JEPA 2 or DINO-WM settings) are reported, leaving open whether the predicted noise-robustness and horizon effects manifest at scale. An additional open question is extending the framework beyond fixed-action conditioning to continuous or temporally extended action abstractions.

Conclusion

The paper supplies the first end-to-end generalization account of JEPA-based world models: JEPA pretraining is equivalent to rank-$g$5 factorization of an action-conditioned co-occurrence matrix, pretraining error upper-bounds planning regret linearly in the horizon, and the resulting finite-sample bound exposes an inherent approximation–sample trade-off in the latent dimension. This trade-off cleanly explains when latent predictive models dominate input-level ones — high observation noise and long planning horizons — a prediction confirmed qualitatively in controlled synthetic experiments. The framework offers a foundation for analyzing broader classes of predictive world-model architectures, subject to the determinism, regularization, and scalability caveats noted above.