Source-linked AI summary
How to Train Your Energy-Based Models
Yang Song, Diederik P. Kingma
TL;DR
EBMs offer flexible probability modeling without requiring a tractable normalizing constant, but that flexibility makes training difficult. This tutorial explains MCMC maximum likelihood, Score Matching, Noise Contrastive Estimation, their connections, and alternative approaches, while noting important scope limitations.
Problem
EBMs can model expressive distributions without tractable normalizers, but the unknown normalizing constant makes training particularly difficult.
Method
The tutorial presents maximum likelihood training with MCMC, MCMC-free Score Matching and Noise Contrastive Estimation, their theoretical connections, and alternative approaches.
Results
Score Matching estimators are consistent under regularity conditions, score-based generative models perform among the best for high-resolution image generation, and NCE can produce a self-normalized EBM.
Takeaways & Limitations
EBM training techniques provide multiple routes around likelihood intractability, with MCMC, Score Matching, and NCE offering distinct training mechanisms and connections.
Takeaways & Limitations
The tutorial is limited by the authors’ knowledge and bias and does not cover latent-variable EBMs or various downstream applications.
Abstract
from arXiv · showhide
Energy-Based Models (EBMs), also known as non-normalized probabilistic models, specify probability density or mass functions up to an unknown normalizing constant. Unlike most other probabilistic models, EBMs do not place a restriction on the tractability of the normalizing constant, thus are more flexible to parameterize and can model a more expressive family of probability distributions. However, the unknown normalizing constant of EBMs makes training particularly difficult. Our goal is to provide a friendly introduction to modern approaches for EBM training. We start by explaining maximum likelihood training with Markov chain Monte Carlo (MCMC), and proceed to elaborate on MCMC-free approaches, including Score Matching (SM) and Noise Constrastive Estimation (NCE). We highlight theoretical connections among these three approaches, and end with a brief survey on alternative training methods, which are still under active research. Our tutorial is targeted at an audience with basic understanding of generative models who want to apply EBMs or start a research project in this direction.
1. Introduction
Tractable-likelihood models enable straightforward optimization but impose structural constraints that are not always natural. EBMs relax these constraints by parameterizing an unnormalized energy, making training difficult because likelihood computation and exact sampling are generally intractable.
- Tractable likelihoods simplify model comparison and log-likelihood optimization but constrain model form.Autoregressive, flow-based, and variational-autoencoder models rely on specific factorization, invertibility, or directed-latent-variable assumptions.
- EBMs specify an unnormalized negative log-probability called an energy function.Because the energy need not integrate to one, it can use flexible nonlinear regression architectures.
- EBM flexibility permits architectures developed for classification or regression and special-purpose designs suited to the data type.
- Exact likelihood computation and exact sampling are generally intractable for EBMs, making training especially difficult.
- The tutorial covers maximum likelihood with MCMC, Score Matching, Noise Contrastive Estimation, their relationships, and other training directions.
2. Energy-Based Models (EBMs)
An unconditional EBM models a single dependent variable through an energy function and a normalizing constant. Although the constant does not depend on x, its dependence on parameters makes likelihood evaluation and differentiation typically intractable.
- The section assumes unconditional EBMs over one dependent variable x, with extensions to multiple variables or conditioning variables described as straightforward.
- The energy Eθ(x) is a nonlinear regression function, while Zθ is the parameter-dependent normalizing constant or partition function.
- Zθ is constant with respect to x but depends on θ, so evaluating and differentiating log pθ(x) typically requires an intractable integral.
3. Maximum Likelihood Training with MCMC
Maximum likelihood for EBMs reduces the intractable normalizing-constant gradient to an expectation under the model, which can be estimated using MCMC samples. Practical methods such as Contrastive Divergence trade convergence for speed, introducing possible gradient bias.
- Maximum likelihood fits pθ to pdata by maximizing expected data log-likelihood, equivalently minimizing their KL divergence up to a θ-independent constant.
- EBM likelihood gradients contain an easily computed energy term and an intractable normalizing-constant term.
- A model sample ˜x ∼ pθ(x) yields an unbiased one-sample Monte Carlo estimate of the log-likelihood gradient, enabling stochastic gradient ascent.
- Langevin MCMC and Hamiltonian Monte Carlo exploit the score identity ∇x log pθ(x) = −∇xEθ(x) to sample from EBMs.
- With ϵ → 0 and K → ∞, Langevin MCMC produces samples distributed as pθ(x) under regularity conditions.
- Finite-step Langevin sampling introduces discretization error that can be negligible in practice or corrected with a Metropolis-Hastings step.
- Contrastive Divergence initializes MCMC at a datapoint and runs a fixed, typically sub-convergent number of steps to make learning practical.
- Truncated MCMC can bias Contrastive Divergence gradients, causing learned EBMs to misrepresent the data distribution and harming learning dynamics.
4. Score Matching (SM)
Score Matching trains EBMs by matching score functions, avoiding the intractable normalizing constant, but its variants trade consistency and computational efficiency against assumptions and estimation costs.
- Matching log-density derivatives identifies the data distribution because normalized densities with equal scores differ only by a constant.An EBM score is −∇xEθ(x), which is directly computable without the normalizing constant.
- Score Matching minimizes Fisher divergence, and integration by parts replaces unknown data-score derivatives with second derivatives of the energy.The resulting estimator is consistent under stated regularity conditions.
- Full second derivatives generally cost O(d^2), limiting implicit Score Matching for high-dimensional data and complex energy functions.Although only the Hessian trace is needed, it remains expensive with automatic differentiation.
- Denoising Score Matching (DSM): Adding smooth noise makes irregular data distributions suitable for Score Matching, but the resulting objective targets the noisy distribution rather than the original data distribution.Digital images are discrete and bounded, so their log density may be discontinuous; noise smooths the distribution but can introduce inconsistency.
- Denoising Score Matching (DSM): Denoising Score Matching avoids second-order derivatives through a tractable denoising objective, yet its estimator is inconsistent unless the perturbation leaves the data distribution effectively unchanged.Small noise reduces inconsistency but can increase objective variance and hinder optimization.
- Sliced Score Matching (SSM): Sliced Score Matching projects vector-valued scores onto random directions, providing a computationally efficient and consistent alternative to Denoising Score Matching.The projection distribution must have a positive-definite second-moment matrix.
- Score-Based Generative Models: Score Matching can misestimate mixture weights when component supports are nearly disjoint, causing samples to concentrate around modes in inappropriate proportions.This issue is especially relevant in high-dimensional real data.
- Score-Based Generative Models: Noise-conditional score models combine multiple noise scales and anneal from large to small noise during sampling, supporting high-resolution image, audio, and shape generation.A single network shares weights across noise scales and is trained with one Score Matching objective per scale.
5. Noise Contrastive Estimation
Noise Contrastive Estimation (NCE) trains an energy-based model by contrasting data with a chosen noise distribution through conditional classification. At the optimum, the model matches the data distribution, and NCE can connect to Score Matching while estimating the normalizing constant.
- NCE contrasts the data distribution with a noise distribution of known density and fits the resulting binary posterior by conditional maximum likelihood.The posterior distinguishes samples from data and noise, and the objective can be optimized with stochastic gradient ascent.
- At the optimum, the learned energy function matches the data distribution and its normalizing constant is recovered.A sufficiently powerful classifier matches the data/noise posterior, yielding an energy model with the corresponding normalizing constant.
- NCE provides the normalizing constant as a by-product of training, unlike Contrastive Divergence and Score Matching.For expressive models, the normalizing constant may instead be absorbed into the energy parameters, producing a self-normalized model.
- The noise distribution is critical for NCE, especially with structured or high-dimensional data, and works best when it is close to the data distribution.A perturbed-data noise distribution is efficient to sample from, although its density is difficult to evaluate directly and is replaced by a model-based parameterization.
- With a suitable noise distribution and parameterization, NCE recovers Score Matching and its objective becomes equivalent to SSM up to small Taylor-expansion terms.This connection arises when the perturbation vector has small norm, ∥v∥2 ≈0.
6. Other Methods
The section surveys EBM training methods beyond the main MCMC, Score Matching, and Noise Contrastive Estimation approaches, emphasizing objectives that avoid or reduce dependence on intractable normalizing constants and costly sampling.
- 6.1 Minimizing Differences/Derivatives of KL Divergences: Alternative objectives based on KL differences, derivatives, or transformations can cancel unknown partition functions during EBM training.Examples include minimum velocity learning, minimum probability flow, and minimum KL contraction.
- 6.2 Stein Discrepancy: Stein-based training avoids the partition function by using the EBM score, but its trace term can be expensive in high dimensions.Kernelization or the Skilling-Hutchinson estimator can address this computational cost.
- 6.3 Adversarial Training: Variational training introduces an auxiliary distribution and optimizes a maximin objective related to adversarial training.The auxiliary distribution must support fast sampling and efficient entropy evaluation.
- 6.3 Adversarial Training: Invertible probabilistic models satisfy the auxiliary distribution’s tractability requirements, while neural samplers provide efficient sampling but make entropy evaluation difficult.Noisy neural samplers make entropy gradients easier to estimate, whereas cooperative sampling can still require multiple MCMC steps.
- 6.3 Adversarial Training: Combining adversarial training with MCMC allows EBMs to be trained using a range of f-divergences, including KL, reverse KL, total variation, and Hellinger.This exposes trade-offs and inductive biases associated with different statistical divergences.
7. Conclusion
The tutorial reviews maximum likelihood with MCMC, Score Matching, and Noise Contrastive Estimation, connects these approaches, and surveys additional methods. Its coverage is limited by the authors’ knowledge and perspective, excluding latent-variable EBMs and many downstream applications.
- The tutorial focuses on maximum likelihood estimation with MCMC, Score Matching, and Noise Contrastive Estimation, while emphasizing their mutual connections.
- It also briefly reviews EBM training approaches outside those three categories.
- The coverage is limited by the authors’ knowledge and bias and excludes latent-variable EBMs and various downstream applications.
- Training techniques for EBMs remain an active direction for future research.