Energy-Based Models and Prediction in Representation Space

An energy-based model does not output a probability or a sample. It outputs one number saying how compatible an input and a candidate answer are — low for good pairs, high for bad ones. Everything hard about the family comes from converting that number into a probability, and everything interesting about the recent work comes from refusing to.

An energy function, and the constant that ruins it

E(x, y) -> a scalar.   Low = compatible, high = not.

To turn it into a probability:

              exp(-E(x, y))
  p(y|x) = ---------------------
                  Z(x)

  Z(x) = sum over EVERY possible y of exp(-E(x, y))

Z, the partition function, is a sum over the entire output space. For a next token it is a sum over the vocabulary — fifty to two hundred thousand terms, which is why a softmax over logits is tractable and why a language model is, technically, an energy-based model with a cheap Z. For an image, a video frame or a continuous action it is a sum over an uncountable space, and it cannot be computed at all.

The framing is worth having because it unifies things that look unrelated. A classifier, a VAE, a GAN discriminator and a reranker are all scoring compatibility; they differ in how they handle the normalisation. Once you see Z as the obstacle, the methods sort themselves into ways of getting around it.

Three ways to avoid the partition function

The second row is the most useful thing on this page for placing the modern landscape. Diffusion did not replace energy-based modelling; it is an energy-based method that found a way to train without ever touching the constant.

Where energy models are already in your stack

The framing is not exotic. Several components of an ordinary language model pipeline are energy models that nobody calls energy models, and recognising them is the fastest way to make the abstraction useful.

• Cross-encoder rerankers. A reranker takes a query and a document together and emits one score. That is E(x, y) exactly. It is trained contrastively on relevant and irrelevant pairs, and it is used by minimising over a small candidate set produced by a cheaper retriever — which is precisely the “search over candidates” inference described below, restricted to a shortlist so that it terminates.
• Reward models. A reward model is a scalar head on a transformer scoring a prompt-completion pair. It is trained on preference comparisons with a loss that pushes the preferred completion’s score up and the rejected one’s down — a contrastive energy objective, with the normalisation avoided in exactly the way described above.
• Best-of-n sampling. Generate several candidates, score them all, keep the best. That is minimising an energy over a finite sample of the output space, and it is why scoring candidates at inference improves results at the cost of compute.
• Constrained decoding. Adding a penalty for outputs that violate a grammar is adding a term to an energy and minimising the sum, which is the composability property the framework promises, in a form people already use.

The pattern in all four is the same: a scoring function is easy to train and hard to sample from, so the sampling is delegated to a cheaper generator and the score is used to choose. That division of labour is the practical shape energy-based modelling actually takes.

Inference is a search, not a forward pass

The other structural difference. A feed-forward model computes y = f(x). An energy model asks you to find the y that minimises E(x, y), typically by gradient descent in the output space.

y = initial guess
repeat n times:
    y = y - lr * dE(x, y)/dy
return y

n is chosen at inference. More steps = more compute = potentially
better answers, from the same weights.

Two things follow. Compute at inference becomes a dial rather than a fixed cost, which is the same idea reasoning models reach by a different route. And constraints compose by addition: to require an output to satisfy something, add a term to the energy and minimise the sum. There is no equivalent trick for a feed-forward model, where a constraint has to be trained in or enforced afterwards.

The cost is that the search is a search. It can be slow, it can find a local minimum, and the answer depends on where it started — which is exactly the property that makes it unattractive for a production system with a latency budget.

JEPA: predict the representation

The Joint Embedding Predictive Architecture is the concrete proposal this idea is currently attached to, set out in Yann LeCun’s 2022 position paper and realised in I-JEPA for images (2023) and V-JEPA for video (2024).

Take an image. Choose a context block and several target blocks.

  encoder(context)                     -> context embeddings
  target_encoder(target blocks)        -> target embeddings
       (target_encoder is an EMA copy of encoder, not trained
        directly, and its gradients are stopped)

  predictor(context embeddings, target positions)
                                       -> predicted target embeddings

  loss = L2 distance in EMBEDDING space

Nothing is reconstructed. No pixels are predicted.

The argument is about where capacity goes. Predicting pixels forces a model to represent everything in the image, including detail that is unpredictable in principle: the exact arrangement of leaves, the grain of a texture, the precise position of a wave. A model trained on that objective spends real capacity on noise, because the loss punishes getting the noise wrong.

Predicting a representation does not. If the target encoder’s embedding of a patch of grass does not record blade positions, the predictor is never asked to guess them. The claim is that the same compute then goes to structure that is actually predictable — objects, layout, motion — which is the part you wanted.

The contrast with a masked autoencoder is exact and is the cleanest way to hold the idea: MAE masks patches and reconstructs pixels; JEPA masks patches and predicts embeddings. Same masking, different target space.

The failure mode, and the machinery against it

There is an obvious problem with a loss defined in embedding space: the encoder chooses that space. If it maps everything to a constant vector, the predictor learns to output that constant, the loss is exactly zero, and the model has learned nothing. This is representation collapse and it is the central difficulty of the whole approach.

The defences are architectural rather than principled:

• An EMA target encoder with stopped gradients. The target is produced by a slowly moving average of the online encoder, so the loss cannot be reduced by moving the target toward the prediction. The asymmetry is the point.
• A predictor, not a direct comparison. The extra network absorbs the positional part of the task and makes the trivial solution less reachable.
• Explicit variance and covariance terms in the related VICReg line of work: penalise a representation whose dimensions have low variance or are correlated.

These work empirically. What is honest to say is that they are empirical: this is a family whose central failure mode is prevented by design choices that were found to work rather than derived, and that is a fair summary of the state of the theory.

What has actually been shown

Precision matters here because this area attracts strong claims.

• Established: joint-embedding predictive methods are effective self-supervised visual representation learners. I-JEPA and V-JEPA produce features that transfer well to downstream image and video tasks with published results, and they train without the heavy hand-designed augmentations that earlier contrastive methods needed.
• Also true: the family is not new in shape. BYOL (2020) and data2vec (2022) already predicted targets produced by an EMA teacher. JEPA is a well-argued position within an existing line rather than a break from it.
• Not shown: that this is a replacement for next-token prediction in language models. There is no public result establishing an energy-based or joint-embedding architecture matching a strong autoregressive language model at equal compute, and the papers do not claim one.
• Open: whether a model that predicts in representation space can be made to generate well. Dropping the pixel objective is what saves the capacity, and it is also what removes the decoder. World models built on this idea need to produce something eventually.