Inside a JEPA: Encoders, Predictors, and Collapse

Lesson 1 made the case for predicting representations. Now open the machine. A JEPA has three components, and their asymmetry is the whole design.

• The context encoder. Takes the visible part of the input (the unmasked patches) and produces embeddings. This is the network you actually care about, the one you keep and use for downstream tasks.
• The target encoder. Takes the full input and produces embeddings for the masked regions. These are the prediction targets. Crucially, it is not trained by gradient descent.
• The predictor. A smaller network that takes the context embedding plus a specification of where the target sits (positional information about the masked region) and outputs a guess at the target's embedding.

The loss is simple: the distance (typically L2) between the predictor's output and the target encoder's output, computed in embedding space, not pixel space. There is no decoder anywhere.

That positional conditioning matters. The predictor is not asked "what is somewhere in this image?" but "what is at this specific location?" Without it, the task would be ill-posed, and the model could not learn spatial structure. It is what turns the objective into a genuine prediction problem about a particular region rather than vague summarization.

The mechanism is clearer as pseudocode. One training step:

# x: a batch of images (or video clips)
context_patches, target_patches, target_pos = mask(x)

# 1. Encode the visible context (this network learns)
z_context = context_encoder(context_patches)

# 2. Encode the FULL input to get targets, no gradients
with torch.no_grad():
    z_target = target_encoder(x)[target_pos]
    z_target = layer_norm(z_target)   # stabilizes scale

# 3. Predict target embeddings from context + position
z_pred = predictor(z_context, target_pos)

# 4. Loss lives in embedding space; no pixels, no decoder
loss = mse(z_pred, z_target)
loss.backward()                        # updates context_encoder + predictor ONLY
opt.step()

# 5. Target encoder follows by EMA, it never gets a gradient
for p_t, p_c in zip(target_encoder.params(), context_encoder.params()):
    p_t.data = m * p_t.data + (1 - m) * p_c.data   # m ~ 0.996 -> 1.0

Every structural choice from the last step appears here. torch.no_grad() on the target branch is the stop-gradient. The EMA update in step 5 is why the target encoder has no optimizer. The loss compares vectors, never pixels.

After training you throw away the predictor and usually the target encoder. The context encoder is the product: the thing that turns an image into a useful representation. Everything else was scaffolding to train it.

Now the danger. Look again at the loss: minimize the distance between the predicted embedding and the target embedding, where both come from networks being trained. Ask the adversarial question: what is the easiest way to make that loss zero?

Not to understand images. The easiest solution is for the encoder to output the same constant vector for every input. Then the target is always, say, the zero vector, the predictor learns to always output the zero vector, and the loss is exactly zero, forever. Perfect score, and the representation carries no information whatsoever.

This is representation collapse, and it is the defining hazard of joint-embedding methods. It is not a rare pathology; it is the global optimum of the naive objective. A gradient descent process handed this objective with two trainable branches will happily find it, because it is the simplest possible way to satisfy the loss.

Notice why the other families in Lesson 1 do not face this. Generative methods cannot collapse: the pixel target is fixed by the data, so a constant prediction is heavily penalized, you cannot change the ground truth to make yourself right. Contrastive methods can collapse in principle, which is exactly why they use negatives: pushing different images apart makes a constant representation maximally bad, so the repulsion term explicitly forbids the cheat.

JEPA has neither defense. It has no fixed target and no negatives. So it needs another answer.

The answer, inherited from the BYOL and SimSiam line of work, is architectural asymmetry. The two branches are deliberately made different, and that difference alone is enough to prevent collapse without a single negative sample. Three ingredients combine.

• The stop-gradient. The target branch receives no gradient (the no_grad() above). This is essential: if gradients flowed to the target encoder, the optimizer could directly optimize the target to be easier to predict, which is precisely the road to a constant. Blocking that path means the model can only improve by making better predictions, never by making the answer easier. It breaks the symmetry that would otherwise let both sides collude on a trivial solution.
• The predictor. An extra network on the context side only, so the two branches are not interchangeable. SimSiam's ablations (Chen and He) showed both the predictor and the stop-gradient are necessary, remove either and the representation collapses.
• The EMA target encoder. The target's weights are a slow exponential moving average of the context encoder's. This makes the target a lagged, smoothed version of the online network: a moving but stable goalpost. Because it changes slowly, the online network is always chasing a target it cannot instantly corrupt.

BYOL and SimSiam are nearly identical except for the momentum encoder, which shows the mechanism is not one trick but a family of ways to break symmetry.

It is worth being honest: exactly why stop-gradient plus a predictor prevents collapse is subtle, and was not obvious even to the researchers who found it empirically. Several complementary explanations exist.

The cleanest intuition is about the direction of optimization. With a stop-gradient, the system is no longer doing plain gradient descent on a single energy that both branches share. Only one branch moves in response to the loss; the other is dragged along slowly by the EMA. The dynamics become non-conservative, they do not simply descend the naive objective, so the collapsed solution, though it is a minimum of that objective, is not what the actual training dynamics converge to. The optimizer is following a modified vector field in which the trivial solution is not the attractor.

A second framing notes that the predictor plus stop-gradient produces an implicit repulsion similar to what contrastive negatives supply explicitly. The asymmetric structure induces contrastive-like effects without a contrastive loss, so the model gets the benefit of "spread the representations out" without paying the price of choosing negatives.

The practical upshot for a practitioner is concrete: collapse is an engineering reality to monitor, not a solved theorem. Teams track the variance or rank of the embeddings during training; variance heading toward zero, or the representation's effective rank falling, is the diagnostic that collapse has begun. Momentum schedules and layer normalization on targets are stabilizers tuned for exactly this reason.

One piece of LeCun's framing completes the architecture. In his 2022 position paper, and developed formally in the latent-variable energy-based model lecture notes by Anna Dawid and Yann LeCun (2023), a JEPA is an energy-based model: it does not output a probability distribution, it assigns a scalar energy, low for compatible context-target pairs, high for incompatible ones. Training shapes that energy landscape, and collapse is precisely the degenerate landscape where energy is low everywhere.

This framing motivates the last component: a latent variable fed to the predictor. The world is genuinely ambiguous. Given the left half of an image, or a video of a hand approaching a cup, several outcomes are compatible with the context, and no deterministic predictor can output all of them. Forced to produce one answer, it outputs the average, which is often a meaningless blur in representation space.

A latent variable gives the predictor a place to put that irreducible uncertainty: as the latent varies, the prediction sweeps the set of plausible outcomes. The predictor is then no longer asked "what is the answer?" but "what answers are compatible?", which is the honest question.

LeCun extends this to H-JEPA, a hierarchical stack where higher levels predict over longer horizons at greater abstraction, because coarse predictions stay valid further into the future. That hierarchy is the proposed route from a representation learner to a world model that can plan, which is where Lesson 3 goes.

Assemble the mechanism into the choices an implementer actually makes.

The unifying reading: JEPA buys the freedom to abstract by taking on the burden of collapse, then pays that burden with asymmetry rather than with negatives or a fixed pixel target. Every row is either exercising that freedom or defending it.

That is the trade in full. Lesson 1 showed why predicting representations is the right target; this lesson showed the machinery that makes it trainable, and the failure mode that machinery exists to prevent. What remains is the payoff: what happens when this architecture is actually built and scaled, on images, then on a million hours of video, and finally used to plan a robot's actions. That is Lesson 3.

The context encoder embeds visible patches and the predictor guesses the masked region's embedding at a given position, while the target encoder embeds the full input under a stop-gradient; the loss compares embeddings, gradients update only the context side, and the target encoder follows by EMA.