PaperlayerView exact revision on arXiv

Source-linked AI summary

Deep Unsupervised Learning using Nonequilibrium Thermodynamics

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

arXiv:1503.03585v8cs.LGcond-mat.dis-nnq-bio.NCstat.ML

TL;DR

Flexible probability models face a tradeoff between tractability and flexibility. This paper learns a reverse process for iterative diffusion and reports effective modeling across toy and real datasets, including challenging natural images.

  • Problem

    Probabilistic models face a tradeoff between tractability and flexibility when modeling complex data distributions.

    Source: §1 · p. 1

  • Method

    The method gradually diffuses data toward an analytically tractable distribution, then learns a reverse diffusion process to restore data structure.

    Source: §2.1 · p. 3, §4 · p. 9

  • Results

    The authors report effective modeling across toy and real datasets, including challenging natural-image datasets, with exact sampling and probability evaluation.

    Source: §4 · p. 9

  • Takeaways & Limitations

    The paper concludes that with more diffusion steps, each reversal step becomes simple to estimate, keeping flexible generative modeling tractable.

    Source: §4 · p. 9

    CIFAR-10 likelihood bounds here use data preprocessed with uniform noise to remove pixel quantization, unlike higher bounds reported in an earlier version.

    Source: §3 · p. 7

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 tradeoff between tractable probabilistic models and flexible but computationally expensive ones by defining a generative model through a diffusion chain. The framework also supports posterior computation and tractable evaluation of model probabilities.

1. Introduction

  • Tractable models are easy to evaluate and fit but cannot capture rich data structure, while flexible models can represent arbitrary structure but require an intractable normalization constant and expensive Monte Carlo computation.Source: §1 · p. 1

1.1. Diffusion probabilistic models

  • Diffusion probabilistic models offer extreme flexibility, exact sampling, multiplication with other distributions for posterior computation, and cheap evaluation of log likelihoods and state probabilities.Source: §1.1 · pp. 1–2
  • The model defines a simple-to-data generative Markov chain as its probability model; analytically evaluable transition probabilities make the full chain analytically evaluable.Learning estimates small perturbations rather than a full non-analytically-normalizable potential, and the approach can capture any smooth target distribution.Source: §1.1 · p. 2

1.2. Relationship to other work

  • The framework draws on physics, quasi-static processes, and annealed importance sampling rather than variational Bayesian methods.Source: §1.2 · p. 2
  • Relative to variational inference techniques, it supports posterior multiplication, matches forward and reverse functional forms, trains thousands of time steps, and bounds entropy production per step.Source: §1.2 · p. 2

2. Algorithm

The algorithm gradually diffuses data toward a tractable distribution, then learns a reverse process to generate data. It uses these trajectories to train the model, evaluate probabilities, and form posterior distributions.

2.1. Forward Trajectory

  • Repeated Markov diffusion transforms the data distribution into a tractable target, using either Gaussian or binomial diffusion in the experiments.The forward trajectory begins at the data distribution and applies diffusion for T steps.Source: §2.1 · pp. 3–4

2.2. Reverse Trajectory

  • The learned reverse trajectory restores data, estimating Gaussian means and covariances or binomial bit-flip probabilities with functions whose cost scales with the number of steps.The experiments use multilayer perceptrons for these functions, though other regression or function-fitting methods are applicable.Source: §2.2 · p. 5

2.3. Model Probability

  • Averaging relative forward and reverse trajectory probabilities over forward samples makes model probability evaluation tractable; identical trajectories allow exact evaluation from one sample.The identical-trajectory case corresponds to a quasi-static process.Source: §2.3 · p. 5

2.4. Training

  • Training maximizes a log-likelihood lower bound with analytically computable entropies and KL divergences, reducing distribution estimation to regression on transition parameters.When forward and reverse trajectories are identical, the lower bound becomes an equality.Source: §2.4 · pp. 5–6

2.4.1. SETTING THE DIFFUSION RATE βt

  • The forward diffusion schedule is learned for Gaussian diffusion but chosen to erase a constant fraction of signal per step for binomial diffusion.For Gaussian diffusion, β1 is fixed small and β2···T are learned by gradient ascent on K; for binomial diffusion, βt = (T −t + 1)−1.Source: §2.4.1 · p. 6

2.5.1. MODIFIED MARGINAL DISTRIBUTIONS

  • Multiplying the model by a second distribution is straightforward under diffusion: it can modify intermediate distributions, enabling posterior computations such as denoising and inpainting.The second distribution can be treated as a perturbation at each step or multiplied into steps directly; the paper demonstrates denoising and inpainting of natural images.Source: §2.5 · p. 6, §2.5.1 · p. 6

2.5.2. MODIFIED DIFFUSION STEPS

  • Modified reverse transitions are chosen to satisfy equilibrium for the perturbed distributions, with smooth perturbations changing Gaussian means or binomial flip rates.When multiplication is available in closed form, it can be applied directly to a reverse transition; delta-function constraints on coordinates support inpainting.Source: §2.5.2 · p. 6, §2.5.3 · p. 7

2.6. Entropy of Reverse Process

  • Because the forward process is known, conditional-entropy bounds for each reverse step and bounds on log likelihood can be computed analytically from the forward distribution.The derivation is provided in Appendix A.Source: §2.6 · p. 7

3. Experiments

The experiments evaluate diffusion probabilistic models on synthetic, binary, and image data, assessing learned distributions, log-likelihood comparisons, sampling, and inpainting. They report strong results on the tested distributions and show posterior-based inpainting of missing image regions.

3. Experiments

  • The experiments trained models on continuous datasets and a binary dataset, computed objectives and gradients with Theano, and used SFO except for CIFAR-10.They demonstrated sampling and inpainting as well as performance comparisons; the earlier, higher CIFAR-10 bounds reflected learning 8-bit pixel quantization.Source: §3 · p. 7

3.2.1. DATASETS

  • For bark textures, the model generated a posterior sample that inpainted a missing 100×100-pixel region with long-range spatial structure.The posterior was conditioned on the rest of the image; the example includes a crack continuing into the inpainted area.Source: §3 · p. 8, Figure 5 · p. 8, §3.2.1 · p. 9

3.1.1. SWISS ROLL

  • The diffusion probabilistic model successfully learned the two-dimensional swiss roll distribution.A radial basis function network generated the model’s mean function.Source: §3.1.1 · p. 8
  • Diffusion scored 317 ± 2.7 bits in the reported log-likelihood comparison, versus 325 ± 2.9 bits for Adversarial net and 349 ± 3.3 bits for Perfect model.The comparison also lists Stacked CAE at 174 ± 2.3 bits, DBN at 199 ± 2.9 bits, and Deep GSN at 309 ± 1.6 bits.Source: §3.1.1 · p. 8, Table 2 · p. 8

3.1.2. BINARY HEARTBEAT DISTRIBUTION

  • Learning the binary heartbeat distribution was nearly perfect, with the true distribution’s log likelihood equal to −2.322 bits per sequence.The sequences had length 20, with a 1 every fifth time bin and zeros elsewhere.Source: §3.1.2 · p. 8

3.2.1. DATASETS

  • MNIST log-likelihood comparisons used Parzen-window estimates to match prior work, although the training algorithm provides an asymptotically consistent lower bound.The comparisons reference Bengio et al. (2012), Bengio & Thibodeau-Laufer (2013), and Goodfellow et al. (2014).Source: Table 2 · p. 8, §3.2.1 · p. 9
  • A probabilistic model was fit to CIFAR-10 training images, and samples from the trained model were provided.The reported CIFAR-10 log-likelihood bounds use uniform-noise preprocessing to remove pixel quantization.Source: §3 · p. 7, §3.2.1 · p. 9
  • The model achieved state-of-the-art performance on dead leaves images, which have analytically tractable structure and capture statistical complexities of natural images.The comparison used identical training and test data to Theis et al. (2012).Source: Table 2 · p. 8, §3.2.1 · p. 9

4. Conclusion · Appendix

The paper introduces a diffusion-based algorithm that models diverse data distributions while enabling tractable training, exact sampling, and probability evaluation. Its reverse process also supports straightforward manipulation of conditional and posterior distributions.

4. Conclusion

  • The algorithm models probability distributions with exact sampling and probability evaluation, and is effective on toy and real datasets, including challenging natural images.The authors report using a similar basic algorithm across these tests to model a wide variety of distributions.Source: §4 · p. 9
  • The method reverses a Markov diffusion chain mapping data to noise; with more steps, each reverse-step distribution becomes simple and easy to estimate.Source: §4 · p. 9
  • Existing density-estimation techniques often sacrifice modeling power for tractability and efficiency, while sampling or evaluation can be extremely expensive.Source: §4 · p. 9
  • The resulting algorithm can fit any data distribution while remaining tractable to train, exactly sample from, and evaluate.It also makes conditional and posterior distributions straightforward to manipulate.Source: §4 · p. 9

A. Conditional Entropy Bounds Derivation

The derivation bounds the conditional entropy of a reverse-trajectory step by combining upper and lower bounds, both analytically computable from the conditional forward trajectory.

A. Conditional Entropy Bounds Derivation

  • The upper bound uses π(y) as the maximum-entropy distribution; for Gaussian data, the stated equalities require variance-1 training data scaled to unit norm, but not whitening.The maximum-entropy observation holds without qualification for the binomial distribution.Source: A. Conditional Entropy Bounds Derivation · p. 12
  • The lower bound follows because adding steps to a Markov chain does not increase information about its initial state, so it does not decrease that state's conditional entropy.Source: A. Conditional Entropy Bounds Derivation · p. 12
  • The combined expressions bound the conditional entropy for a single step using bounds that depend only on the conditional forward trajectory q and can be computed analytically.Source: A. Conditional Entropy Bounds Derivation · p. 12

B. Log Likelihood Lower Bound

This section introduces a lower bound on the log likelihood.

B. Log Likelihood Lower Bound

B.1. Entropy of p X(T ) … C. Perturbed Gaussian Transition

The derivation rewrites diffusion-model terms as entropies and KL divergences, then approximates a perturbed Gaussian transition through an energy-function expansion. It also removes an edge effect and expresses terms using the forward process’s posterior and marginals.

B.1. Entropy of p X(T )

  • The contribution involving p_{X(T)} is rewritten as an entropy that remains constant under the diffusion kernels and equals the entropy of p.Source: §B.1 · p. 13

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, allowing the first time-step contribution to be removed from the sum.Source: §B.2 · p. 13

B.3. Rewrite in terms of posterior q x(t−1)|x(0)

  • Using the forward trajectory’s Markov property, the derivation rewrites terms via a posterior and forward marginals, with a resulting term constant across t.Source: §B.3 · p. 13

B.4. Rewrite in terms of KL divergences and entropies

  • Several terms in the derivation are recognized as conditional entropies, and a log ratio of distributions is transformed into a KL divergence.Source: §B.4 · p. 14
  • The conditional entropies can be computed analytically, and the KL divergence can be computed analytically given x(0) and x(t).Source: §B.4 · p. 14

C. Perturbed Gaussian Transition

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

A 40-step Gaussian diffusion model was trained to generate a two-dimensional swiss roll distribution, which the paper reports was successfully learned.

D.1.1. SWISS ROLL

  • The target distribution was a two-dimensional swiss roll.A probabilistic model was built for this distribution.Source: §D.1.1 · p. 17
  • The generative model used 40 Gaussian diffusion time steps, initialized from an identity-covariance Gaussian.Source: §D.1.1 · p. 17
  • A normalized radial basis function network with one hidden layer and 16 hidden units generated the reverse-trajectory mean and covariance functions.Source: §D.1.1 · p. 17
  • The readout layer was learned independently at each time step, while other layers shared weights across time steps and both functions.Source: §D.1.1 · p. 17
  • The model successfully learned the swiss roll distribution.The passage points to Figure 1 as showing the result.Source: §D.1.1 · p. 17
  • A quantity was passed through a sigmoid to restrict it between 0 and 1; the swiss roll distribution was successfully learned, as shown in Figure 1.Source: §D.1.1 · p. 17

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

A 2000-step binomial diffusion model learned the length-20 binary heartbeat distribution nearly perfectly, whose true-generating-process log likelihood is −2.322 bits per sequence.

D.1.2. BINARY HEARTBEAT DISTRIBUTION

  • Training used length-20 binary sequences with a 1 every fifth bin and 2000 binomial diffusion steps initialized at p = 0.2.The initial independent binomial distribution had the same mean activity as the data.Source: §D.1.2 · p. 17
  • A sigmoid multilayer perceptron with 20 inputs and three 50-unit hidden layers generated reverse-trajectory Bernoulli rates.The readout layer was learned separately at each time step, while weights in the other layers were shared across time steps.Source: §D.1.2 · p. 17
  • The heartbeat distribution was learned nearly perfectly, with a true-generating-process log likelihood of −2.322 bits per sequence.The passage points to Figure 2 and Table 1 for the learning result.Source: §D.1.2 · p. 17

D.2.1. ARCHITECTURE

The architecture produces full-resolution, per-pixel mean and covariance predictions using convolutional features and time-dependent readout functions. Multi-scale convolution captures long-range dependencies without downsampling the final feature map.

Readout; Temporal Dependence

  • A convolutional network produces 2J outputs per pixel, split into equal-sized mean and variance coefficient subsets for time-dependent readout.The mean coefficients weight time-dependent bump functions; the variance readout is formed analogously.Source: §D.2.1 · p. 17

Mean and Variance

  • The per-pixel diffusion mean and variance predictions are parameterized as perturbations around the forward diffusion kernel, and the covariance matrix is diagonal.The equilibrium-distribution mean is used in the parameterization.Source: §D.2.1 · p. 17

Multi-Scale Convolution

  • Multi-scale convolution downsamples images to powers-of-two scales, convolves at each scale, upsamples and sums the results, then applies a soft-ReLU transformation.The design aims to capture long-range and multi-scale dependencies while retaining a full-resolution output for every pixel.Source: §D.2.1 · p. 17

Network architecture

  • For Section 3.2 experiments, the image passes through multi-scale convolution and 1 × 1 convolutional layers before a linear layer generates per-pixel mean and covariance readout coefficients.CIFAR-10 also uses a parallel dense pathway, whereas MNIST uses the dense pathway instead of multi-scale convolution.Source: §D.2.1 · p. 18

Dense Layers

  • Dense layers and width-1 convolutional layers both apply a linear transformation followed by a tanh nonlinearity.Source: §D.2.1 · p. 18

Related papers

Loading 1503.03585v8…