Learning Latent Space Energy-Based Prior Model

Bo Pang; Tian Han; Erik Nijkamp; Song-Chun Zhu; Ying Nian Wu

Source-linked AI summary

Learning Latent Space Energy-Based Prior Model

Bo Pang, Tian Han, Erik Nijkamp, Song-Chun Zhu, Ying Nian Wu

arXiv:2006.08205v2Jun 15, 2020stat.MLcs.LG

TL;DR

The paper addresses the limited informativeness of fixed latent priors and the difficulty of MCMC-based EBM learning. It learns an EBM prior in a generator’s latent space jointly with the top-down network using short-run MCMC, and reports strong results in image and text generation and anomaly detection.

Problem
Generator models commonly use non-informative latent priors, while EBM learning is limited by the difficulty of MCMC sampling.
Method
The method learns a latent-space EBM prior and the generator’s top-down network jointly by maximum likelihood with short-run MCMC sampling of prior and posterior latent variables.
Results
The model shows strong performance in image and text generation and anomaly detection, and is competitive with prior art.
Takeaways & Limitations
Moving the EBM into the low-dimensional latent space makes MCMC-based learning feasible while preserving effective data-regularity modeling.
Takeaways & Limitations
MCMC sampling makes the method approximately four times slower than VAEs on image datasets, and the model leaves the top-down likelihood un-tilted.

Abstract

from arXiv · show

We propose to learn energy-based model (EBM) in the latent space of a generator model, so that the EBM serves as a prior model that stands on the top-down network of the generator model. Both the latent space EBM and the top-down network can be learned jointly by maximum likelihood, which involves short-run MCMC sampling from both the prior and posterior distributions of the latent vector. Due to the low dimensionality of the latent space and the expressiveness of the top-down network, a simple EBM in latent space can capture regularities in the data effectively, and MCMC sampling in latent space is efficient and mixes well. We show that the learned model exhibits strong performances in terms of image and text generation and anomaly detection. The one-page code can be found in supplementary materials.

1 Introduction

The paper augments generator models with a learned latent-space EBM prior, jointly trained with the top-down network using short-run MCMC. This approach targets effective latent regularity modeling and efficient sampling, with competitive results across image and text generation and anomaly detection.

Motivation: The latent space energy-based prior model adds an informative learned prior atop the generator’s expressive top-down network.It replaces a non-informative latent prior with an EBM learned from observed data.
Learning method: The EBM prior and top-down network are jointly learned by maximum likelihood using MCMC samples from prior and posterior latent distributions.Prior parameters use statistical differences between the two sample sets, while the top-down network uses posterior samples and observed data.
Computational design: Low-dimensional latent variables allow a small MLP energy function to capture data regularities while short-run MCMC remains efficient and mixes well.The method runs a fixed number of MCMC steps from a fixed initial distribution.
Empirical scope: The method is evaluated for image synthesis, text generation, and anomaly detection, where it is reported as competitive with prior art.The authors also provide a theoretical foundation for short-run-MCMC learning and report strong empirical results.

2 Model and learning

The model replaces a fixed latent prior with an energy-based prior atop a top-down generator, then jointly learns both through maximum likelihood using prior and posterior MCMC samples. Short-run MCMC provides a tractable approximation whose updates have a theoretical connection to perturbed likelihood and estimating-equation learning.

2.1 Model: The latent prior is an energy-based correction of a known reference distribution, while the top-down generator maps latent variables to observed data.The reference prior is isotropic Gaussian in this paper, and the energy function is parameterized by a small multilayer perceptron.
2.2 Maximum likelihood: Joint maximum-likelihood learning updates the prior using posterior-versus-prior latent statistics and updates the generator from posterior samples and observed data.The prior gradient compares inferred latent variables under the posterior with samples from the current prior, whereas the generator gradient uses posterior expectations of the conditional model.
2.3 Short-run MCMC: Short-run Langevin MCMC starts from the fixed reference prior and runs a fixed number of steps to approximate both prior and posterior sampling.The resulting distributions are well defined for fixed initialization, step size, and step count; their divergence from the target decreases as the number of steps increases.
2.4 Learning algorithm: The algorithm alternates prior and posterior short-run sampling, then uses Monte Carlo estimates to update the prior and top-down generator parameters.Prior and posterior sampling serve as the negative and positive phases of latent-space EBM learning.
2.5 Theoretical understanding: Short-run-MCMC learning is a perturbation of maximum-likelihood learning, and its stochastic updates solve an estimating equation with Robbins–Monro convergence under regularity conditions.Replacing short-run distributions by their exact target distributions recovers the maximum-likelihood estimating equation.
2.6 Amortized inference and synthesis: The paper also describes amortized inference and synthesis as joint optimization with extra networks, but the initial model intentionally omits those learned computation networks.The later formulation combines variational inference and adversarial training, whereas the initial method remains deliberately simple.

3 Experiments

Experiments evaluate the latent-space EBM on image and text generation, latent-chain behavior, anomaly detection, and computational cost. The model produces strong generation and reconstruction results, improves text fluency and diversity, supports meaningful latent-space transitions, and remains more costly than amortized VAEs on image data.

3.1 Image modeling: The model achieves superior image generation performance and accurate reconstructions across SVHN, CIFAR-10, and CelebA compared with listed baselines.Generation is measured by FID and reconstruction by MSE in Table 1.
3.2 Text modeling: The model outperforms SA-VAE, FB-VAE, and ARAE on FPPL and RPPL across SNLI, PTB, and Yahoo, indicating fluent and diverse generated text.FPPL measures fluency, while RPPL measures distributional coverage.
3.3 Analysis of latent space: Langevin chains initialized from the reference prior improve image synthesis quality and textual semantics and syntax as the number of steps increases.The textual improvement is especially clear from step 0 to 40.
3.4 Anomaly detection: For anomaly detection, the method uses the unnormalized log-posterior as a decision function after training on normal data and testing on normal and anomalous examples.Performance is evaluated with AUPRC on MNIST against BiGAN, MEG, and VAE.
3 Experiments: MCMC-based inference is approximately four times slower than VAEs on image datasets, although total training time is not disadvantageous on text datasets.The paper also reports faithful samples from a larger 128 × 128 CelebA generator.

4 Discussion and conclusion

The paper situates its latent-space EBM as an integration of top-down generation and energy-based modeling, using short-run MCMC for feasible prior and posterior sampling. Moving the EBM into latent space and placing it over an expressive generator is presented as a way to capture data regularities while reducing sampling difficulty.

The model combines a directed top-down generator with an energy-based latent prior that corrects simple independence assumptions and captures additional regularities.
Latent-space EBM learning is more feasible than data-space EBM learning because data-space distributions can be highly multimodal and difficult to sample.
Short-run MCMC is used to sample both prior and posterior latent variables, extending earlier uses of short-run sampling in EBMs and top-down models.
The approach addresses limitations of fixed simple priors, which may cause ineffective generator learning, by learning a flexible prior jointly with the generator.
The authors conclude that latent-space EBMs can make MCMC-based learning more sound and feasible while capturing regularities in data effectively.

Broader Impact

The paper identifies potential relevance for researchers studying generator models, EBMs, MCMC sampling, unsupervised learning, image synthesis, and text generation.

The work may interest researchers in generator models, energy-based models, MCMC sampling, unsupervised learning, image synthesis, and text generation.

A Theoretical derivations

This appendix begins the theoretical derivations from a basic score-function identity and connects it to policy-gradient notation. The identity follows directly by differentiating a normalized expectation.

The derivations begin from the identity Eθ[∇θ log pθ(x)] = 0 for expectations under the model distribution.
The identity is a special case of the generalized relation ∇θEθ[R(x)] = Eθ[R(x)∇θ log pθ(x)] when R(x) = 1.

A.2 Maximum likelihood estimating equation

The appendix derives maximum-likelihood learning for the latent-variable generator and relates short-run MCMC learning to a perturbation of the exact MLE estimating equation. It also connects the derivation to EM and energy-based-model learning.

The resulting stochastic-gradient procedure is a perturbation of the MLE estimating equation, while replacing short-run distributions with exact prior and posterior distributions recovers MLE.
For pθ(x,z) = pα(z)pβ(x|z), the observed-data gradient equals the posterior expectation of the joint log-density gradient.
The latent-variable identity follows because the expected posterior score is zero, leaving the observed-data gradient unchanged.
The derivation connects this gradient to EM, where x is observed data, z is missing data, and the joint log-density is the complete-data log-likelihood.
For the EBM prior, the normalization-gradient identity relates the energy-gradient expectation to the derivative of log Z(α).
The short-run MCMC method is formulated as a computationally tractable perturbation of log likelihood whose added KL-divergence terms have zero derivative at the current parameters.

A.8 Learning with short-run MCMC as perturbation of log-likelihood

Short-run MCMC learning is framed as a perturbation of maximum likelihood, with two KL-divergence terms capturing bias from approximate prior and posterior sampling.

Short-run MCMC replaces exact prior and posterior distributions with approximate distributions, yielding a perturbed log-likelihood objective.The learning rule is derived as a perturbation involving approximate prior and posterior sampling distributions.
The perturbation contains two KL terms that measure divergence between short-run and target posterior and prior distributions.These are DKL(˜pθt(z|x)∥pθ(z|x)) and DKL(˜pαt(z)∥pα(z)).
The posterior KL term relates the procedure to VAE-style inference, while the prior KL term relates it to contrastive divergence initialized from Gaussian noise.The method uses no separate inference network in the posterior comparison and initializes prior sampling from p0(z).
The KL terms bias learning relative to exact MLE because exact MCMC sampling is impractical.The paper notes that this bias may nevertheless bring short-run posterior and prior distributions closer to their useful counterparts.
The prior update has an adversarial interpretation: the EBM compares short-run prior samples against posterior samples inferred from real data.The prior is updated toward the data-associated posterior samples and away from its own short-run samples.
The authors also describe amortizing both short-run samplers with an inference network qφ(z|x) and a synthesis network qψ(z).The inference network follows VAE-style learning, while the synthesis network variationally approximates the EBM prior.

B.1 Experiment details

Experiments cover image and text datasets using specified EBM, generator, and short-run dynamics configurations, with optimization continued until convergence.

Experiments use SVHN, CIFAR-10, and CelebA for images, plus PTB, Yahoo, and SNLI for text generation.Image training uses resized data scaled to [−1, 1], with the stated training-set sizes for SVHN, CIFAR-10, and CelebA.
The EBM uses dataset-specific architectures, while text generators are one-layer unidirectional LSTMs and image generators follow the listed convolutional designs.For text, the latent dimensionality is set to 32.
Short-run dynamics specify separate prior and posterior step counts and step sizes, with the same settings across models and modalities except CIFAR-10.CIFAR-10 uses K1 = 40 posterior sampling steps.
All models use Adam optimization and are trained until convergence, requiring approximately 70,000 image updates and 40,000 text updates.Initialization differs between image and text generators, while Xavier normal initialization is used for the EBM and image generators.

C Ablation study

Ablations show that the latent EBM prior outperforms fixed Gaussian alternatives, while more MCMC steps and greater model complexity generally improve FID.

The ablation reference configuration is the architecture and hyperparameter setting whose FID is reported in the main text.The study uses SVHN and highlights the reference FID in Tables 8, 9, and 10.
Fixed prior: The EBM prior outperforms a fixed Gaussian prior even when the Gaussian model uses a generator with four times more parameters.Fixed Gaussian models improve as generator complexity increases, but remain inferior to the EBM-prior reference model.
MCMC steps: Increasing prior or posterior short-run MCMC steps improves synthesis quality as measured by FID.Table 9 separates the effects of prior steps K0 and posterior steps K1.
Prior EBM and generator complexity: Increasing prior EBM hidden features or generator channel width generally improves generation according to FID.Table 10 reports FID as a function of prior EBM feature count nef and generator channel factor ngf.

D PyTorch code

The PyTorch implementation samples prior and posterior latent states with Langevin dynamics, then updates the generator and EBM from reconstruction and energy differences.

The reference image configuration uses nz = 100, K0 = 60, K1 = 40, and 70,000 training iterations.The implementation uses Adam with separate learning rates for the EBM and generator.
The generator is a transposed-convolution network that maps latent vectors to images through Leaky-ReLU layers and a final Tanh output.The shown image generator uses four transposed-convolutional stages.
The EBM is a small multilayer perceptron mapping latent vectors to scalar energies.Its displayed architecture uses three linear layers with Leaky-ReLU activations between them.
The implementation initializes Gaussian latent samples, draws data batches, and runs separate prior and posterior Langevin chains.The posterior chain uses the generator likelihood and EBM energy gradients, while the prior chain uses the EBM gradient.
The generator update minimizes reconstruction MSE using posterior samples, while the EBM update contrasts energies on posterior and prior samples.The EBM loss is computed as the mean posterior energy minus the mean prior energy.