Predictive Embeddings, JEPA, and Energy-Based Models › Energy Functions and Compatibility Learning

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:

$$p_\psi(y \mid x) = \frac{\exp\big(-E_\psi(x, y)\big)}{Z_\psi(x)}, \qquad Z_\psi(x) = \int \exp\big(-E_\psi(x, y')\big)\, d\mu(y') .$$

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

$$\nabla_\psi \log p_\psi(y \mid x) = -\nabla_\psi E_\psi(x, y) + \mathbb{E}_{y' \sim p_\psi(\cdot \mid x)}\big[\nabla_\psi E_\psi(x, y')\big].$$

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:

$$\mathcal{L}_{\mathrm{NCE}} = -\,\mathbb{E}\left[ \log \frac{\exp\big(-E_\psi(x, y^+)\big)}{\exp\big(-E_\psi(x, y^+)\big) + \sum_{j=1}^{K} \exp\big(-E_\psi(x, y_j^-)\big)} \right],$$

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

$$E_{\theta,\phi,\bar\theta}(x, y; c) = \big\| g_\phi(f_\theta(x), c) - \mathrm{sg}(f_{\bar\theta}(y)) \big\|_2^2 ,$$

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.