Predictive Embeddings, JEPA, and Energy-Based Models › Joint-Embedding Predictive Architectures

Yann LeCun introduced JEPA in the 2022 position paper as a broad architecture for prediction in representation space, alongside generative and other joint-embedding architectures. That generic formulation does not prescribe an EMA target or stop-gradient. The concrete recipe developed here follows prominent I-JEPA- and V-JEPA-style implementations: a context encoder, an EMA target encoder, a predictor, and detached target embeddings.

The components are:

• Context encoder $f_\theta$: maps the observed part of the input to an embedding and is trained by gradient descent. Which encoder is retained for evaluation is method-dependent; BYOL uses the online representation, while an official V-JEPA frozen-evaluation configuration loads the target encoder.
• Target encoder $f_{\bar\theta}$: processes the target input without gradient updates. In masked I-JEPA, the official training implementation processes the full unmasked image before target locations are selected from its output. Its parameters $\bar\theta$ follow $\theta$ via an exponential moving average.
• Predictor $g_\phi$: maps the context embedding (plus optional conditioning, such as an action) to a prediction of the target embedding.
• Stop-gradient: the target embedding is detached from the computation graph before the loss is applied.

For the squared-error instantiation used in this chapter's analysis and toy system, the objective is an embedding-space discrepancy. The loss is not canonical across JEPA implementations: the official I-JEPA training code uses Smooth L1, while the cited official V-JEPA configuration sets loss_exp to $1$, which gives an L1 objective in its training implementation.

$$\mathcal{L}(\theta, \phi) \;=\; \mathbb{E}_{(x, y) \sim \mathcal{D}} \Big\| g_\phi\big(f_\theta(x), c\big) \;-\; \mathrm{sg}\big(f_{\bar\theta}(y)\big) \Big\|_2^2 ,$$

where:

• $\mathrm{sg}(\cdot)$: the stop-gradient operator, which detaches its argument from the computation graph
• $x$: the context (the observed part of the input)
• $y$: the target (the part of the world the model must predict)
• $c$: any conditioning information, such as an action
• $f_\theta$: the context encoder, parameterized by $\theta$
• $\theta$: the parameters of the context encoder
• $f_{\bar\theta}$: the target encoder, parameterized by $\bar\theta$
• $\bar\theta$: the parameters of the target encoder (updated by EMA)
• $g_\phi$: the predictor, parameterized by $\phi$
• $\phi$: the parameters of the predictor
• $\mathcal{D}$: the data distribution
• $\mathcal{L}(\theta, \phi)$: the training loss for this chapter's squared-error instantiation
• $\mathbb{E}_{(x, y) \sim \mathcal{D}}$: the expectation over context-target pairs drawn from the data distribution
• $\|\cdot\|_2^2$: the squared Euclidean norm

The target parameters follow

$$\bar\theta \;\leftarrow\; \tau \bar\theta + (1 - \tau)\,\theta.$$

where:

• $\bar\theta$: the target encoder parameters
• $\theta$: the context encoder parameters
• $\tau$: the EMA decay rate in $[0,1]$, chosen close to one relative to the online-update timescale. Values below one update the target; $\tau=1$ freezes it. An official I-JEPA configuration, for example, schedules it from $0.996$ toward $1.0$; the toy example below uses a fixed value of $0.99$.

The following details play distinct roles, although related methods show that one mechanism can sometimes substitute for another.

There is no decoder. Nothing in $\mathcal{L}$ requires $z$ to be invertible or decodable. The embedding only has to be predictable and internally consistent. This removes the pixel-rendering burden and, at the same time, removes any external constraint that would keep the embedding from collapsing on its own.

The gradient updates are asymmetric. Gradients flow to $\theta$ and $\phi$ but not to $\bar\theta$. If you remove the stop-gradient and use the same encoder on both sides, both encoder evaluations can move through the same update toward a constant solution, as we will see in the seeded experiment. Detaching the target prevents that instantaneous gradient cooperation, but does not eliminate constant solutions mathematically.

The predictor is not a formality. In the simplified linear setting below, it can absorb part of the context-to-target regression that would otherwise constrain the encoder. This is an intuition under stated assumptions, not a general guarantee that predictors prevent shrinkage or collapse.

Consider the linear case to see the mechanics cleanly. Assume for this analysis that the context and target inputs have the same dimension and that $x$ and $y$ are centered, so that $W$ and $\bar W$ are directly comparable and the second moments below are covariance matrices. Suppose $f_\theta(x) = W^\top x$, $f_{\bar\theta}(y) = \bar W^\top y$, and the predictor is the identity. The population loss is

$$\mathcal{L}(W) = \mathbb{E}\big\| W^\top x - \bar W^\top y \big\|_2^2 .$$

where:

• $W$: the online encoder weight matrix
• $\bar W$: the target encoder weight matrix
• $x$: the context input
• $y$: the target input
• $\mathcal{L}(W)$: the population squared-error loss for the online encoder when the predictor is fixed to the identity
• $\mathbb{E}$: the expectation over context-target pairs from the data distribution
• $\|\cdot\|_2^2$: the squared Euclidean norm

Taking the gradient with respect to $W$ gives $2 \mathbb{E}[x (W^\top x - \bar W^\top y)^\top] = 0$. Setting this to zero and using the definitions of the covariance matrices (assuming $\Sigma_{xx}$ is invertible) gives

$$\begin{aligned} \Sigma_{xx} W &= \Sigma_{xy} \bar W && \text{(from } \mathbb{E}[xx^\top]\,W = \mathbb{E}[xy^\top]\,\bar W \text{)} \\ \Longrightarrow\quad W &= \Sigma_{xx}^{-1} \Sigma_{xy} \bar W . && \text{(solve for } W \text{)} \end{aligned}$$

where:

• $\Sigma_{xx} = \mathbb{E}[x x^\top]$: the covariance matrix of the context inputs
• $\Sigma_{yx} = \mathbb{E}[y x^\top]$: the cross-covariance matrix between the targets and the contexts
• $\Sigma_{xy} = \mathbb{E}[x y^\top]$: the cross-covariance matrix between the contexts and the targets
• $\Sigma_{xx}^{-1}$: the inverse of the context covariance matrix, which exists when $\Sigma_{xx}$ is invertible
• $W$: the online encoder weight matrix
• $\bar W$: the target encoder weight matrix

Write $A = \Sigma_{xx}^{-1}\Sigma_{xy}$. For the centered variables assumed here, $A$ maps each target-encoder direction to the least-squares coefficient vector for predicting that target coordinate from $x$. If the target encoder tracks the online encoder via an EMA and a stationary point also has $W = \bar W$, then the normal equation reduces to $W = A W$. Thus, the columns of a nonzero stationary $W$ must lie in the eigenspace of $A$ with eigenvalue one. When that eigenspace is absent, $W=0$ is the only stationary solution satisfying this simplified fixed-point condition. This calculation establishes the availability of a collapsed fixed point under the stated linear assumptions; it does not prove that every nonlinear predictor-free training run converges to collapse. The useful conclusion is that the predictor-free objective does not rule out the trivial stationary point.

Now add a predictor $g_\phi(z) = M^\top z$ with $M$ learnable. The loss becomes

$$\mathcal{L}(W, M) = \mathbb{E}\big\| M^\top W^\top x - \bar W^\top y \big\|_2^2 .$$

where:

• $M$: the predictor weight matrix
• $W$: the online encoder weight matrix
• $\bar W$: the target encoder weight matrix
• $x$: the context input
• $y$: the target input
• $\mathcal{L}(W, M)$: the population squared-error loss with a learnable linear predictor
• $\mathbb{E}$: the expectation over context-target pairs from the data distribution
• $\|\cdot\|_2^2$: the squared Euclidean norm

The composite map $M^\top W^\top$ can allocate part of the context-to-target regression to $M$, so the encoder is not required to absorb the entire regression map. In this simplified analysis, that extra degree of freedom permits nonzero stationary solutions that the identity-predictor fixed-point equation can exclude. A sufficiently capable predictor can represent systematic context-to-target mismatch here, so the encoder need not represent the whole regression map itself. This limited result helps explain the predictor in BYOL, SimSiam, and JEPA-style training; normalization, learning rates, regularization, target updates, and data geometry still affect the outcome. Representation-only evaluation may discard the predictor, whereas a latent world model needs a transition predictor to roll forward.

So far we have described JEPA abstractly with inputs $x$ and $y$. Spatial masking and temporal or action-conditioned prediction illustrate two different constructions. The choice determines which invariances the encoder can learn and which information it can discard.

Masked spatial prediction. In I-JEPA, the context encoder receives visible image patches. The target encoder processes the full unmasked image, after which embeddings at masked target locations are selected for prediction. The model therefore predicts contextualized hidden-region features from visible surroundings. This encourages representations that support masked-region feature prediction; whether they encode amodal object shape or help track occluded entities requires a separate evaluation.

Temporal and action-conditioned prediction. Original V-JEPA predicts masked tube-like regions whose targets span a clip's temporal dimension; it is not specifically a future-prediction model. In V-JEPA 2-AC, the predictor instead forecasts future frame embeddings from past states and actions. A deterministic action-conditioned model implements a map $\hat{z}_{t+1}=g_\phi(z_{\le t},a_{\le t})$ rather than a full conditional distribution $p(z_{t+1}\mid z_{\le t},a_{\le t})$. With sufficiently informative interaction data, conditioning on actions lets it model how transitions vary with the agent's actions; action conditioning alone does not guarantee causal attribution.

The action-conditioned temporal case is the one that connects directly to the rest of this book. The model now has the same interface as the RSSMs of Part V, Ch 1: a latent state, an action input, and a predicted next latent, but with a different objective and no reconstruction term. Planning then requires a cost defined over embeddings, for example the distance from a predicted embedding to an embedding of a goal observation. Sampling-based planners and model predictive control over such learned costs are the subject of Part VII: Planning and Agency. In V-JEPA 2-AC, the planner searches over action sequences while predicted embeddings serve as rollout state and define the objective. The embedding geometry matters directly even though arbitrary embedding vectors are not the optimization variables.

There is an obvious objection. Squared error in embedding space is a point prediction, and point predictions cannot express multimodality. If a pedestrian might go left or right, the optimal point prediction under squared error is the mean embedding, which may correspond to no plausible future at all. We discussed this failure mode in general terms under stochasticity and uncertainty in Part IV: Space, Physics, and Causality, and it is worth restating here because JEPA handles it in a way that is easy to miss.

The target encoder does not solve multimodality. Under squared error, the optimal deterministic prediction is the conditional mean of the target embedding. Learning the embedding can make nuisance variation less prominent, but it does not guarantee that this mean corresponds to any attainable future or remains useful for control. Representing distinct alternatives requires a predictor or objective that can produce or score multiple targets, for example through sets, multiple heads, stochastic outputs, energies, or probabilities. An EMA target encoder is one component of some anti-collapse recipes; it does not address ambiguity in a deterministic point predictor.