Predictive Embeddings, JEPA, and Energy-Based Models

Imagine a world model for a delivery robot navigating a sidewalk. To decide whether to step left or right, the robot does not need to know the exact pixel values of an oak leaf trembling in the wind, the precise visual texture of wet concrete, or the sub-millimeter position of every pebble. It needs to know that the pedestrian ahead is drifting right at walking speed, that the gap between the parked car and the curb is narrowing, and that its own left wheel is about to clip the curb. Much of the raw image can be irrelevant to this decision, and some of it can be unpredictable from the robot's context. A model that spends finite capacity reproducing the leaf and the concrete may be learning detail its controller will never use.

Every observation-generating architecture we have met so far in Part V: World-Model Architectures pays some price for that irrelevance. Reconstruction-based recurrent state-space models optimize an observation-reconstruction term, so they can allocate parameters and training signal to detail the agent will never use; how much depends on the architecture, data, and objective weights. Autoregressive token models predict the next discrete symbol, which means prediction errors in task-irrelevant regions still contribute to the training signal. Diffusion models typically generate by iteratively reversing a noising process, while flow models integrate a learned vector field from a base distribution to the data distribution. Both observation-generating approaches spend computation on relevant and irrelevant visual detail.

This chapter takes a different stance. Instead of asking the model to reproduce the world, we ask it to predict a representation of the world. The model never renders an image. It maps observations into a learned embedding, which may be made compact by architectural or capacity constraints; in temporal or action-conditioned versions, a predictor learns dynamics in that embedding space. The prediction target is a vector in a learned latent space. No decoder, no pixels, no per-symbol likelihood. The model's job is to answer a narrower question: given its context, what target representation is consistent with it?

The idea immediately runs into a serious problem. If the model chooses both the representation and the prediction target, it can cheat. Map every observation to the constant zero vector and the prediction is perfect. The loss goes to zero, the model learns nothing, and the agent is blind. An unprotected joint-embedding matching objective admits this constant solution, and standard symmetric implementations can collapse rapidly after training begins. The main anti-collapse responses in this chapter come in three forms:

• Explicit-negative contrastive methods, which train compatible pairs against sampled decoys.
• Asymmetric methods, which change the two branches' gradient updates and can discourage collapse in specific training recipes without making the constant solution impossible.
• Regularized methods, which constrain embedding statistics; VICReg, for example, directly penalizes low per-coordinate standard deviation and off-diagonal covariance.

The Joint-Embedding Predictive Architecture, or JEPA, is a recent predictive-embedding framework introduced in a 2022 position paper. Practical variants combine predictors and asymmetric target handling; explicit variance or covariance regularization is a separate, optional mechanism. This chapter builds up the EMA-target structure found in I-JEPA and V-JEPA, studies a squared-error instantiation for its mathematical analysis and toy system, connects that illustrative discrepancy to the broader energy-based view, and compares collapsed and non-collapsed representations after training.

By the end you should be able to look at a "predict in latent space" objective and ask three questions: what discourages collapse, what kind of uncertainty the prediction target can express, and what evidence shows that the representation retains information an agent needs to act. Those questions expose most of the important design tradeoffs.

Start with the most basic design question. A world model is a function that takes what the agent has seen and done and produces something that stands in for what comes next. The question is what that "something" should be. In the observation-prediction view, that "something" is an observation, predicted by a separate observation-space model $h_\psi$: $\hat{o}_{t+1} = h_\psi(o_{\le t}, a_{\le t})$, where $\psi$ are the parameters of the observation-space model and $\hat{o}_{t+1}$ is the predicted next observation. This is the classical prediction problem: you predict the quantity you measured, and you are graded on how closely you reproduced it. In the representation-prediction view, we interpose an encoder $f_\theta : \mathcal{O} \to \mathcal{Z}$ and predict in $\mathcal{Z}$ instead:

where:

• $\hat{z}_{t+1}$: the predicted latent at time $t+1$
• $g_\phi$: the dynamics model (or predictor) that maps the context embedding and actions to a predicted latent
• $\phi$: the parameters of the dynamics model (or predictor)
• $f_\theta$: the online encoder that maps an observation sequence (or context) to a latent, parameterized by $\theta$
• $\theta$: the parameters of the online encoder
• $\mathcal{O}$: the observation space
• $\mathcal{Z}$: the latent embedding space
• $o_{\le t}$: the sequence of observations up to time $t$
• $a_{\le t}$: the sequence of actions up to time $t$
• $\mathcal{L}$: the loss function
• $f_{\bar\theta}$: the target encoder, a separate copy of the encoder with parameters $\bar\theta$
• $\bar\theta$: the parameters of the target encoder
• $o_{t+1}$: the observation at the next time step
• $\|\cdot\|_2^2$: squared Euclidean norm
• $t$: the current time step

The notation already contains the seeds of the whole chapter. There are two encoders, $f_\theta$ and $f_{\bar\theta}$, not one. The second one, the target encoder, is what makes the objective meaningful, and we will spend a full section on it. For now, notice what has changed structurally: the prediction target is no longer the measured observation itself. It is a learned embedding vector expressed in the encoder's coordinate system, and the meaning of "predicting well" depends entirely on what that coordinate system chose to preserve. Two different encoders trained on the same data can yield the same numerical loss while preserving completely different information, because the loss only measures agreement in whatever space the encoders happen to occupy.

The cleanest way to see why this matters is with an information budget. Let $S_t$ be a state sufficient to determine the conditional distribution of future observations and rewards given future actions, as discussed in Part II: Inference, Decisions, and Control, and let $O_t$ be the observation. Observations can contain large amounts of nuisance or redundant information even when the task-relevant state is compact. For a predictive world model, the representation should retain information needed to predict task-relevant future distributions under candidate actions. Modeling information with no downstream use can consume finite capacity and optimization effort without improving the decision objective.

An observation-reconstruction objective pressures $z_t = f_\theta(o_{\le t})$ to retain information about $O_t$. Perfect reconstruction by a deterministic decoder would imply that the observation can be recovered from the latent; finite-capacity, lossy training provides no such guarantee. In either case, minimizing observation error can spend capacity on pixels that do not help a decision. There is no general guarantee that a reconstruction-optimal representation is aligned with a decision-optimal one. The reconstruction objective does not know which pixels matter; it treats a trembling leaf and a braking car with the same weight, because both contribute to the pixel error in the same way.

A latent-prediction objective with a target encoder instead pushes toward a different quantity. If $z_t$ and $z_{t+1}$ are produced by encoders that must be mutually predictive, the encoder has an incentive to discard whatever is unpredictable between consecutive observations, because carrying unpredictable information only inflates the loss. If a pixel's next flicker value is independent of the entire available context, that value is unpredictable from preceding frames. Same-observation reconstruction can preserve and reproduce the realized flicker; future prediction may drop it. In this sense the prediction objective can perform automatic feature selection driven by temporal structure rather than explicit supervision.

This is not a free lunch. Discarding what does not help at the trained timescale can remove variables whose decision-relevant effects or dependencies appear only farther into the future. A pedestrian's crossing intention may become clear only after several observations; the consequences of wheel slip may also emerge over a longer rollout. Prediction horizon and masking directly shape which content the objective rewards. The target-encoder timescale is instead an optimization choice controlled by the EMA hyperparameter; it can influence the learned representation through training dynamics without directly selecting semantic content.

It helps to lay out the design space before diving into any one method. Latent-space prediction methods differ in what pressures keep the representation informative. Without one of these anchors or another effective constraint, a constant representation remains an available optimum, and training can collapse to it. The table separates explicit-negative contrast from the broader energy-based framework and separates architectural asymmetry from statistical regularization.

Reconstruction trains a decoder to recover observation content to the fidelity imposed by its loss; only perfect reconstruction with a deterministic decoder implies exact recoverability. This tethers the latent to the observation but can reintroduce pressure and cost from predicting nuisance detail. Explicit-negative contrastive objectives avoid the decoder by requiring compatible pairs to score above sampled decoys; their behavior depends on the loss and negative-sampling scheme. Asymmetric methods use combinations such as a predictor plus stop-gradient, often with an EMA target. Statistical methods instead penalize low variance or redundant coordinates. These mechanisms can discourage collapse, but none is an architecture-independent guarantee. Energy-based modeling is the broader scalar-compatibility framework: explicit-negative contrast is one way to train an energy, not the definition of the framework.

Handing the model freedom over the target representation has three concrete costs, and it is worth naming them now so the rest of the chapter reads as a response to them rather than a catalogue. These costs are not incidental; they are the reason that reconstruction and contrastive methods remain in active use despite the practical advantages of the asymmetric family.

The first cost is loss of an observation likelihood. If there is no normalized observation model, you cannot compute $p(o_{t+1} \mid o_{\le t})$. You can compute distances or energies between embeddings and train classifiers or planners on top of them, but those scores do not say how probable a particular observation was. This matters when normalized, calibrated probabilities or likelihood-based anomaly scores are required. Distance- and energy-based anomaly detection remains possible, but it needs separate validation or calibration.

The second cost is opacity of the prediction target. When a diffusion model predicts a frame, you can look at the frame. When a JEPA predicts an embedding, you interpret it through probes, nearest neighbors, or downstream task performance. Some semantic failures are therefore less immediately visible, and a non-degenerate embedding can still omit a variable the policy needs. The prediction loss does not know which variables the policy values unless the task construction makes them relevant.

The third cost is the trivial solution, which is the subject of the collapse section. A constant decoder cannot match varied observations, and a constant representation gives every contrastive candidate the same score. For an InfoNCE loss with $N$ candidates, equal scores give chance probability $1/N$ and loss $\log N$; that is not optimal whenever the model can score the positive above sampled negatives. Positive-pair matching can collapse because a constant output trivially matches a constant target, so its anti-collapse mechanisms must be evaluated explicitly.

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.

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

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

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

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

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.

Energy functions supply a scalar language for compatibility and multimodal prediction. They cover normalized probabilistic models, explicit-negative contrastive objectives, and unnormalized scores, but writing a stop-gradient discrepancy as an energy does not by itself explain the optimization dynamics that prevent collapse. The distinctions between these uses matter more here than a single historical lineage.

An energy-based model assigns a scalar energy $E_\psi(x, y)$ to a pair. When the energy is measurable and the partition function satisfies $0 < Z_\psi(x) < \infty$, it defines a density over continuous $y$ relative to a chosen base measure $\mu$ through the Boltzmann, or Gibbs, form:

where:

• $p_\psi(y \mid x)$: the conditional density over targets $y$ given context $x$
• $x$: the context input
• $y$: the target variable
• $\psi$: the parameters of the energy function
• $E_\psi(x, y)$: the scalar energy assigned to the pair $(x, y)$, parameterized by $\psi$
• $Z_\psi(x)$: the partition function, which normalizes the distribution over $y$
• $y'$: a dummy variable of integration over possible targets under base measure $\mu$

For a discrete target space, the integral in $Z_\psi(x)$ is replaced by a sum. Adding an $x$-dependent constant to every energy leaves the normalized conditional distribution unchanged. Only energy differences matter, so the energy zero point, or additive offset, is not identifiable; multiplying all energies would generally change the distribution's effective temperature. A low energy is not by itself a probability. In continuous spaces, probability mass depends on the relative depths and integrated shapes or volumes of the wells, not depth alone.

The quantity $Z_\psi(x)$ is the partition function. Computing or differentiating it is a central practical difficulty in likelihood-based energy learning. Training by maximum likelihood gives

where:

• $\nabla_\psi \log p_\psi(y \mid x)$: the gradient, with respect to $\psi$, of the model log-density of the observed target
• $x$: the context input
• $y$: the observed target
• $\psi$: the parameters of the energy function
• $\nabla_\psi E_\psi(x, y)$: the gradient of the energy of the observed pair
• $\mathbb{E}_{y' \sim p_\psi(\cdot \mid x)}[\nabla_\psi E_\psi(x, y')]$: the expected energy gradient under the model's current distribution over $y'$
• $y'$: a sample from the model distribution $p_\psi(\cdot \mid x)$

The first term is the positive phase: it pushes the energy of observed pairs down. The second term is the negative phase: it pushes up the energy of pairs the model currently considers likely. Without the negative phase, nothing in this maximum-likelihood objective opposes lowering observed energies through a broad or nearly uniform downward shift, so the model need not learn meaningful relative normalization. Approximating the model expectation is the main computational problem for likelihood-trained EBMs. Other energy objectives replace maximum likelihood rather than approximate this particular term.

Maximum-likelihood sampling. Run Langevin dynamics or a Markov chain targeting $p_\psi$ to draw $y'$ for the negative phase. Under suitable convergence conditions this is exact in the limit, but high-dimensional mixing can be slow or unstable. Denoising diffusion uses a different construction in which iterative sampling is part of the model; for a normalized Gibbs density the score identity is $\nabla_y \log p(y) = -\nabla_y E(y)$.

Contrastive estimation. Train a sampled classification objective over one positive and a finite candidate set of negatives drawn from a chosen marginal, minibatch, noise distribution, or proposal. The candidates are generally resampled during training; they are not equilibrium samples from the current EBM. A canonical form is InfoNCE:

where:

• $x$: the context input
• $y^+$: the true successor (positive example)
• $y_j^-$: the $j$-th negative example drawn from a proposal distribution
• $K$: the number of negatives
• $\psi$: the parameters of the energy function
• $E_\psi(x, y)$: the energy assigned to the pair $(x, y)$
• the expectation is over the data distribution and the negative sampling distribution
• $\mathcal{L}_{\mathrm{NCE}}$: the InfoNCE contrastive loss

Under the standard sampling construction, with one positive from the joint distribution and negatives drawn independently from the target marginal, InfoNCE yields the lower bound $I(X;Y) \ge \log(K+1)-\mathcal{L}_{\mathrm{NCE}}$, where $X$ and $Y$ are the context and target random variables. Increasing the candidate count can raise the ceiling of this particular bound, but a tighter numerical bound is not guaranteed for every finite model or negative sampler. Computing more candidate scores also raises training cost.

Asymmetry. The third option, used by several joint-embedding methods, is to avoid explicit negatives and use architectural asymmetry together with other anti-collapse mechanisms. The illustrative squared prediction discrepancy used here can be read as a compatibility score

where:

• $x$: the context input
• $y$: the target input
• $c$: 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
• $g_\phi$: the predictor, parameterized by $\phi$
• $\phi$: the parameters of the predictor
• $g_\phi(f_\theta(x), c)$: the predictor applied to the context embedding and conditioning input
• $\mathrm{sg}(f_{\bar\theta}(y))$: the stop-gradient target embedding
• $E_{\theta,\phi,\bar\theta}(x, y; c)$: the compatibility energy assigned to the pair $(x, y)$ under conditioning $c$, parameterized by $\theta$, $\phi$, and $\bar\theta$
• $\mathrm{sg}(\cdot)$: the stop-gradient operator, which detaches its argument from the computation graph
• $\|\cdot\|_2^2$: the squared Euclidean norm

but stop-gradient and EMA do not estimate the maximum-likelihood negative-phase expectation above. They change the optimization dynamics of the positive-pair regression problem. Because the target branch is detached and, with an EMA decay close to one, updated on a slower timescale, the two branches cannot co-adapt through the same instantaneous gradient. Collapse can still be a mathematical solution, so practical methods rely on combinations of architecture, normalization, predictor dynamics, data augmentation or masking, and, in some variants, explicit variance regularization. Calling the method self-distillation describes the teacher-student relation; it should not be mistaken for likelihood training of a normalized energy model.

Three practical benefits come from thinking in energies rather than in losses.

First, you can compose energies. If $E_1$ scores progress toward a goal and $E_2$ scores obstacle avoidance, then, for $\lambda \geq 0$, $E_1 + \lambda E_2$ expresses a weighted soft tradeoff. For nonnegative, comparably scaled penalties with aligned zero points, the sublevel set $\{E_1 + E_2 \leq \epsilon\}$ lies inside $\{E_1 \leq \epsilon\} \cap \{E_2 \leq \epsilon\}$, which supports an approximate joint-constraint interpretation. Arbitrary energies do not implement an exact intersection, and new terms must be calibrated against the existing objective.

Second, an energy can represent multiple separated low-energy wells, although sampling among them may still be difficult. A squared-error point predictor instead returns the conditional mean. Depending on the mode locations and probabilities, that mean can lie between the wells, including in an implausible high-energy region.

Third, the energy lens covers several objectives without making them equivalent. Contrastive predictive coding and InfoNCE use sampled classification; triplet and margin losses use discriminative ranking; score matching uses a derivative-based criterion that avoids the partition function. For likelihood-trained EBMs, ask how the negative phase is estimated. For alternative objectives, ask what replaces likelihood training and what prevents a flat, uninformative score.

The cost is that sampling or minimizing a general high-dimensional energy is often iterative, and evaluating $Z_\psi(x)$ is intractable in general. Tractable structure or an amortized proposal can reduce that cost. When iterative inference requires many model evaluations, its latency can matter in planners with tight per-query budgets.

Everything converges on one question. What stops the encoder from outputting a constant?

Complete collapse is the extreme case: $f_\theta(o) = \mathbf{v}$ for all $o$, for a fixed vector $\mathbf{v}$. The prediction loss is exactly zero. The representation carries no information at all, and any downstream probe will fail. Complete collapse is easy to detect because the per-dimension variance is zero for every dimension.

Dimensional collapse is subtler. The embedding retains variation but occupies an effectively lower-dimensional subspace; strong coordinate correlations can be one symptom, but correlation alone is not the definition. If a nominally 64-dimensional latent has entropy effective rank 3, it uses roughly three dimensions' worth of spectral capacity, a deficit of 61 relative to uniform use of all 64. That does not mean exactly 61 singular directions are zero. The loss can still fall smoothly while effective dimensionality disappears.

Diagnostics matter here. Cheap moments can run every few epochs on a fixed validation subset; approximate effective-rank checks and trained probes can run at a cadence their cost permits. Useful signals include:

• Per-dimension standard deviation. After fixing the representation scale, a coordinate whose standard deviation is near zero over representative data is nearly constant. Plot every coordinate because a mean can hide coordinate-wise failure.
• Effective rank. When the singular values are not all zero, center the embedding matrix, normalize its singular values as $\tilde\sigma_i = \sigma_i / \sum_j \sigma_j$, and compute $\exp\big(-\sum_i \tilde\sigma_i \log \tilde\sigma_i\big)$, using the continuous-extension convention $0\log 0=0$ for zero components. The normalized spectrum is undefined when every singular value is zero; handle that case separately as complete collapse. Otherwise, effective rank is a smooth measure of spectral dimensionality, not a count of exactly nonzero directions.
• Linear probe performance on a held-out variable the representation should contain, such as the true state or a task label. It tests linear decodability under the chosen split and probe. Poor performance does not prove that the information is absent, and high performance does not establish downstream usefulness or causal relevance.
• Uniformity and alignment, if you are working in the contrastive tradition: how evenly the embeddings spread on the unit sphere, and how close positive pairs are.

The distinction between "non-collapsed" and "useful" is the one that matters. A representation can have full effective rank and still fail at control. Effective rank and variance diagnose degeneracy; a held-out linear probe tests access to one chosen variable. Downstream usefulness ultimately requires task-level evaluation.

If the loss contains terms that explicitly push decoys away, the constant solution is not optimal whenever the model can score positives above sampled negatives. A constant embedding makes all $K+1$ candidate scores equal, so the per-example InfoNCE loss is $\log(K+1)$, the chance-level value. It is not the largest possible loss, because assigning a positive pair a much worse score than its negatives can make the loss arbitrarily larger. The price is the cost of negatives and sensitivity to the sampling distribution. In practice this can motivate large batches, memory banks, or hard-negative mining.

One widely used asymmetric recipe combines a predictor, stop-gradient, and an EMA target. The components interact, but no single three-part recipe is universally necessary: SimSiam, for example, uses a predictor and stop-gradient without an EMA target.

In the simplified centered linear regression above, a sufficiently capable predictor can absorb part of the context-to-target map. Direct unnormalized matching can favor shrinkage in directions where that map is contractive, but omitting a predictor does not universally imply shrinkage.

The stop-gradient breaks the symmetry of the gradient updates while leaving the forward scalar loss unchanged. Without it, the same parameters can receive updates through both encoding branches and move cooperatively toward a constant solution. With it, the online encoder cannot directly update the target representation through that loss.

The EMA target encoder controls the target timescale. A permanently frozen target would reduce training to regression toward fixed, randomly initialized target features, which need not encode useful invariances. Copying online weights into the target after every update would preserve stop-gradient within each step yet remove EMA smoothing. A larger $\tau$ lengthens the nominal smoothing timescale and gives less weight to the current online parameters; a smaller value shortens that timescale and gives them more weight. Individual target-step magnitudes and the pointwise online-target gap also depend on the evolving online trajectory, so neither comparison guarantees monotone realized changes or pointwise closeness. A smaller value may reduce stability or representation quality in recipes that rely on a slow target. Whether this appears as collapse depends on the other mechanisms.

The third family adds explicit statistical pressure. VICReg penalizes per-dimension standard deviations below a threshold and off-diagonal covariance entries. Barlow Twins pushes the cross-correlation matrix between two augmented-view embeddings toward the identity. Whitening methods such as W-MSE normalize embedding covariance using a transformation such as Cholesky, eigendecomposition, or an iterative inverse square root. Sinkhorn normalization instead balances assignment matrices in methods such as SwAV; it is not covariance whitening. These losses are broadly composable, but their effectiveness still depends on batch statistics, normalization, dimension, loss weights, and the surrounding objective.

The hinge form of the variance term is important. Penalizing deviation from a target standard deviation would be symmetric and would fight the encoder's natural scale, whereas the hinge $\max(0, \gamma - \mathrm{std}(z_j))$, where $\gamma>0$ is the target standard-deviation threshold and $z_j$ is the $j$-th embedding coordinate, only fires below that threshold and is silent otherwise. This is what makes the regularizer compatible with a wide range of architectures. A dimension at or above the threshold contributes nothing to the penalty, so the regularizer does not impose an upper target scale on the representation; it only refuses to tolerate dimensions below the chosen minimum.

There is no single best answer, but the tradeoffs are clear enough to guide a first choice.

• If you have a natural contrastive structure, an explicit-negative objective makes its anti-collapse pressure easy to inspect, but requires a sampling scheme and candidate scores.
• If you want to avoid explicit negatives and memory banks, an asymmetric recipe is an option. Its target encoder and predictor still add computation, and successful recipes depend on their surrounding architecture and optimization.
• If you want explicit anti-collapse pressure that composes with many architectures, add variance and covariance regularization when batch statistics and loss weights are reliable.
• Combining mechanisms is common. The full toy variant below uses an EMA target and a variance hinge; that combination is stable in this seeded setup, not a general optimum for small datasets.

Now let's build one and watch the failure mode happen.

We will use a small system where the true state is recorded exactly. The environment is a two-dimensional linear dynamical system with a rotating, contracting transition matrix. The agent observes a 32-dimensional nonlinear function of the state plus Gaussian noise. The observation is much bigger than the state, and the encoder must decide what to keep. Because the probe targets are ground-truth simulated states, held-out $R^2$ measures their linear decodability under a specified ridge-regression protocol. It does not measure all retained information or control usefulness. The experiment and evaluation cells compute the arrays used by the later presentation-only figure cells.

Two patterns stand out. First, the raw-observation reference has the strongest linear probe: the state is substantially linearly decodable from informative directions in the noisy observation, although the noise prevents $R^2=1$. Second, the pixel predictor is competitive on this toy problem. The setup lacks the structured, high-dimensional nuisance content of real images even though all 32 observation coordinates contain independent noise. Latent prediction shows no probe advantage here, and the experiment does not measure compute efficiency or scaling. Predictive objectives can become more useful as unpredictable nuisance variation grows while task-relevant structure remains predictable, but that is a conditional design argument rather than a monotonic guarantee.

The worked example exposes several parameters that control the predictive embedding:

• D_LATENT: Dimension of the learned embedding. It limits coordinate capacity and should be selected with task-level evaluation. A smaller latent can strengthen a capacity bottleneck, but it does not by itself determine whether nuisance detail or decision-relevant state is discarded.
• ema: Exponential moving average decay for the target encoder. Smaller values track faster and may reduce stability in recipes that rely on a slow target. The appropriate schedule depends on the update timescale; this toy uses 0.99.
• var_weight: Weight on the VICReg-style variance hinge. It sets the strength of the direct anti-collapse pressure.
• epochs: Number of optimization steps. Too few leaves the predictor undertrained; too many can expose slow collapse.
• lr: Adam learning rate for the online encoder and predictor.
• Predictor hidden size: Capacity of the predictor MLP. In the simplified analysis, a sufficiently capable predictor can represent part of the context-to-target regression; practical capacity must be validated with the rest of the recipe.

Reconstruction and explicit-negative contrastive learning are two prominent self-supervised approaches, alongside clustering, masked prediction, and other pretext tasks. Reconstruction provides an inspectable observation-space target and resists constant collapse through its decoder, but it can spend training signal on irrelevant detail. Explicit-negative objectives provide direct discrimination and mutual-information analyses, but depend on candidate construction and can motivate large batches or memory banks. Negative-free asymmetric methods avoid those candidate stores and observation decoders, while target encoders and predictors add their own computation. I-JEPA and V-JEPA demonstrate EMA-target predictive-embedding systems at image and video scale, with the nonsquared loss choices noted earlier. V-JEPA 2 further demonstrates action-conditioned latent post-training and robot planning after large-scale action-free video pretraining. Those results establish concrete uses; they do not isolate the objective from architecture, data, scale, masking, loss choice, or planning design as the cause.

The limitations cluster around likelihood, sufficiency, and multimodality. A basic predictive embedding can rank candidate futures by energy but cannot report a normalized probability for a specific observation. This removes direct likelihood-based anomaly scoring and probability calibration; energy or latent-error anomaly scores remain possible but require separate validation or calibration. Evaluation therefore uses explicit protocols such as held-out probes, task performance, retrieval, or rollouts, and conclusions can vary with the chosen protocol.

The basic JEPA objective also offers no task-agnostic sufficiency guarantee. An encoder can discard a variable that is hard to predict at the trained horizon even when its later effect matters for braking. Careful horizons, masks, conditioning, or task-aware constraints can reduce this risk, but guarantees require additional assumptions. Finally, squared error is a point objective: its conditional-mean prediction can be misleading for ambiguous futures. General high-dimensional energies can represent multiple modes, but sampling or minimization is often iterative unless the model has tractable structure or an amortized proposal. Some systems, including V-JEPA 2-AC, keep latent prediction deterministic and use it inside an action-sequence planner; that ranks actions but does not itself represent exogenous multimodal uncertainty.

The cost balance depends on architecture and execution phase. In common recipes, anti-collapse mechanisms can cost less than a large observation decoder or iterative sampler, but a target network still adds computation and state, and statistical losses add tunable weights. EMA timescale, variance regularization, predictor capacity, and prediction horizon can each affect training and should be validated jointly for the intended setup. An EMA decay that is too small for one recipe may reduce stability; an overly strong variance penalty can inflate coordinate variation without making it useful; a short horizon can miss variables whose decision-relevant effects appear later. Cheap per-coordinate statistics can run often, while effective-rank checks and trained probes can run at a cadence their cost permits. The worked example reports all three, but task-level evaluation remains necessary.

This family foregrounded a choice earlier world models also made: what the model predicts has consequences for the representation it learns. Reconstruction, explicit-negative discrimination, and asymmetric prediction hold a model accountable for different targets. The next chapter examines hierarchical, hybrid, omnimodal, and modular architectures, including multimodal systems that combine several training objectives. The same target-choice question carries into Part VI: Learning World Models at Scale.

• Predicting representations instead of pixels shifts the model's accountability from rendering the world to capturing what the target construction rewards. It can discard unrewarded nuisance detail and information that is unpredictable at the trained timescale, but neither outcome is guaranteed; auxiliary objectives can reward retaining either.
• The EMA-target JEPA recipe studied here combines a context encoder, an EMA target encoder, a predictor, and detached targets. Generic JEPA does not mandate this exact recipe, and there is no observation decoder.
• In a simplified linear fixed-point analysis, a learnable predictor can absorb part of the context-to-target regression and permit nonzero stationary solutions. It is an important optimization component, not a general mathematical guarantee against collapse.
• Energy functions provide a common compatibility language. Maximum-likelihood EBM training has positive and negative phases; practical methods either approximate the model expectation or replace maximum likelihood with objectives such as sampled classification, margin ranking, or score matching.
• Contrastive estimation uses explicit decoys. Asymmetric methods use stop-gradient and, in some recipes, a slowly updated target to alter positive-pair optimization dynamics; they do not estimate the maximum-likelihood negative phase.
• Complete collapse maps every input to a constant, while dimensional collapse concentrates variation in a low-effective-rank subspace. Effective rank and per-dimension variance diagnose degeneracy. A held-out linear probe tests whether one chosen variable is linearly recoverable; task usefulness requires downstream evaluation.
• Anti-collapse mechanisms fall into three families: contrastive negatives, architectural asymmetry (predictor, stop-gradient, EMA target), and variance or covariance regularization. They compose, and the worked example uses two at once.
• The costs include no direct observation likelihood, no task-agnostic sufficiency guarantee, and the weakness of point predictions for multimodal futures. The next chapter examines hierarchical, hybrid, omnimodal, and modular architectures, including multimodal systems that combine reconstruction and predictive objectives.

Ready to test your understanding of predictive embeddings, JEPA, and energy-based models? This quiz checks whether you can reason about the design choices that keep latent prediction from collapsing and about what energy functions do and do not buy you.