How to Train Your Energy-Based Models

Introduction

Probabilistic models with a tractable likelihood are a double-edged sword. On one hand, a tractable likelihood allows for straightforward comparison between models, and straightforward optimization of the model parameters w.r.t. the log-likelihood of the data. Through tractable models such as autoregressive (Graves, 2013; Germain et al., 2015; Van Oord et al., 2016) or flow-based generative models (Dinh et al., 2014, 2016; Rezende and Mohamed, 2015), we can learn flexible models of high-dimensional data. In some cases even though the likelihood is not completely tractable, we can often compute and optimize a tractable lower bound of the likelihood, as in the framework of variational autoencoders (Kingma and Welling, 2014; Rezende et al., 2014).

Still, the set of models with a tractable likelihood is constrained. Models with a tractable likelihood need to be of a certain form: for example, in case of autoregressive models, the model distribution is factorized as a product of conditional distributions, and in flow-based generative models the data is modeled as an invertible transformation of a base distribution. In case of variational autoencoders, the data must be modeled as a directed latent-variable model. A tractable likelihood is related to the fact that these models assume that exact synthesis of pseudo-data from the model can be done with a specified, tractable procedure. These assumptions are not always natural.

Energy-based models (EBM) are much less restrictive in functional form: instead of specifying a normalized probability, they only specify the unnormalized negative log-probability, called the energy function. Since the energy function does not need to integrate to one, it can be parameterized with any nonlinear regression function. In the framework of EBMs, density estimation is thus basically reduced to a nonlinear regression problem. One is generally free to choose any nonlinear regression function as the energy function, as long as it remains normalizable in principle. It is thus straightforward to leverage advances in architectures originally developed for classification or regression, and one may choose special-purpose architectures that make sense for the type of data at hand. For example, If the data are graphs (such as molecules) then one could use graph neural networks (Scarselli et al., 2008); if the data are spherical images one can in principle use spherical CNNs (Cohen et al., 2018). As such, EBMs have found wide applications in many fields of machine learning, including, among others, image generation (Ngiam et al., 2011; Xie et al., 2016; Du and Mordatch, 2019), discriminative learning (Grathwohl et al., 2020b; Gustafsson et al., 2020a, b), natural language processing (Mikolov et al., 2013; Deng et al., 2020), density estimation (Wenliang et al., 2019; Song et al., 2019) and reinforcement learning (Haarnoja et al., 2017, 2018).

Although this flexibility of EBMs can provide significant modeling advantages, both computation of the exact likelihood and exact synthesis of samples from these models are generally intractable, which makes training especially difficult. There are three major ways for training EBMs: (i) maximum likelihood training with MCMC sampling; (ii) Score Matching (SM); and (iii) Noise Constrastive Estimation (NCE). We will elaborate on these methods in order, explain their relationships to each other, and conclude by an overview to other directions for EBM training.

Energy-Based Models (EBMs)

For simplicity we will assume unconditional Energy-Based Models over a single dependent variable ${\mathbf{x}}$. It is relatively straightforward to extend the models and estimation procedures to the case with multiple dependent variables, or with conditioning variables. The density given by an EBM is

where $E_{{\bm{\theta}}}({\mathbf{x}})$ (the energy) is a nonlinear regression function with parameters ${\bm{\theta}}$, and $Z_{{\bm{\theta}}}$ denotes the normalizing constant (a.k.a. the partition function):

which is constant w.r.t ${\mathbf{x}}$ but is a function of ${\bm{\theta}}$. Since $Z_{{\bm{\theta}}}$ is a function of ${\bm{\theta}}$, evaluation and differentiation of $\log p_{{\bm{\theta}}}({\mathbf{x}})$ w.r.t. its parameters involves a typically intractable integral.

Maximum Likelihood Training with MCMC

We cannot directly compute the likelihood of an EBM as in the maximum likelihood approach due to the intractable normalizing constant $Z_{\bm{\theta}}$. Nevertheless, we can still estimate the gradient of the log-likelihood with MCMC approaches, allowing for likelihood maximization with gradient ascent (Younes, 1999). In particular, the gradient of the log-probability of an EBM (Eq. 1) decomposes as a sum of two terms:

The first gradient term, $-\nabla_{{\bm{\theta}}}E_{{\bm{\theta}}}({\mathbf{x}})$, is straightforward to evaluate with automatic differentiation. The challenge is in approximating the second gradient term, $\nabla_{{\bm{\theta}}}\log Z_{{\bm{\theta}}}$, which is intractable to compute exactly. This gradient term can be rewritten as the following expectation:

where steps (i) and (ii) are due to the chain rule of gradients, and (iii) and (iv) are from definitions in Eqs. 1 and 2. Thus, we can obtain an unbiased one-sample Monte Carlo estimate of the log-likelihood gradient by

Since drawing random samples is far from being trivial, much of the literature has focused on methods for efficient MCMC sampling from EBMs. Some efficient MCMC methods, such as Langevin MCMC (Parisi, 1981; Grenander and Miller, 1994) and Hamiltonian Monte Carlo (Duane et al., 1987; Neal et al., 2011), make use of the fact that the gradient of the log-probability w.r.t. ${\mathbf{x}}$ (a.k.a., score) is equal to the (negative) gradient of the energy, therefore easy to calculate:

For example, when using Langevin MCMC to sample from $p_{{\bm{\theta}}}({\mathbf{x}})$, we first draw an initial sample ${\mathbf{x}}^{0}$ from a simple prior distribution, and then simulate an (overdamped) Langevin diffusion process for $K$ steps with step size $\epsilon>0$:

When $\epsilon\to 0$ and $K\to\infty$, ${\mathbf{x}}^{K}$ is guaranteed to distribute as $p_{{\bm{\theta}}}({\mathbf{x}})$ under some regularity conditions. In practice we have to use a small finite $\epsilon$, but the discretization error is typically negligible, or can be corrected with a Metropolis-Hastings (Hastings, 1970) step, leading to the Metropolis-Adjusted Langevin Algorithm (Besag, 1994).

Running MCMC till convergence to obtain a sample ${\mathbf{x}}\sim p_{{\bm{\theta}}}({\mathbf{x}})$ can be computationally expensive. Therefore we typically need approximation to make MCMC-based learning of EBMs practical. One popular method for doing so is Contrastive Divergence (CD) (Hinton, 2002). In CD, one initializes the MCMC chain from the datapoint ${\mathbf{x}}$, and perform a fixed number of MCMC steps; typically fewer than required for convergence of the MCMC chain. One variant of CD that sometimes performs better is persistent CD (Tieleman, 2008), where a single MCMC chain with a persistent state is employed to sample from the EBM. In persistent CD, we do not restart the MCMC chain when training on a new datapoint; rather, we carry over the state of the previous MCMC chain and use it to initialize a new MCMC chain for the next training step. This method can be further improved by keeping multiple historical states of the MCMC chain in a replay buffer and initialize new MCMC chains by randomly sampling from it (Du and Mordatch, 2019). Other variants of CD include mean field CD (Welling and Hinton, 2002), and multi-grid CD (Gao et al., 2018).

EBMs trained with CD may not capture the data distribution faithfully, since truncated MCMC can lead to biased gradient updates that hurt the learning dynamics (Schulz et al., 2010; Fischer and Igel, 2010; Nijkamp et al., 2019). There are several methods that focus on removing this bias for improved MCMC training. For example, one line of work proposes unbiased estimators of the gradient through coupled MCMC (Jacob et al., 2017; Qiu et al., 2019); and Du et al. (2020) propose to reduce the bias by differentiating through the MCMC sampling algorithm and estimating an entropy correction term.

Score Matching (SM)

where $d$ is the dimensionality of ${\mathbf{x}}$. The constant does not affect optimization and thus can be dropped for training. It is shown by Hyvärinen (2005) that estimators based on Score Matching are consistent under some regularity conditions, meaning that the parameter estimator obtained by minimizing Eq. 7 converges to the true parameters in the limit of infinite data.

An important downside of the objective Eq. 8 is that, in general, computation of full second derivatives is quadratic in the dimensionality $d$, thus does not scale to high dimensionality. Although SM only requires the trace of the Hessian, it is still expensive to compute even with modern hardware and automatic differentiation packages (Martens et al., 2012). For this reason, the implicit SM formulation of Eq. 8 has only been applied to relatively simple energy functions where computation of the second derivatives is tractable.

Score Matching assumes a continuous data distribution with positive density over the space, but it can be generalized to discrete or bounded data distributions (Hyvärinen, 2007; Lyu, 2012). It is also possible to consider higher-order gradients of log-PDFs beyond first derivatives (Parry et al., 2012).

Vincent (2011) propose one elegant and scalable solution to the above difficulty, by showing that:

When estimating the above expectation with samples, the variances of $({\mathbf{z}}^{(i)})^{\mkern-1.5mu\mathsf{T}}\nabla_{\mathbf{x}}\log p_{{\bm{\theta}}}({\mathbf{x}}^{(i)})/\sigma$ and $\left\lVert{\mathbf{z}}^{(i)}\right\rVert_{2}^{2}/\sigma^{2}$ will both grow unbounded as $\sigma\to 0$ due to division by $\sigma$ and $\sigma^{2}$. This enlarges the variance of DSM and makes optimization challenging.

Substracting it from Eq. 11 will yield an estimator with reduced variance for DSM training:

Variance-reducing variables like $c_{\bm{\theta}}({\mathbf{x}},{\mathbf{z}})$ are called control variates (Owen, 2013). For large $\sigma$, the Taylor series might be a bad approximator, in which case $c_{\bm{\theta}}({\mathbf{x}}^{(i)},{\mathbf{z}}^{(i)})$ might not be sufficiently correlated with Eq. 11, and could actually increase variance.

2 Sliced Score Matching (SSM)

Sliced Score Matching (SSM) (Song et al., 2019) is one alternative to Denoising Score Matching that is both consistent and computationally efficient. Instead of minimizing the Fisher divergence between two vector-valued scores, SSM randomly samples a projection vector $\mathbf{v}$, takes the inner product between $\mathbf{v}$ and the two scores, and then compare the resulting two scalars. More specifically, Sliced Score Matching minimizes the following divergence called the sliced Fisher divergence

All expectations in the above objective can be estimated with empirical means, and again the constant term can be removed without affecting training. The second term involves second-order derivatives of $E_{{\bm{\theta}}}({\mathbf{x}})$, but contrary to SM, it can be computed efficiently with a cost linear in the dimensionality $d$. This is because

where $f({\mathbf{x}})$ is the same for different values of $i$. Therefore, we only need to compute it once with $O(d)$ computation, plus another $O(d)$ computation for the outer sum to evaluate Eq. 16, whereas the original SM objective requires $O(d^{2})$ computation.

For many choices of $p({\mathbf{v}})$, part of the SSM objective (Eq. 15) can be evaluated in closed form, potentially leading to lower variance. For example, when $p({\mathbf{v}})=\mathcal{N}(\bm{0},{\bm{I}})$, we have

The above objective Eq. 17 can also be obtained by approximating the sum of second-order gradients in the standard SM objective (Eq. 8) with the Skilling-Hutchinson trace estimator (Skilling, 1989; Hutchinson, 1989). It often (but not always) has lower variance than Eq. 15, and can perform better in some applications (Song et al., 2019).

3 Connection to Contrastive Divergence

Though Score Matching and Contrastive Divergence are seemingly very different approaches, they are closely connected to each other. In fact, Score Matching can be viewed as a special instance of Contrastive Divergence in the limit of a particular MCMC sampler (Hyvarinen, 2007). Moreover, the Fisher divergence optimized by Score Matching is related to the derivative of KL divergence (Cover, 1999), which is the underlying objective of Contrastive Divergence.

Contrastive Divergence requires sampling from the Energy-Based Model $E_{{\bm{\theta}}}({\mathbf{x}})$, and one popular method for doing so is Langevin MCMC (Parisi, 1981). Recall from Section 3 that given any initial data point ${\mathbf{x}}^{0}$, the Langevin MCMC method executes the following

iteratively for $k=0,1,\cdots,K-1$, where ${\mathbf{z}}^{k}\sim\mathcal{N}(\bm{0},{\bm{I}})$ and $\epsilon>0$ is the step size.

Suppose we only run one-step Langevin MCMC (i.e., $K=1$) for Contrastive Divergence. In this case, the gradient of the log-likelihood is given by

After Taylor series expansion with respect to $\epsilon$ followed by some algebraic manipulations, the above equation can be transformed to the following (see Hyvarinen (2007))

When $\epsilon$ is sufficiently small, it corresponds to the re-scaled gradient of the Score Matching objective.

4 Score-Based Generative Models

One typical application of EBMs is creating new samples that are similar to training data. Towards this end, we can first train an EBM with Score Matching, and then sample from it with MCMC approaches. Many efficient sampling methods for EBMs, such as Langevin MCMC, rely on just the score of the EBM (see Eq. 6). In addition, Score Matching objectives (Eqs. 8, 10 and 15) depend sorely on the scores of EBMs. Therefore, we only need a model for score when training with Score Matching and sampling with score-based MCMC, and do not have to model the energy explicitly. By building such a score model, we save the gradient computation of EBMs and can make training and sampling more efficient. These kind of models are named score-based generative models (Song and Ermon, 2019, 2020; Song et al., 2021).

Song and Ermon (2019, 2020) and Song et al. (2021) overcome this difficulty by perturbing training data with different scales of noise, and learn a score model for each scale. For a large noise perturbation, different modes are connected due to added noise, and estimated weights between them are therefore accurate. For a small noise perturbation, different modes are more disconnected, but the noise-perturbed distribution is closer to the original unperturbed data distribution. Using a sampling method such as annealed Langevin dynamics (Song and Ermon, 2019, 2020; Song et al., 2021) or leveraging reverse diffusion processes (Sohl-Dickstein et al., 2015; Ho et al., 2020; Song et al., 2021), we can sample from the most noise-perturbed distribution first, then smoothly reduce the magnitude of noise scales until reaching the smallest one. This procedure helps combine information from all noise scales, and maintains the correct portion of modes from larger noise perturbations when sampling from smaller ones. In practice, all score models share weights and are implemented with a single neural network conditioned on the noise scale, named a Noise-Conditional Score Network. Scores of different scales are estimated by training a mixture of Score Matching objectives, one per noise scale. This method are amongst the best generative modeling approaches for high-resolution image generation (see samples in Fig. 1), audio synthesis (Chen et al., 2020; Kong et al., 2020), and shape generation (Cai et al., 2020).

Noise Contrastive Estimation

A third principle for learning the parameters of EBMs is Noise Contrastive Estimation (NCE), introduced by Gutmann and Hyvärinen (2010). It is based on the idea that we can learn an Energy-Based Model by contrasting it with another distribution with known density.

Suppose our Energy-Based Model $p_{{\bm{\theta}}}({\mathbf{x}})$ has the form

As one unique feature that Contrastive Divergence and Score Matching do not have, NCE provides the normalizing constant of an Energy-Based Model as a by-product of its training procedure. When the EBM is very expressive, e.g., a deep neural network with many parameters, we can assume it is able to approximate a normalized probability density and absorb $Z_{\bm{\theta}}$ into the parameters of $E_{{\bm{\theta}}}({\mathbf{x}})$ (Mnih and Teh, 2012), or equivalently, fixing $Z_{\bm{\theta}}=1$. The resulting EBM trained with NCE will be self-normalized, i.e., having a normalizing constant close to 1.

In this case, the NCE objective (Eq. 21) reduces to

At ${\bm{\theta}}^{*}$, we have a solution where

i.e., the optimal model matches the data distribution.

As noted in Gutmann and Hirayama (2012) and Song et al. (2019), when $\left\lVert{\mathbf{v}}\right\rVert_{2}\approx 0$, the NCE objective Eq. 26 has the following equivalent form by Taylor expansion

Comparing against Eq. 15, we immediately see that the above objective equals that of SSM, if we ignore small additional terms hidden in $o(\left\lVert{\mathbf{v}}\right\rVert_{2}^{2})$ and take the expectation with respect to ${\mathbf{v}}$ over a user-specified distribution $p({\mathbf{v}})$.

Other Methods

Aside from MCMC-based training, Score Matching and Noise Contrastive Estimation, there are also other methods for learning EBMs. Below we briefly survey some examples of them. Interested readers can learn more details from references therein.

The overarching strategy for learning probabilistic models from data is to minimize the KL divergence between data and model distributions. However, because the normalizing constants of EBMs are typically intractable, it is hard to directly evaluate the KL divergence when the model is an EBM (see the discussion in Section 3). One generic idea that has frequently circumvented this difficulty is to consider differences or derivatives (i.e., infinitesimal differences) of KL divergences. It turns out that the unknown partition functions of EBMs are often cancelled out after taking the difference of two closely related KL divergences, or computing the derivatives.

in minimum velocity learning and minimum probability flow respectively. These objectives typically do not require computing the normalizing constant $Z_{{\bm{\theta}}}$.

This objective does not require computing the partition function $Z_{{\bm{\theta}}}$ whenever $\Phi$ is linear.

Minimum velocity learning, minimum probability flow, and minimum KL contraction are all different generalizations to Score Matching and Noise Contrastive Estimation (Movellan, 2008; Sohl-Dickstein et al., 2011; Lyu, 2011).

2 Minimizing the Stein Discrepancy

We can train EBMs by minimizing the Stein discrepancy, defined by

There are two common methods to sidestep this difficulty. Gorham and Mackey (2015) and Liu et al. (2016) discovered that when $\mathcal{F}$ is a unit ball in a Reproducing Kernel Hilbert Space (RKHS) with a fixed kernel, the Stein discrepancy becomes kernelized Stein discrepancy, where the trace term is a constant and does not affect optimization. Otherwise, $\operatorname{trace}(\nabla_{\mathbf{x}}\mathbf{f}({\mathbf{x}}))$ can be approximated with the Skilling-Hutchinson trace estimator (Skilling, 1989; Hutchinson, 1989; Grathwohl et al., 2020c).

3 Adversarial Training

Recall from Section 3 that when training EBMs with maximum likelihood estimation (MLE), we need to sample from the EBM per training iteration. However, sampling using multiple MCMC steps is expensive and requires careful tuning of the Markov chain. One way to avoid this difficulty is to use non-MLE methods that do not need sampling, such as Score Matching and Noise Contrastive Estimation. Here we introduce another family of methods that sidestep costly MCMC sampling by learning an auxiliary model through adversarial training, which allows fast sampling.

Using the definition of EBMs, we can rewrite the maximum likelihood objective by introducing a variational distribution $q_{\bm{\phi}}({\mathbf{x}})$ parameterized by ${\bm{\phi}}$:

where $H(q_{\bm{\phi}}({\mathbf{x}}))$ denotes the entropy of $q_{\bm{\phi}}({\mathbf{x}})$. Step (i) is due to Jensen’s inequality. Eq. 31 provides an upper bound to the expected log-likelihood. For EBM training, we can first minimize the upper bound Eq. 31 with respect to $q_{\bm{\phi}}({\mathbf{x}})$ so that it is closer to the likelihood objective, and then maximize Eq. 31 with respect to $E_{{\bm{\theta}}}({\mathbf{x}})$ as a surrogate for maximizing likelihood. This amounts to using the following maximin objective

Optimizing the above objective is similar to training Generative Adversarial Networks (GANs) (Goodfellow et al., 2014) and can be achieved by adversarial training. The variational distribution $q_{\phi}({\mathbf{x}})$ should allow both fast sampling and efficient entropy evaluation to make Eq. 32 tractable. This limits the model family of $q_{\phi}({\mathbf{x}})$, and usually restricts our choice to invertible probabilistic models, such as inverse autoregressive flow (Kingma et al., 2016), and NICE/RealNVP (Dinh et al., 2014, 2016). See Dai et al. (2019b) for an example on designing $q_{\bm{\phi}}({\mathbf{x}})$ and training EBMs with Eq. 32.

Grathwohl et al. (2020a) represent $q_{\bm{\phi}}({\mathbf{x}})$ as a noisy neural sampler, where samples are obtained via $g_{\bm{\phi}}({\mathbf{z}})+\sigma{\bm{\epsilon}}$, assuming $\mathbf{z},{\bm{\epsilon}}\sim\mathcal{N}(\bm{0},{\bm{I}})$. With a noisy neural sampler, $\nabla_{\bm{\phi}}H(q_{\bm{\phi}}({\mathbf{x}}))$ becomes particularly easy to estimate, which allows gradient-based optimization for the maximin objective in Eq. 32. A related approach is proposed in Xie et al. (2018), where authors train a noisy neural sampler with samples obtained from MCMC, and initialize new MCMC chains with samples generated from the neural sampler. This cooperative sampling scheme improves the convergence of MCMC, but may still require multiple MCMC steps for sample generation. It does not directly optimize the objective in Eq. 31.

When using both adversarial training and MCMC sampling, Yu et al. (2020) observe that EBMs can be trained with an arbitrary $f$-divergence, including KL, reverse KL, total variation, Hellinger, etc. The method proposed by Yu et al. (2020) allows us to explore the trade-offs and inductive bias of different statistical divergences for more flexible EBM training.

Conclusion

We reviewed some of the modern approaches for EBM training. In particular, we focused on maximum likelihood estimation with MCMC sampling, Score Matching, and Noise Contrastive Estimation. We emphasized their mutual connections, and concluded by a short review on other EBM training approaches that do not directly fall into these three categories. The contents of this tutorial is of course limited to the authors’ knowledge and bias in the field; we did not cover many other important aspects of EBMs, including EBMs with latent variables, and various downstream applications of EBMs. Training techniques are crucial to problem solving with EBMs, and will remain an active direction for future research.

References