JEPA-Style World Models

A world model answers one question: if I do this, what happens next? The obvious way to build one is to predict the future frame by frame, pixel by pixel. The Joint-Embedding Predictive Architecture, or JEPA, which Yann LeCun proposed in 2022 as the world-model module of an autonomous agent, argues that this is the wrong target. Much of the next frame, such as the ripple of water or the flutter of leaves, cannot be predicted and does not matter for acting.

A JEPA predicts in representation space instead. It encodes the present and the future, and learns to predict the encoding of the future from the encoding of the present, so the encoder is free to drop whatever is unpredictable. That one change brings its own problem, since an encoder that maps everything to the same vector makes prediction trivially perfect. It has also produced a family of models: I-JEPA for images, V-JEPA for video, and V-JEPA 2, whose action-conditioned version plans robot-arm motions without generating a single pixel.

---

1. The Architecture

Take a pair of related inputs: a context \(x\) and a target \(y\). They might be the visible and hidden parts of an image, the first and second halves of a video clip, or the current camera frame and the frame after an action. A JEPA has three networks, each written with its weights as a subscript:

• the context encoder \(f_\theta\), with weights \(\theta\), which turns \(x\) into an embedding: a vector, or one vector per image or video patch;
• the target encoder \(f_{\bar\theta}\), the same architecture with its own weights \(\bar\theta\), which turns \(y\) into an embedding the same way;
• the predictor \(g_\phi\), with weights \(\phi\), which takes the context embedding and guesses the target embedding.

Training runs all three and compares the guess with the real target embedding:

\[s_x = f_\theta(x), \qquad s_y = f_{\bar\theta}(y), \qquad \hat s_y = g_\phi(s_x, z), \qquad \mathcal{L} = D\big(\hat s_y,\; s_y\big).\]

Reading left to right: \(s_x\) is the context embedding and \(s_y\) the target embedding; \(\hat s_y\) (“s-y hat”) is the predictor’s guess of \(s_y\); and the loss \(\mathcal{L}\) is the distance \(D\) between guess and target, squared L2 in I-JEPA and L1 in V-JEPA. The remaining input \(z\) tells the predictor which target to produce: the position of a hidden block, an action the agent is about to take, or a latent variable standing in for whatever the context cannot determine. The loss trains \(\theta\) and \(\phi\); how \(\bar\theta\) is set is the subject of section 3.

LeCun frames this as an energy-based model. The energy of a pair is the prediction error, minimized over the latent:

\[E(x, y) = \min_{z}\; D\big(g_\phi(f_\theta(x), z),\; f_{\bar\theta}(y)\big).\]

Compatible pairs, a scene and a plausible continuation of it, should get low energy; incompatible pairs should get high energy. Training shapes this landscape, and planning searches it: find the actions whose predicted outcome has low energy with respect to a goal.

Two older families of self-supervised learning bracket the design. Generative models, such as masked autoencoders and video predictors, reconstruct \(y\) itself, so every pixel counts. Joint-embedding methods, such as SimCLR and DINO, pull together the embeddings of two augmented views of the same image; nothing tells the model where or when the second view comes from, so they learn invariance to the augmentations rather than prediction. A JEPA keeps the embedding-space loss of the second family and the conditional prediction of the first.

2. Why Not Predict Pixels?

Consider the simplest unpredictable future. A ball sits in the middle of the frame, and in the next frame it will have moved left or right with equal probability. Say it covers \(A\) pixels, and train a predictor with squared error in pixel space. A sharp guess, the ball on one side, is right half the time and otherwise wrong on \(2A\) pixels, since the ball is missing where it is and drawn where it is not. Its expected error is

\[\tfrac12 \cdot 0 + \tfrac12 \cdot 2A = A.\]

The blurry guess, half a ball on each side, is off by \(\tfrac12\) on \(2A\) pixels whichever future happens:

\[2A \cdot \big(\tfrac12\big)^2 = \tfrac{A}{2}.\]

Squared error rewards the blur by a factor of two. In general the minimizer of expected squared error is the conditional mean \(\mathbb{E}[\,y \mid x\,]\), the average of every future consistent with the context \(x\), and when several futures are possible, their average is usually not one of them.

Now make the target an embedding. If the encoder describes the frame by what is predictable, say “one ball, one step from where it was, same color”, both futures get the same embedding and the predictor can match it exactly. The encoder has learned to drop the direction, which it cannot predict. If the direction matters, for instance to an agent that has to catch the ball, it moves into the latent variable: the predictor outputs one sharp embedding for each value of \(z\), and both futures get low energy instead of sharing a blur.

That is the JEPA bet in a sentence: spend model capacity on what is predictable, and let the representation or the latent absorb the rest. Real video is full of the unpredictable, such as water, foliage, crowds and sensor noise, so the saving is large.

3. The Collapse Problem

The loss in section 1 has a trivial minimum. If both encoders output the same constant vector for every input, the predictor copies it and the loss is zero. Nothing in “predict the embedding of \(y\)” forbids this, because the target is itself learned. Every JEPA needs something that keeps the embeddings informative, and there are four main answers.

Contrastive Losses

Methods like SimCLR also push the embeddings of unrelated inputs apart, so a constant encoder is penalized. They work, but they need many negative examples per batch, and LeCun argues that the number needed grows quickly with the dimension of the embedding. JEPA work has mostly moved to the three non-contrastive options below.

Variance and Covariance Regularization

VICReg (Bardes, Ponce and LeCun, 2022) penalizes the collapse directly. Stack a batch of \(n\) embeddings as the rows of a matrix \(Z\), so that it has \(d\) columns, one per embedding dimension, and write \(Z_{:,j}\) for column \(j\). VICReg adds a variance term \(v(Z)\), which wants every dimension to keep a standard deviation of at least \(\gamma\) across the batch, and a covariance term \(c(Z)\), which wants different dimensions to be uncorrelated:

\[v(Z) = \frac{1}{d}\sum_{j=1}^{d} \max\Big(0,\; \gamma - \sqrt{\operatorname{Var}(Z_{:,j}) + \epsilon}\Big), \qquad c(Z) = \frac{1}{d}\sum_{i \ne j} C(Z)_{ij}^2.\]

Here \(\operatorname{Var}(Z_{:,j})\) is the variance of dimension \(j\) over the batch; \(C(Z)\) is the \(d \times d\) covariance matrix of the columns, so \(C(Z)_{ij}\) is the covariance between dimensions \(i\) and \(j\); and \(\epsilon = 10^{-4}\) keeps the square root well behaved near zero. Both terms are applied to the embeddings on each side of the loss, \(Z\) and \(Z'\). With the default weights (25 on the prediction term, 25 on variance, 1 on covariance) and \(\gamma = 1\), a collapsed encoder has zero variance in every dimension, so \(v(Z) = v(Z') = 1 - \sqrt{\epsilon} = 0.99\) and the variance penalty totals \(25 \times (0.99 + 0.99) = 49.5\). Features with unit standard deviation pay nothing. The covariance term blocks a subtler failure: copying one informative feature into every dimension.

Stop-Gradient and a Moving-Average Target

I-JEPA and V-JEPA use the recipe from BYOL instead. The loss is never backpropagated into the target encoder; its weights just track the context encoder’s,

\[\bar\theta \leftarrow \tau\,\bar\theta + (1 - \tau)\,\theta,\]

where \(\tau\), the momentum, is ramped from 0.996 to 1.0 over training. At 0.996, each step keeps 99.6% of the old target weights and mixes in 0.4% of the current context weights, which amounts to averaging over roughly the last 250 steps. The target moves slowly and the predictor has to chase it. In practice this is enough to prevent collapse, although the theory of why is still incomplete.

Pushing Embeddings Toward a Gaussian

LeJEPA (Balestriero and LeCun, 2025) replaces these heuristics with a single regularizer, SIGReg, which pushes the distribution of embeddings toward an isotropic Gaussian, the shape they show is optimal for downstream prediction. Testing whether a high-dimensional distribution is Gaussian is expensive, so SIGReg projects the embeddings onto random directions and applies a one-dimensional test of normality along each, which keeps the cost linear. There is no stop-gradient, no teacher network and no schedule, and one weight trades the regularizer against prediction. LeWorldModel uses the same kind of Gaussian regularizer to train a world model end to end from pixels (section 5).

4. I-JEPA: Predicting Hidden Parts of an Image

I-JEPA (Assran et al., 2023) is the image version. A Vision Transformer with 16-pixel patches cuts a 224 × 224 image into a 14 × 14 grid. Each training example samples four target blocks, each covering 15–20% of the image with an aspect ratio between 0.75 and 1.5, and one context block covering 85–100% of the image, from which every patch that falls inside a target is removed.

The target encoder sees the whole image, and the targets are its output embeddings at the target patches. The context encoder sees only the context patches. The predictor, a narrower transformer, receives the context embeddings plus one placeholder token per target patch carrying that patch’s position, and fills in the missing embeddings. The loss averages squared error over the \(M = 4\) blocks,

\[\mathcal{L} = \frac{1}{M}\sum_{i=1}^{M}\;\sum_{j \in B_i} \big\|\hat s_{y_j} - s_{y_j}\big\|_2^2,\]

where \(B_i\) is the set of patches in block \(i\), and \(\hat s_{y_j}\) and \(s_{y_j}\) are the predicted and the target embedding of patch \(j\), each a single vector.

Two choices make this work. The targets are large blocks rather than scattered patches, so predicting them takes knowing what object is there, not interpolating texture from the neighbors. And the loss lives in embedding space, so the model is never asked to reproduce the exact fur of a dog, only what the target encoder makes of it. I-JEPA needs no hand-designed augmentations either, and it is cheap: the paper trains a ViT-H/14 on ImageNet with 16 A100 GPUs in under 72 hours.

5. From Video to Action

V-JEPA (Bardes et al., 2024) applies the same recipe to video. The clip is cut into space-time patches, and the hidden regions are tubes, the same spatial blocks masked in every frame, so the model cannot solve the task by copying from the previous frame. The loss is L1 between predicted and target embeddings, and the encoder is evaluated frozen, with only a light probe trained on top for each downstream task.

V-JEPA 2 (Meta, 2025) scales this to a ViT-g encoder of about a billion parameters, pretrained on more than a million hours of internet video. A second stage turns it into a world model. With the encoder frozen, an action-conditioned predictor, V-JEPA 2-AC, is trained on under 62 hours of unlabeled robot video from the DROID dataset to predict the embedding of the next frame from the current embedding and the robot’s action, a change in end-effector pose and gripper state.

Planning in Embedding Space

To reach a goal given as an image \(x_g\), the robot encodes it into a goal embedding \(s_g = f_\theta(x_g)\), then searches over sequences of \(T\) future actions, \(a_{1:T} = (a_1, \dots, a_T)\), for the one whose predicted outcome lands closest to the goal:

\[a^\star_{1:T} = \arg\min_{a_{1:T}}\; \big\| g_\phi(s_t, a_{1:T}) - s_g \big\|_1.\]

Here \(s_t\) is the embedding of the current camera frame, \(g_\phi(s_t, a_{1:T})\) applies the predictor \(T\) times to roll that embedding forward under the chosen actions, \(\|\cdot\|_1\) is the L1 distance, and \(a^\star_{1:T}\) is the best sequence found. \(T\) is the planning horizon, how many steps ahead the robot looks. The search uses the cross-entropy method:

1. Sample many candidate action sequences from a Gaussian.
2. Roll each one out with the predictor, entirely in embedding space, and score it by its distance to the goal.
3. Refit the Gaussian to the best-scoring sequences and repeat for a few iterations.
4. Execute only the first action, observe the new frame, and plan again.

The last step, model-predictive control, is what makes an imperfect model usable: errors never compound for long, because every step starts again from a real observation. With this loop, V-JEPA 2-AC picked and placed objects with Franka arms in two labs whose scenes it had never seen, without task-specific training or rewards.

Frozen or End to End?

V-JEPA 2-AC keeps its encoder frozen, and so does DINO-WM, which plans on top of frozen DINOv2 patch features: the representation is learned once from large data, and only the predictor learns the dynamics. LeWorldModel (Maes, Le Lidec, Scieur, LeCun and Balestriero) goes the other way and trains the encoder and predictor together from raw pixels, which has historically been unstable. The Gaussian regularizer from section 3 keeps it from collapsing, with two loss terms and one hyperparameter where the only earlier end-to-end method needed six. The model has about 15 million parameters, trains on a single GPU in a few hours, and plans up to 48 times faster than world models built on large pretrained encoders, on small control tasks such as pushing a T-shaped block and navigating between two rooms.

6. How JEPA Differs from Other World Models

All world models predict the future from the present; they differ in what they predict and what trains them.

Pixel models are the easiest to inspect, because you can watch what they imagine, but they spend most of their capacity on appearance. MuZero sits at the other extreme: its latent state only has to predict reward and value, so it learns nothing the task does not reward. A JEPA sits between them. It learns from observation alone, like the generative models, but keeps only what is predictable, and it needs no reward to do so.

Limitations

• There is no decoder, so there is nothing to look at. Debugging a JEPA world model means probing its embeddings, or training a separate decoder after the fact.
• What the encoder drops as unpredictable may be exactly what a task needs. Contact geometry is a good example: a grasp planner cares about surface normals to within a few degrees (see Friction Cone and Antipodal Grasps), and a representation trained on internet video has no particular reason to keep them.
• Goals have to live in the same space, which in current systems usually means an image of the goal state rather than an instruction.
• Long horizons remain hard, since errors still compound across predictor steps. LeCun’s proposed answer is a hierarchical JEPA that predicts coarse, abstract states far ahead and detailed ones only in the near term.

Intuitively: Predicting the Gist

A goalkeeper facing a penalty does not imagine the ball’s seams, the grass it kicks up or the crowd behind it. They predict one thing, where the ball will be when it reaches the line, and they predict it from the few cues that carry that information: the kicker’s hips, the angle of the run-up. Everything else is noise they have learned to ignore.

A pixel world model tries to imagine the whole scene and, when it is unsure, paints the average of everything that might happen. A JEPA learns what the goalkeeper learned: which parts of the future can be predicted, and how to predict only those. The hard part, as section 3 showed, is making sure it does not decide that the easiest thing to predict is nothing at all.