DEV Community

Umesh Tharuka Malaviarachchi
Umesh Tharuka Malaviarachchi

Posted on

Energy-Based Models: Why an Unnormalised Density Can Still Define a Distribution

A probability density is usually introduced as a function that integrates to one. That normalisation condition feels non-negotiable: if the total probability is not one, the numbers are not probabilities. Energy-based models deliberately break that habit. They define a distribution through an unnormalised score, accept that the normalising constant is unknown and often intractable, and still learn, sample, and make decisions. The trick is that many of the operations we actually care about do not need the constant at all.

This is not a historical curiosity. Energy-based models are the conceptual parent of score matching, of much of modern diffusion modelling, and of a family of methods for structured prediction where a flexible score is easier to write down than a normalised density. Understanding them clarifies what a density really is, what training by maximum likelihood really costs, and why sampling is often the hard part of learning.

The energy view of a distribution

An energy-based model assigns every configuration an energy, a single real number. Low energy means plausible, high energy means implausible. The model turns energy into an unnormalised density by exponentiating its negative:

p~θ(x)=exp⁡(−Eθ(x)) \tilde{p}{\theta}(x) = \exp(-E{\theta}(x))

Here Eθ(x)E_{\theta}(x) is the energy function with parameters θ\theta , and the tilde marks the fact that p~θ\tilde{p}_{\theta} is not yet a density. To make it one, divide by the total mass over all configurations:

pθ(x)=1Zθexp⁡(−Eθ(x)),Zθ=∫exp⁡(−Eθ(x)) dx p_{\theta}(x) = \frac{1}{Z_{\theta}} \exp(-E_{\theta}(x)), \quad Z_{\theta} = \int \exp(-E_{\theta}(x)) \, dx

The constant ZθZ_{\theta} is called the partition function. The name comes from statistical physics, where the same construction describes systems in thermal equilibrium: states with lower energy are exponentially more likely. In machine learning, the energy function is usually a neural network, so EθE_{\theta} can be almost any flexible function of xx . That flexibility is the appeal. You do not need to design a density whose integral you can compute in closed form. You only need to design a score that is low where data live and high elsewhere.

The price is immediate. For a high-dimensional xx , such as an image, the integral defining ZθZ_{\theta} has no closed form and cannot be computed by brute force. The model is therefore defined by a quantity we cannot evaluate. Much of the theory of energy-based models is about working around that fact without pretending it away.

Why the partition function matters, and when it does not

Maximum likelihood training makes the role of ZθZ_{\theta} explicit. The log-likelihood of a data point is:

log⁡pθ(x)=−Eθ(x)−log⁡Zθ \log p_{\theta}(x) = -E_{\theta}(x) - \log Z_{\theta}

Differentiating with respect to θ\theta gives the classic two-term gradient:

∇θlog⁡pθ(x)=−∇θEθ(x)+Ex′∼pθ[∇θEθ(x′)] \nabla_{\theta} \log p_{\theta}(x) = -\nabla_{\theta} E_{\theta}(x) + \mathbb{E}{x' \sim p{\theta}}[\nabla_{\theta} E_{\theta}(x')]

The first term is the positive phase: it pushes the energy of observed data down. The second term is the negative phase: it pushes the energy of samples drawn from the model itself up. Training is a tug of war. Without the negative phase, the model would simply assign low energy to everything. Without the positive phase, it would assign high energy to everything. At the optimum, the two phases balance, and the model's samples look like the data.

That gradient is elegant and, in its exact form, useless. The expectation is over the model's own distribution, which we cannot sample from exactly and whose normalising constant we cannot compute. Every practical training method for energy-based models is a way of approximating that negative phase, or of avoiding it altogether.

Some operations, however, never need ZθZ_{\theta} . Compare two configurations under the same model:

log⁡pθ(x1)−log⁡pθ(x2)=−Eθ(x1)+Eθ(x2) \log p_{\theta}(x_1) - \log p_{\theta}(x_2) = -E_{\theta}(x_1) + E_{\theta}(x_2)

The partition function cancels. Ranking, classification by relative score, and optimisation over xx for a fixed θ\theta all work with energies alone. The gradient of the log-density with respect to xx , known as the score, also cancels it:

∇xlog⁡pθ(x)=−∇xEθ(x) \nabla_x \log p_{\theta}(x) = -\nabla_x E_{\theta}(x)

This observation is the bridge to score-based methods. If you only need the score, you never need to know ZθZ_{\theta} . Score matching trains a model to match the data score directly, and diffusion models learn scores at many noise levels. Both inherit the flexibility of energy-based models while sidestepping the partition function.

Sampling with Langevin dynamics

If the negative phase needs samples from the model, how do we get them when we cannot even normalise the density? The standard answer is Markov chain Monte Carlo, and in particular Langevin dynamics, which uses only the score. Starting from some x0x_0 , iterate:

xt+1=xt+ϵ2∇xlog⁡pθ(xt)+ϵ zt,zt∼N(0,I) x_{t+1} = x_t + \frac{\epsilon}{2} \nabla_x \log p_{\theta}(x_t) + \sqrt{\epsilon} \, z_t, \quad z_t \sim \mathcal{N}(0, I)

Each step climbs uphill in log-density, following the score, and then adds Gaussian noise. The noise is essential: without it, the chain would collapse to a local maximum of the density. With the right step size and enough steps, the chain's stationary distribution is pθp_{\theta} , even though ZθZ_{\theta} never appears in the update. In energy terms, the update is gradient descent on energy plus noise, which is why samples wander toward low-energy regions but keep exploring.

A minimal PyTorch sketch makes the mechanics concrete. Suppose energy is a neural network returning a scalar for a batch of inputs:

import torch

def langevin_sample(energy, x_init, steps=100, step_size=0.01):
    x = x_init.clone().detach().requires_grad_(True)
    for _ in range(steps):
        e = energy(x).sum()
        (grad,) = torch.autograd.grad(e, x)
        # Score is minus the energy gradient.
        noise = torch.randn_like(x) * (step_size ** 0.5)
        with torch.no_grad():
            x = x - 0.5 * step_size * grad + noise
        x = x.detach().requires_grad_(True)
    return x.detach()
Enter fullscreen mode Exit fullscreen mode

This code is short, and that brevity hides the difficulty. Langevin chains mix slowly when the distribution is multimodal, because moving between modes requires crossing high-energy barriers that the gradient actively pushes away from. Step size is a genuine trade-off: too large and the discretised chain diverges or produces biased samples, too small and it barely moves. In practice, chains are run for many steps, often from a replay buffer of past samples, and the samples are still approximate. The negative phase of training is therefore approximate too, and that approximation error is a first-class part of the method, not an implementation detail.

Contrastive divergence and its descendants

Contrastive divergence, originally developed for restricted Boltzmann machines, is the most famous shortcut for the negative phase. Instead of running a chain to convergence, run it for only a few steps starting from a data point, and use the resulting sample as if it came from the model. The gradient estimate becomes:

∇θlog⁡pθ(x)≈−∇θEθ(x)+∇θEθ(xk) \nabla_{\theta} \log p_{\theta}(x) \approx -\nabla_{\theta} E_{\theta}(x) + \nabla_{\theta} E_{\theta}(x_k)

where xkx_k is the point reached after kk MCMC steps from xx . With k=1k = 1 , this is CD-1. It is biased: the chain has not converged, so xkx_k is not really a model sample. Yet it often works well in practice, because the bias pushes the model in roughly the right direction: raise the energy of points the chain drifts toward, which are typically near the data but not quite on it.

Persistent contrastive divergence improves on this by keeping chains alive across parameter updates instead of restarting them from data each time. As the parameters change slowly, the chains track the evolving distribution, so a few steps per update are closer to convergence than a cold start would be. Modern energy-based image models use this idea with replay buffers and careful regularisation of the energy landscape, and they can generate recognisable images from an unnormalised model trained this way.

The lesson generalises. Training an energy-based model is a coupled system: the sampler and the energy function co-evolve. If the sampler is poor, the gradient is wrong, and the energy landscape learns to exploit the sampler's blind spots. A model that looks good under short chains can collapse under long ones, producing samples that drift to spurious low-energy regions the short chains never visited. Checking a trained energy-based model therefore means sampling with longer, better-mixed chains than the ones used in training, not just reporting training likelihood estimates.

A worked intuition: a two-mode density

Consider data drawn from two well-separated clusters on a line. A Gaussian mixture could fit this easily, but suppose we instead parameterise an energy as a small neural network and train by contrastive divergence. Early in training, the energy is roughly flat, so Langevin chains wander freely and the negative phase raises energy almost everywhere except near the data, where the positive phase lowers it. Two basins form.

Now suppose the chains used in training are short and start from data points. A chain starting in the left cluster rarely crosses to the right cluster, because the region between them has high energy and the gradient pushes the chain back. The negative phase therefore never asks the model to raise the energy of the valley between the clusters relative to the clusters themselves in a balanced way, and the learned energy may leave a spurious low-energy bridge, or an entire spurious third basin that short chains never visit. A longer chain at test time, or a chain started from noise, may fall into that basin and produce samples that look nothing like the data, even though the training objective appeared to converge.

This is not a reason to avoid energy-based models. It is a reason to treat sampling diagnostics as part of model evaluation: run chains from multiple starting points, check that samples cover all data modes in roughly the right proportions, and distrust a model whose samples depend strongly on chain length. The same diagnostic discipline applies, with different details, to diffusion models and other score-based generators.

Where energy-based models still win

Normalised generative models such as normalizing flows and autoregressive models pay for tractability with architectural constraints: invertibility, triangular Jacobians, or a fixed generation order. Energy-based models pay with an intractable constant instead, and keep full architectural freedom. That trade is attractive in several settings.

First, structured prediction. If the output is a segmentation, a parse, or a set of interacting labels, writing a flexible energy over the whole output is often natural, while writing a normalised distribution over it is hard. Inference then means minimising energy, which needs no partition function.

Second, composition. Energies add. If you have one energy that prefers realistic images and another that prefers images matching a text prompt, their sum is an energy for images satisfying both, up to a new normalising constant you again do not need for sampling by Langevin dynamics. This compositional property underlies classifier guidance and related conditioning tricks, where a classifier's log-probability is added to a generative score.

Third, anomaly detection and density-ratio reasoning. An energy trained on normal data assigns high energy to outliers, giving a usable anomaly score without ever computing a calibrated probability. The score is only meaningful relative to the same model, so thresholds must be calibrated on held-out data, but the ranking is often what matters.

The cost remains sampling and evaluation. You cannot report an exact held-out likelihood for an energy-based model without estimating ZθZ_{\theta} , for example by annealed importance sampling, and such estimates are themselves noisy in high dimensions. Comparisons between energy-based models and normalised models on likelihood alone are therefore treacherous. Compare them on samples, on downstream tasks, and on the cost of getting those samples, not on a number that one side can compute exactly and the other can only estimate.

Limitations and common pitfalls

Several failure modes recur often enough to name explicitly.

Spurious low-energy regions are the first. An energy function is only constrained where training data and training chains visit. Everywhere else, its value is an extrapolation of the network, and it can dip arbitrarily low. Regularisation, spectral normalisation, and gradient penalties on the energy are commonly used to keep the landscape sane, but none of them removes the issue entirely.

MCMC mixing is the second. Langevin dynamics struggles with multimodality and with sharp, narrow modes. If your sampler cannot move between modes, your negative phase cannot teach the model the correct relative weights of those modes, and generated samples will over-represent whichever mode the chains happen to start near.

Evaluation is the third. It is tempting to report samples from short chains, or to cherry-pick attractive samples, because long-chain samples are worse. Resist that. The honest evaluation of an energy-based model includes long chains, multiple initialisations, and quantitative checks such as mode coverage and held-out energy rankings, alongside qualitative samples.

A fourth pitfall is confusing the energy with a probability. Energies are not calibrated, not comparable across models, and not meaningful in absolute terms. A statement like "the energy is 3.2" carries no information without a reference point in the same model. Only differences of energy, gradients of energy, and samples drawn with a stated sampler are interpretable.

Finally, do not confuse an energy-based model with a score model. A score model parameterises ∇xlog⁡p(x)\nabla_x \log p(x) directly and need not correspond to the gradient of any single energy function, since a general vector field need not be conservative. Energy-based models are the conservative special case, which buys composition and a well-defined density at the cost of differentiating through the energy network. Score models buy flexibility at the cost of that guarantee. Modern diffusion models are closer to score models, and their success does not imply that the underlying energy view is obsolete: it is the view that explains why the score is sufficient in the first place.

Takeaway

Energy-based models separate two questions that normalised models fuse: what configurations are plausible, and what the total probability mass is. The energy answers the first question, the partition function answers the second, and many tasks need only the first. That separation buys architectural freedom and compositionality, and it costs an intractable constant that forces approximate sampling into the heart of training. If you learn one thing from energy-based models, let it be this: a generative model is not just a density, it is a density plus a sampler, and the sampler's quality is part of the model.


Connect with me on LinkedIn: linkedin.com/in/umeshtharukaofficial

Top comments (0)