52  World Models and JEPA

TipTL;DR

JEPA (Assran et al., 2023; Bardes et al., 2024) predicts a representation of missing input, not raw pixels or tokens. World models extend the same prediction target to actions and time, enabling planning by simulating latent futures. Both are proposed as alternatives to autoregressive token/pixel prediction for planning-heavy tasks (LeCun, 2022).

Depends on: Self-Supervised & Contrastive Learning

52.1 Why this matters

Every generative model in this book predicts in the space it was given: an LLM predicts the next token, diffusion predicts the noise added to a pixel grid. Embodied AI applied that same recipe to actions and closed with an open question: is token-level (or action-level) prediction the right training objective for an agent that must plan over a long horizon? This note covers two proposed alternatives, JEPA’s representation-level prediction and world models built on it, and the specific critique of autoregression that motivates them (LeCun, 2022).

52.2 The mechanism

52.2.1 Predicting representations, not pixels

I-JEPA (Assran et al., 2023) and V-JEPA (Bardes et al., 2024) apply one change to the self-supervised recipe already covered in Self-Supervised Learning: the prediction target is a representation, not a reconstruction.

Component	Role
Context encoder	Encodes the visible patches of an image or video clip
Target encoder	An exponential moving average of the context encoder (same mechanism as DINO); encodes the masked target patches
Predictor	A small network mapping the context encoding to a predicted target-patch representation
Loss	Distance between the predictor’s output and the target encoder’s actual output, for the masked patches only

Written as a loss, with \(s_\theta\) the context encoder, \(\bar{s}_\theta\) the (stop-gradient) target encoder, and \(p_\phi\) the predictor:

\[ \mathcal{L} = \big\lVert\, p_\phi\big(s_\theta(x_{\text{context}})\big) - \bar{s}_\theta(x_{\text{target}}) \,\big\rVert^2. \]

No pixels are reconstructed anywhere in this loss. V-JEPA applies the identical objective to masked spatio-temporal patches in video.

52.2.2 Energy-based models

JEPA is one instance of an energy-based model: a function \(E(x, y)\) trained to output low energy for compatible pairs and high energy otherwise, without necessarily normalizing to a probability distribution over \(y\). The central design problem for any energy-based model is preventing collapse: a constant \(E(x,y) = 0\) trivially minimizes the loss on positive pairs. JEPA avoids collapse architecturally rather than with explicit negative examples:

• The predictor’s limited capacity makes a constant output a poor fit to varied targets.
• The stop-gradient on the target encoder removes the trivial solution of both encoders collapsing to the same constant together.

This contrasts with contrastive methods (SimCLR, MoCo), which prevent collapse using explicit negative pairs instead.

52.2.3 World models for planning

A world model extends the same prediction target to dynamics:

\[ z_{t+1} = f(z_t, a_t), \]

a learned transition function over latent states \(z_t\) and actions \(a_t\), with no pixel generation at any step. Planning becomes a search over action sequences, simulated entirely in latent space using \(f\), scored against a goal, rather than executed in the real or rendered environment to observe the outcome.

52.2.4 The critique of autoregressive prediction

Two specific limitations of token-by-token (or pixel-by-pixel) generation have been raised as motivation for representation-level prediction (LeCun, 2022):

Limitation	Description
Error accumulation	Each generated token is a chance to deviate; small errors compound over a long sequence, so the probability of a fully coherent path shrinks with length
Wasted capacity	Predicting exact low-level detail (the precise pixel noise in a frame, the exact phrasing of a sentence) spends capacity on content that is often genuinely unpredictable and irrelevant to higher-level structure

JEPA’s representation-level target sidesteps the second limitation directly: the encoder discards low-level, unpredictable detail before prediction happens, rather than forcing the model to predict it.

52.2.5 What to observe

• JEPA changes the prediction target (representations, not pixels), not the underlying self-supervised recipe; it reuses the context/target-encoder structure from DINO.
• A world model is representation-level prediction applied to actions and time, the same extension Embodied AI’s closing question raised.
• The critique in (LeCun, 2022) targets the training objective (next-token prediction), not the Transformer architecture; JEPA models can and do use Transformer blocks internally.

52.3 Outlook

JEPA and world models remain far behind autoregressive LLMs and VLA models in scale and demonstrated real-world capability as of this writing. Whether representation-level prediction displaces next-token prediction for planning-heavy agents is an open research question.

NoteKey takeaway

JEPA predicts representations instead of raw signals, the same principle as every self-supervised method in this book, applied specifically to address error accumulation and wasted capacity in autoregressive generation. World models extend this to actions and time, the most direct alternative to next-token-style planning covered here.

← Back to: Embodied AI & VLA

LeCun, Y. (2022). A path towards autonomous machine intelligence. OpenReview.

Assran, M., Duval, Q., Misra, I., Bojanowski, P., Vincent, P., Rabbat, M., LeCun, Y., & Ballas, N. (2023). Self-supervised learning from images with a joint-embedding predictive architecture. Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR).

Bardes, A., Garrido, Q., Ponce, J., Chen, X., Rabbat, M., LeCun, Y., Assran, M., & Ballas, N. (2024). Revisiting feature prediction for learning visual representations from video. arXiv Preprint arXiv:2404.08471.