Deep Unsupervised Learning using Nonequilibrium Thermodynamics

Jascha Sohl-Dickstein; Eric A. Weiss; Niru Maheswaranathan; Surya Ganguli

Source-linked AI summary

Deep Unsupervised Learning using Nonequilibrium Thermodynamics

Jascha Sohl-Dickstein, Eric A. Weiss, Niru Maheswaranathan, Surya Ganguli

arXiv:1503.03585v8Mar 12, 2015cs.LGcond-mat.dis-nnq-bio.NCstat.ML

TL;DR

Flexible probabilistic models are often difficult to evaluate, train, and sample from tractably. This paper learns a reverse diffusion process that restores structure from noise, yielding generative models that can be trained, sampled, and evaluated efficiently across varied datasets.

Problem
Probabilistic modeling seeks distributions that are both flexible and tractable, but flexible models often make evaluation, training, and sampling computationally expensive.
Method
The method gradually diffuses data into an analytically tractable noise distribution, then learns the reverse Markov diffusion process to restore the data distribution.
Results
The algorithm accurately models varied toy and real datasets while enabling exact sampling and probability evaluation, including on challenging natural image datasets.
Takeaways & Limitations
The learned model remains tractable to train, exactly sample from, evaluate, and manipulate for conditional and posterior distributions.
Takeaways & Limitations
Reported CIFAR-10 log-likelihood bounds use uniformly dequantized data rather than the original 8-bit pixel values.

Abstract

from arXiv · show

A central problem in machine learning involves modeling complex data-sets using highly flexible families of probability distributions in which learning, sampling, inference, and evaluation are still analytically or computationally tractable. Here, we develop an approach that simultaneously achieves both flexibility and tractability. The essential idea, inspired by non-equilibrium statistical physics, is to systematically and slowly destroy structure in a data distribution through an iterative forward diffusion process. We then learn a reverse diffusion process that restores structure in data, yielding a highly flexible and tractable generative model of the data. This approach allows us to rapidly learn, sample from, and evaluate probabilities in deep generative models with thousands of layers or time steps, as well as to compute conditional and posterior probabilities under the learned model. We additionally release an open source reference implementation of the algorithm.

1. Introduction

The paper addresses the tractability–flexibility tradeoff by defining generative models through a diffusion process that gradually transforms a simple distribution into the data distribution. The framework supports flexible structure, exact sampling, distribution multiplication, and cheap likelihood and state-probability evaluation.

Motivation: Tractable models are easy to evaluate and fit but cannot describe rich datasets, whereas flexible models face generally intractable normalization and expensive computation.Flexible models typically require costly Monte Carlo procedures for evaluation, training, or sampling.
Contributions: The proposed probabilistic models combine extreme flexibility in structure, exact sampling, easy multiplication with other distributions, and cheap evaluation of log likelihoods and individual-state probabilities.Multiplication with another distribution can support posterior computation.
Method: A generative Markov chain uses diffusion to convert a simple known distribution, such as a Gaussian, into a target data distribution.The model is explicitly defined through the Markov chain rather than using it only to approximate an otherwise-defined model.
Method: Learning estimates small perturbations to the diffusion process, which is more tractable than specifying a single non-analytically-normalizable potential and can capture any smooth target distribution.The framework therefore targets data distributions of arbitrary form when the target is smooth.
Novelty relative to prior work: Unlike variational Bayesian approaches, the framework is developed from physics, quasi-static processes, and annealed importance sampling, while simplifying the forward inference process and enabling posterior computation.The forward process is restricted to a simple functional form to address asymmetry-related training difficulty in variational inference.

2. Algorithm

The algorithm gradually diffuses a complex data distribution into a tractable one, then learns a reverse diffusion process for generation, likelihood evaluation, and posterior computation. It reduces distribution estimation to regression over transition parameters and supports perturbation-based denoising and inpainting.

Forward diffusion: The forward process repeatedly applies Gaussian or binomial Markov diffusion kernels to convert the data distribution into a tractable target distribution.Gaussian diffusion approaches an identity-covariance Gaussian, while binomial diffusion approaches an independent binomial distribution.
Reverse diffusion: The reverse generative process uses learned Gaussian means and covariances or binomial bit-flip probabilities at each diffusion step.Multi-layer perceptrons define these functions in the reported experiments, although other regression or function-fitting methods are applicable.
Probability evaluation: A quasi-static process makes forward and reverse trajectory distributions identical, allowing exact probability evaluation from a single forward-trajectory sample.This follows from evaluating relative forward and reverse trajectory probabilities averaged over forward trajectories.
Training and likelihood: Estimating the generative distribution is reduced to regression on transition functions that parameterize sequences of Gaussians or Bernoulli trials.Training finds reverse Markov transitions that maximize a lower bound on log likelihood, with analytically computable entropies and KL divergences.
Posterior computation: Diffusion models straightforwardly multiply a learned distribution by a second distribution or positive function, enabling denoising and inpainting.The second distribution can be treated as a perturbation at each diffusion step or multiplied exactly into many steps.

3. Experiments

Experiments trained diffusion probabilistic models on continuous and binary datasets, demonstrating successful learning, sampling, likelihood evaluation, and posterior inpainting. The models learned simple synthetic distributions nearly perfectly, achieved strong image-modeling results, and enabled posterior-based reconstruction of missing bark texture.

Image datasets: CIFAR-10 likelihood bounds were recomputed after adding uniform noise to remove 8-bit pixel quantization, correcting earlier inflated results.The revised preprocessing followed the recommendation cited in the paper.
Synthetic distributions: The model successfully learned a two-dimensional swiss roll distribution.A radial basis function network generated the reverse-trajectory mean.
Synthetic distributions: −2.322 bits per sequence is the true-distribution log likelihood for the 20-bit binary sequence task, where learning was nearly perfect.The sequence contains a 1 every fifth time bin and zeros elsewhere.
Image datasets: On MNIST, diffusion achieved 317 ± 2.7 bits, compared with 309 ± 1.6 bits for Deep GSN and 325 ± 2.9 bits for an adversarial net.The comparison used Parzen-window estimates in bits and included several previous methods.
Image datasets: The model achieved state-of-the-art performance on dead leaf images, which test natural-image statistical complexity despite analytically tractable structure.Evaluation used identical training and test data as prior work.
Image datasets: Posterior sampling enabled straightforward inpainting of a large missing bark-texture region while recovering long-range spatial structure.The missing region was reconstructed by sampling from the posterior conditioned on the observed image.

4. Conclusion · Appendix

The paper introduces a diffusion-reversal algorithm that models probability distributions while enabling exact sampling and probability evaluation. Its tractability also supports fitting broad data distributions and manipulating conditional and posterior distributions.

4. Conclusion: The algorithm models probability distributions with exact sampling and probability evaluation.It was demonstrated on toy and real datasets, including challenging natural image datasets.
4. Conclusion: The method was effective across toy and real datasets, including challenging natural image datasets.The tests used a similar basic algorithm across datasets.
4. Conclusion: The core algorithm estimates the reversal of a Markov diffusion chain that maps data to a noise distribution.As the number of diffusion steps increases, each reversal distribution becomes simpler to estimate.
4. Conclusion: With many diffusion steps, each reversal distribution becomes simple and easy to estimate.This property underlies the algorithm’s tractable learning and generation procedures.
4. Conclusion: The resulting algorithm can learn a fit to any data distribution while remaining tractable to train, sample from exactly, and evaluate.The conclusion presents these capabilities as properties of the diffusion-reversal construction.
4. Conclusion: The learned model makes conditional and posterior distributions straightforward to manipulate.This capability is stated alongside tractable training, exact sampling, and probability evaluation.

A. Conditional Entropy Bounds Derivation

The section derives upper and lower bounds on the conditional entropy of a single reverse-trajectory step. The bounds rely on maximum-entropy observations and Markov-chain information monotonicity, and depend only on the analytically computable conditional forward trajectory q.

Upper bound: The upper entropy-change bound uses π(y) as the maximum-entropy distribution, with Gaussian equalities requiring unit-variance training data scaled to unit norm.Whitening is not required for the Gaussian case.
Lower bound: The lower entropy-difference bound follows because additional Markov-chain steps cannot increase information about the initial state or decrease its conditional entropy.This establishes the lower bound through information monotonicity along the chain.
Combined bound: Combining the bounds yields a single-step conditional-entropy bound whose upper and lower limits depend only on the conditional forward trajectory q and can be computed analytically.The derivation applies these expressions to the conditional entropy of a reverse-trajectory step.

B. Log Likelihood Lower Bound … B.3. Rewrite in terms of posterior q x(t−1)|x(0)

The appendix derives a lower bound on log likelihood, simplifies entropy and edge contributions, and rewrites the remaining terms using forward-process posteriors and marginals. It also notes that the final reverse step matches the corresponding forward step and includes true diffusion-model samples on MNIST.

B. Log Likelihood Lower Bound: The appendix begins by stating a lower bound on the log likelihood.
B.1. Entropy of p X(T ): The contribution from p X(T) is separated and rewritten as an entropy.
B.1. Entropy of p X(T ): The entropy associated with the forward diffusion kernels remains constant under the diffusion process.
B.2. Remove the edge effect at t = 0: To avoid edge effects, the final reverse-trajectory step is set equal to its corresponding forward diffusion step.
B.2. Remove the edge effect at t = 0: This equivalence removes the contribution of the first time-step from the sum.
B.3. Rewrite in terms of posterior q x(t−1)|x(0): Bayes’ rule expresses the rewritten term through a posterior and marginals from the forward trajectory.
B.3. Rewrite in terms of posterior q x(t−1)|x(0): The appendix figure shows true samples from a diffusion probabilistic model trained on MNIST digits, rather than Gaussian or binomial distribution means.

B.4. Rewrite in terms of KL divergences and entropies · C. Perturbed Gaussian Transition

The derivation rewrites terms using conditional entropies and KL divergences, with these quantities analytically computable under stated conditions. For the perturbed Gaussian transition, an energy-based Taylor expansion yields a Gaussian form when the energy is sufficiently smooth relative to the covariance precision.

B.4. Rewrite in terms of KL divergences and entropies: Several terms in the derivation are recognized as conditional entropies.
B.4. Rewrite in terms of KL divergences and entropies: The log ratio of probability distributions is transformed into a KL divergence.
B.4. Rewrite in terms of KL divergences and entropies: The entropies are analytically computable, and the KL divergence is analytically computable given x(0) and x(t).
C. Perturbed Gaussian Transition: The perturbed reverse diffusion kernel is computed by rewriting the transition in terms of energy functions.
C. Perturbed Gaussian Transition: A Taylor expansion around µ is justified when the Hessian eigenvalues of Er(y) are everywhere much smaller in magnitude than those of Σ−1.
C. Perturbed Gaussian Transition: Under this approximation, the full energy becomes a quadratic expression involving y−µ+Σg and a constant.
C. Perturbed Gaussian Transition: The resulting expression corresponds to a Gaussian and is then substituted back into the original formalism.

D. Experimental Details · D.1. Toy Problems · D.1.1. SWISS ROLL

The study modeled a two-dimensional swiss roll distribution with a 40-step Gaussian diffusion process and a radial basis function reverse model. The learned model successfully reproduced the swiss roll distribution.

D.1.1. SWISS ROLL: A probabilistic model was built for a two-dimensional swiss roll distribution.
D.1.1. SWISS ROLL: 40 time steps of Gaussian diffusion were initialized at an identity-covariance Gaussian distribution.
D.1.1. SWISS ROLL: A normalized radial basis function network with one hidden layer and 16 hidden units generated the reverse trajectory’s mean and covariance functions.
D.1.1. SWISS ROLL: The readout layer for each function was learned independently at every time step.
D.1.1. SWISS ROLL: Weights in all other layers were shared across time steps and between both functions.
D.1.1. SWISS ROLL: The top-layer output of fΣ was passed through a sigmoid to restrict it between 0 and 1.
D.1.1. SWISS ROLL: The swiss roll distribution was successfully learned, as shown in Figure 1.

D.1.2. BINARY HEARTBEAT DISTRIBUTION · D.2. Images

The model learned a periodic binary heartbeat distribution from 20-bin sequences using 2000-step binomial diffusion and a shared-weight multilayer perceptron. Learning was nearly perfect, with the learned distribution successfully reproducing the heartbeat pattern and the true process having log2 = −2.322 bits per sequence.

D.1.2. BINARY HEARTBEAT DISTRIBUTION: The data comprised length-20 binary sequences with a 1 every fifth time bin and 0 elsewhere.
D.1.2. BINARY HEARTBEAT DISTRIBUTION: 2000 time steps of binomial diffusion initialized from an independent binomial distribution with matching mean activity were used.The initialization probability was p = 0.2.
D.1.2. BINARY HEARTBEAT DISTRIBUTION: A multilayer perceptron with sigmoid nonlinearities used 20 input units and three hidden layers of 50 units each.
D.1.2. BINARY HEARTBEAT DISTRIBUTION: The network generated Bernoulli rates fb of the reverse trajectory, with the readout layer learned separately at each time step.
D.1.2. BINARY HEARTBEAT DISTRIBUTION: Weights in all non-readout layers were shared across time steps, while the top-layer sigmoid constrained outputs to [0, 1].
D.1.2. BINARY HEARTBEAT DISTRIBUTION: The heartbeat distribution was successfully learned, and learning was nearly perfect according to Figure 2 and Table 1.The true generating process has log2 = −2.322 bits per sequence.

D.2.1. ARCHITECTURE

The architecture uses a convolutional network with per-pixel outputs split into mean and variance coefficients, then applies time-dependent readouts to produce diffusion predictions. Multi-scale convolution captures long-range and multi-scale dependencies while preserving full-resolution outputs.

A convolutional network produces y_i ∈ R^2J for each image pixel, splitting its entries into equal-sized y_µ and y_Σ subsets.
Time-dependent bump functions weight the convolution outputs over time to generate z_µ and z_Σ.The z_Σ construction is identical to z_µ but uses y_Σ.
Image experiments use T = 1000 timesteps, except the bark dataset, which uses T = 500.
Multi-scale convolution downsamples images at powers-of-two scales, convolves each scale, upsamples and sums them, then applies a soft-ReLU transformation.This design targets long-range and multi-scale dependencies while producing full-resolution feature maps.