6. Energy-Based Models and Score Matching

Energy-based models

Unnormalized densities

The most flexible way to define a distribution with a network is to let the network assign a number to every input and make low numbers likely. An energy-based model (EBM) sets

\[p_\theta(x)=\frac{e^{-E_\theta(x)}}{Z_\theta},\qquad Z_\theta=\int e^{-E_\theta(x)}\,dx,\]

where the energy \(E_\theta\) can be any network with a scalar output (LeCun et al., 2006). There is no constraint of invertibility, no ordering of dimensions, and no latent variable to infer. Markov networks are energy-based models whose energy is a sum of potentials (AI chapter 8), and the Boltzmann machine of 1985 was one over binary units with pairwise interactions. The price of this flexibility is the partition function \(Z_\theta\), an integral over the whole space that cannot be computed for any interesting energy. Without it, the density cannot be evaluated, samples cannot be drawn directly, and even comparing \(p_\theta(x)\) with \(p_\theta(x')\) requires only the energy difference, which is all one can get.

The maximum-likelihood gradient

The log-likelihood is \(-E_\theta(x)-\log Z_\theta\), and its gradient, although \(Z_\theta\) is unknown, has a simple form:

\[\nabla_\theta\log p_\theta(x)=-\nabla_\theta E_\theta(x)+\mathbb E_{x'\sim p_\theta}\bigl[\nabla_\theta E_\theta(x')\bigr]\]

(Appendix A). The first term, the positive phase, lowers the energy of the data; the second, the negative phase, raises the energy of the model’s own samples. At the optimum, the two balance: the model’s samples and the data have the same expected energy gradient, the moment matching of log-linear models (AI chapter 14). The difficulty has moved from the partition function to the negative phase, which needs samples from the model at every training step, and the model can only be sampled by Markov chain Monte Carlo.

Sampling with Langevin dynamics

For continuous data, the natural sampler follows the gradient of the log-density with added noise. The Langevin update

\[x_{k+1}=x_k+\frac\epsilon2\nabla_x\log p_\theta(x_k)+\sqrt\epsilon\,\xi_k,\qquad\xi_k\sim\mathcal N(0,I),\]

is a discretized diffusion whose stationary distribution is \(p_\theta\) as \(\epsilon\to0\), and it needs only \(\nabla_x\log p_\theta=-\nabla_xE_\theta\), in which the partition function has disappeared. With a finite step size it is slightly biased, which a Metropolis–Hastings correction removes (AI chapter 10; Roberts and Tweedie, 1996). The same update with stochastic gradients of a posterior is a method for Bayesian learning (Welling and Teh, 2011).

Running a chain to convergence at every training step is too expensive, so EBM training uses short chains. Contrastive divergence starts each chain at a training example and runs a single step, or a few (Hinton, 2002), which trained restricted Boltzmann machines well enough to pretrain the first deep networks (Hinton, Osindero, and Teh, 2006); persistent contrastive divergence keeps the chains running across parameter updates instead of restarting them (Tieleman, 2008). For images, EBMs with convolutional energies trained with Langevin chains and a replay buffer of past samples (Du and Mordatch, 2019), or with short chains from noise that never converge (Nijkamp et al., 2019), generated recognizable images, but their training was unstable and their samples lagged behind those of other families.

Why Langevin sampling struggles

Langevin dynamics explores by following local gradients, and it inherits the weaknesses of every local sampler. Chains cross the low-density regions between modes rarely, so their samples reflect where they started rather than how much mass each mode holds. The problem is worse than slow mixing: in a region where one mode dominates, the score of a mixture is almost exactly the score of that mode alone, so the gradient carries almost no information about the relative weights of separated modes (Song and Ermon, 2019; Appendix C). The code samples a known mixture that puts 80% of its mass in the left mode, using its exact score.

import numpy as np

rng = np.random.default_rng(0)
# Target: 0.8 N(-4, 0.5^2) + 0.2 N(4, 0.5^2). Its score is known in closed form, at any added noise level sigma:
# the noisy density is the same mixture with variance 0.5^2 + sigma^2.
w, m, s = np.array([0.8, 0.2]), np.array([-4.0, 4.0]), 0.5


def score(x, sigma=0.0):
    v = s ** 2 + sigma ** 2
    comp = w * np.exp(-(x[:, None] - m) ** 2 / (2 * v))  # unnormalized responsibilities (equal variances)
    r = comp / comp.sum(1, keepdims=True)
    return (r * (m - x[:, None]) / v).sum(1)


def langevin(x, sigma, eps, steps):
    for _ in range(steps):
        x = x + 0.5 * eps * score(x, sigma) + np.sqrt(eps) * rng.standard_normal(len(x))
    return x


x0 = rng.uniform(-8, 8, 5000)                             # chains start spread over both modes
plain = langevin(x0.copy(), 0.0, 0.01, 2000)
print(f"plain Langevin, 2,000 steps:  share of samples in the left mode {np.mean(plain < 0):.3f} (target 0.800)")

x = x0.copy()                                             # annealed: large noise first, then smaller
sigmas = np.geomspace(10, 0.01, 20)
for sigma in sigmas:                                      # step size proportional to the noisy variance
    x = langevin(x, sigma, 0.2 * (s ** 2 + sigma ** 2), 100)
print(f"annealed Langevin, 20 levels: share of samples in the left mode {np.mean(x < 0):.3f}")
print(f"annealed samples: mean {x[x < 0].mean():.2f} and sd {x[x < 0].std():.2f} in the left mode (target -4.00, 0.50)")
# plain Langevin, 2,000 steps:  share of samples in the left mode 0.505 (target 0.800)
# annealed Langevin, 20 levels: share of samples in the left mode 0.785
# annealed samples: mean -4.00 and sd 0.51 in the left mode (target -4.00, 0.50)

After 2,000 steps from starting points spread evenly over both modes, the chains still split evenly between the modes, 50.5% on the left instead of 80%, although every chain has settled into the correct shape of its mode. The fix, which the second half of the code applies, uses the scores of noisy versions of the distribution. With a large amount of added Gaussian noise, the two modes blur into one broad distribution whose score does reflect the weights; as the noise is reduced in stages, each chain is guided into the mode that the larger-noise distribution assigned it. This annealed Langevin dynamics puts 78.5% of samples in the left mode, with the right mean and spread. Noise is what makes the score informative, and it is the idea on which the rest of this chapter and chapters 7–9 are built.

Noise-contrastive estimation

A different way around the partition function is to avoid sampling from the model at all. Noise-contrastive estimation (NCE; Gutmann and Hyvärinen, 2010) trains a logistic regression to distinguish data from samples of a known noise distribution \(q\), using \(\log p_\theta(x)-\log q(x)\) as the classifier’s logit, with \(\log Z_\theta\) treated as one more parameter. Since the optimal logit is the true log ratio (chapter 1), the fitted model is the data density, normalized. It works when the noise is close enough to the data that the classification is hard; in high dimensions a simple noise distribution is too easy to tell apart, and the classifier learns little. The negative sampling used to train word vectors is a simplified version (NLP chapter 3).

Score matching

The score function

The score of a distribution is the gradient of its log-density with respect to the data, \(s(x)=\nabla_x\log p(x)\), a vector field that points toward higher density. For an energy-based model it is \(-\nabla_xE_\theta(x)\), independent of the partition function. The score determines the distribution on a connected domain, since integrating it recovers \(\log p\) up to the constant fixed by normalization, and it is all that Langevin dynamics needs. So instead of modeling the density, one can model the score directly with a network \(s_\theta:\mathbb R^D\to\mathbb R^D\), and fit it by minimizing the Fisher divergence,

\[\frac12\,\mathbb E_{p_{\mathrm{data}}}\bigl\|s_\theta(x)-\nabla_x\log p_{\mathrm{data}}(x)\bigr\|^2 .\]

The target score is unknown, since estimating it is the whole problem, but the objective can be rewritten without it.

Hyvärinen’s score matching

Integrating by parts, and assuming the density vanishes at infinity, the Fisher divergence equals, up to a constant,

\[\mathbb E_{p_{\mathrm{data}}}\Bigl[\operatorname{tr}\bigl(\nabla_xs_\theta(x)\bigr)+\frac12\|s_\theta(x)\|^2\Bigr],\]

which involves only the model and the data (Hyvärinen, 2005; Appendix B). The estimator is consistent, but the trace of the Jacobian costs \(D\) backward passes, as in continuous flows (chapter 4). Sliced score matching replaces the trace with \(v^\top\nabla_xs_\theta\,v\) for random directions \(v\), the same random-projection trick as Hutchinson’s estimator (Song et al., 2019).

Denoising score matching

A cheaper route comes from denoising. Perturb each data point with Gaussian noise, \(\tilde x=x+\sigma\epsilon\) with \(\epsilon\sim\mathcal N(0,I)\). The score of the conditional distribution of \(\tilde x\) given \(x\) is known exactly, \(\nabla_{\tilde x}\log q_\sigma(\tilde x\mid x)=-(\tilde x-x)/\sigma^2=-\epsilon/\sigma\), and Vincent (2011) showed that regressing a network onto it,

\[\mathbb E_{x,\epsilon}\Bigl\|s_\theta(x+\sigma\epsilon)+\frac\epsilon\sigma\Bigr\|^2,\]

has the same minimizer as matching the score of the noisy data distribution \(p_\sigma=p_{\mathrm{data}}*\mathcal N(0,\sigma^2I)\) (Appendix B). Denoising score matching needs one forward pass per example and no Jacobian. Its target is the noise itself, scaled: the network learns to predict which way the added noise pushed each point, which is the denoising autoencoder’s task (DL chapter 10). The connection is exact in the other direction as well: Tweedie’s formula gives the best denoiser in terms of the noisy score,

\[\mathbb E[x\mid\tilde x]=\tilde x+\sigma^2\nabla_{\tilde x}\log p_\sigma(\tilde x)\]

(Efron, 2011), so learning to denoise and learning the score of noisy data are the same problem. The code learns the score of eight Gaussians on a circle, blurred with noise of standard deviation 0.3, compares it with the exact score, and uses it to denoise.

import numpy as np
import torch
from torch import nn

torch.manual_seed(0)
# Data: eight Gaussians (sd 0.1) on a circle of radius 2. Noise level for denoising score matching: sigma = 0.3.
angles = torch.arange(8) * (2 * np.pi / 8)
centers = torch.stack([torch.cos(angles), torch.sin(angles)], 1) * 2.0
sd, sigma = 0.1, 0.3
data = lambda n: centers[torch.randint(8, (n,))] + sd * torch.randn(n, 2)


def true_score(x, var):                                   # score of the mixture with component variance var
    logits = -((x[:, None, :] - centers) ** 2).sum(-1) / (2 * var)
    r = torch.softmax(logits, 1)
    return (r[:, :, None] * (centers - x[:, None, :])).sum(1) / var


net = nn.Sequential(nn.Linear(2, 128), nn.SiLU(), nn.Linear(128, 128), nn.SiLU(), nn.Linear(128, 2))
opt = torch.optim.Adam(net.parameters(), lr=1e-3)
for step in range(4000):
    x = data(512)
    noise = torch.randn_like(x)
    x_noisy = x + sigma * noise
    # Denoising score matching: regress the score onto -(x_noisy - x) / sigma^2 = -noise / sigma.
    loss = ((net(x_noisy) + noise / sigma) ** 2).sum(1).mean() * sigma ** 2
    opt.zero_grad()
    loss.backward()
    opt.step()

with torch.no_grad():
    var = sd ** 2 + sigma ** 2                            # the noisy data are the same mixture, wider
    box = 8 * torch.rand(20000, 2) - 4
    for name, pts in [("at noisy data points", data(5000) + sigma * torch.randn(5000, 2)),
                      ("inside the circle, radius < 1", box[box.norm(dim=1) < 1]),
                      ("outside, radius > 3.5", box[box.norm(dim=1) > 3.5])]:
        err = (net(pts) - true_score(pts, var)).norm(dim=1) / true_score(pts, var).norm(dim=1)
        print(f"relative error of the learned score {name:30s}: median {err.median():.3f}")
    x = data(5000)
    x_noisy = x + sigma * torch.randn_like(x)
    denoised = x_noisy + sigma ** 2 * net(x_noisy)       # Tweedie: E[x | x_noisy] = x_noisy + sigma^2 score
    best = x_noisy + sigma ** 2 * true_score(x_noisy, var)
    print(f"mean squared distance to the clean points: noisy {((x_noisy - x) ** 2).sum(1).mean():.4f}, "
          f"denoised with the learned score {((denoised - x) ** 2).sum(1).mean():.4f}, "
          f"with the exact score {((best - x) ** 2).sum(1).mean():.4f}")
# relative error of the learned score at noisy data points          : median 0.083
# relative error of the learned score inside the circle, radius < 1 : median 0.198
# relative error of the learned score outside, radius > 3.5         : median 0.061
# mean squared distance to the clean points: noisy 0.1798, denoised with the learned score 0.0434, with the exact score 0.0419

At noisy data points the learned score is within about 8% of the exact one. Inside the circle, where no data fall and the true score turns sharply between competing modes, its error is more than twice as large, about 20%; outside, where the score simply points back toward the nearest mode, the network extrapolates well. Applying Tweedie’s formula with the learned score reduces the mean squared distance to the clean points from 0.180, the noise’s variance times two dimensions, to 0.043, nearly the 0.042 that the exact score achieves, which is the best any denoiser can do on average.

Noise-conditional score networks

Why one noise level is not enough

A single noise level forces a compromise. Small noise gives a score close to that of the data, but the data are concentrated near a thin set (chapter 1), so the noisy data leave most of the space empty and the score there is learned from nothing; the score of the clean data is not even defined off the data manifold. Large noise fills the space and makes the score easy to learn everywhere, but it describes a blurred distribution. And as the code showed, Langevin dynamics with the score of the nearly clean distribution cannot find the right proportions of separated modes.

Noise-conditional score networks (NCSN; Song and Ermon, 2019) use many noise levels at once. One network \(s_\theta(x,\sigma)\) takes the noise level as an input and is trained by denoising score matching at every level of a geometric sequence \(\sigma_1>\sigma_2>\dots>\sigma_L\), with the loss at level \(\sigma\) weighted by \(\sigma^2\) so that all levels contribute comparably. Sampling runs annealed Langevin dynamics: a few steps with the score at \(\sigma_1\), starting from noise, then a few at \(\sigma_2\) from where those ended, and so on down to \(\sigma_L\), with the step size shrinking in proportion to \(\sigma^2\). The large-noise stages move samples across the space and apportion them among the modes; the small-noise stages refine them. With ten noise levels from 1 down to 0.01, NCSN reached an Inception score of 8.87 on CIFAR-10, then the best for unconditional generation, and an improved version with a largest noise level as large as the greatest distance between training images generated faces and scenes at up to \(256\times256\) (Song and Ermon, 2020). The figure trains a small noise-conditional score network on eight modes of unequal weight.

A noise-conditional score network trained by denoising score matching at 12 noise levels from 3 down to 0.05 on eight Gaussians (standard deviation 0.1) whose weights grow from 1/36 at the right, counterclockwise, to 8/36 at the lower right; orange disks are sized by weight. First two panels: the direction of the learned score on a grid at σ = 0.05, where each arrow points to the nearest mode, and at σ = 0.98, where the field pulls toward the heavier modes of the lower half. Last two panels: 4,000 samples from uniform starting points, by Langevin dynamics with the smallest-noise score alone (1,200 steps) and by annealed Langevin dynamics (100 steps per level). Both put 96.5% of samples near a mode, but plain Langevin gives each mode nearly the same share, a total-variation distance of 0.223 from the true weights, while annealing reduces it to 0.074.

From noise levels to diffusion

A noise-conditional score network is a diffusion model in all but name. Chapter 7 arrives at the same model from a different direction, as a latent-variable model that learns to reverse a gradual noising process, and shows that its training objective is denoising score matching summed over noise levels with a particular weighting. Chapter 8 then takes the number of noise levels to infinity: the noising becomes a stochastic differential equation, annealed Langevin dynamics becomes the reverse-time equation, and a deterministic version of the sampler turns the score model into a continuous normalizing flow with exact likelihoods. Energy-based models themselves remain useful where an explicit, composable energy matters, for example adding the energies of several models to combine their constraints, and a trained classifier can be read as an energy-based model of its inputs (Grathwohl et al., 2020).

Appendices