Variational Diffusion Models

Diederik P. Kingma; Tim Salimans; Ben Poole; Jonathan Ho

Source-linked AI summary

Variational Diffusion Models

Diederik P. Kingma, Tim Salimans, Ben Poole, Jonathan Ho

arXiv:2107.00630v6Jul 1, 2021cs.LGstat.ML

TL;DR

Diffusion models had strong perceptual synthesis but had not matched autoregressive models for likelihood-based density estimation. This paper introduces VDMs with learnable noise schedules, a simplified SNR-based VLB, and architectural improvements, achieving state-of-the-art image likelihoods and enabling near-optimal lossless compression.

Problem
Diffusion models had shown impressive perceptual generation but had yet to match autoregressive models on density estimation benchmarks, despite the importance of likelihood-based modeling for applications such as compression.
Method
VDMs jointly learn the diffusion process and denoising model, analyze the VLB through the signal-to-noise ratio, and add Fourier features for fine-scale prediction.
Results
VDMs achieve state-of-the-art likelihoods on standard image density estimation benchmarks and outperform contemporary autoregressive models, with faster optimization in a reported CIFAR-10 comparison.
Takeaways & Limitations
The continuous-time VLB’s invariance to the noise schedule enables variance-minimizing schedule learning, while the model also supports lossless bits-back compression with rates close to the theoretical optimum.
Takeaways & Limitations
The reverse model assumes a spherical Gaussian marginal at the most-noisy endpoint when the endpoint SNR is sufficiently small, and the diffusion process assumes strictly decreasing SNR with smooth coefficients.

Abstract

from arXiv · show

Diffusion-based generative models have demonstrated a capacity for perceptually impressive synthesis, but can they also be great likelihood-based models? We answer this in the affirmative, and introduce a family of diffusion-based generative models that obtain state-of-the-art likelihoods on standard image density estimation benchmarks. Unlike other diffusion-based models, our method allows for efficient optimization of the noise schedule jointly with the rest of the model. We show that the variational lower bound (VLB) simplifies to a remarkably short expression in terms of the signal-to-noise ratio of the diffused data, thereby improving our theoretical understanding of this model class. Using this insight, we prove an equivalence between several models proposed in the literature. In addition, we show that the continuous-time VLB is invariant to the noise schedule, except for the signal-to-noise ratio at its endpoints. This enables us to learn a noise schedule that minimizes the variance of the resulting VLB estimator, leading to faster optimization. Combining these advances with architectural improvements, we obtain state-of-the-art likelihoods on image density estimation benchmarks, outperforming autoregressive models that have dominated these benchmarks for many years, with often significantly faster optimization. In addition, we show how to use the model as part of a bits-back compression scheme, and demonstrate lossless compression rates close to the theoretical optimum. Code is available at https://github.com/google-research/vdm .

1 Introduction

Likelihood-based modeling is important across applications, yet diffusion models had not matched autoregressive models on density estimation benchmarks. Variational Diffusion Models address this gap with theoretical and modeling advances that achieve state-of-the-art likelihoods.

Likelihood-based generative modeling supports applications including speech synthesis, translation, and compression.
Diffusion models had shown strong perceptual generation but had yet to match autoregressive models on density estimation benchmarks.Autoregressive models had long been dominant because they combine tractable likelihoods with expressivity.
VDMs achieve new state-of-the-art log-likelihoods on CIFAR-10 and ImageNet density estimation benchmarks.The approach incorporates Fourier features, a learnable diffusion process, and other modeling innovations.
The VLB simplifies in terms of the diffusion process’s signal-to-noise ratio, yielding an invariance result and equivalence among several diffusion models.The continuous-time model and its VLB are invariant to the diffusion-process specification, while models from the literature are equivalent up to time-dependent data rescaling.

2 Related work

The paper builds on diffusion models, denoising score matching, and continuous-time formulations while positioning its VLB analysis as a complementary perspective. Its modeling approach jointly learns the diffusion process rather than fixing it.

Diffusion probabilistic models can be viewed as variational autoencoders whose structure supports efficient training of arbitrarily deep models.
Earlier diffusion work emphasized image-generation quality, while later methods improved likelihoods and extended denoising score matching to continuous time.
Concurrent work also derived continuous-time diffusion likelihood bounds, but this paper emphasizes an intuitive VLB expression and comparisons across formulations.The related approaches use different derivations, including stochastic differential equations and infinitely deep variational autoencoders.
Unlike previous approaches with fixed diffusion processes, the proposed model jointly optimizes diffusion-process parameters and the rest of the model.The diffusion process is parameterized through marginal means and variances, and Fourier features improve likelihoods over prior diffusion models.

3 Model

The model defines a Gaussian diffusion from data to increasingly noisy latents and reverses it into a hierarchical generative model. It learns the noise schedule and uses denoising, Fourier features, and a VLB objective for likelihood modeling.

3.1 Forward time diffusion process: A Gaussian forward process maps data x to latent variables z_t that become increasingly noisy from t=0 to t=1.The conditional latent distribution is defined for every t in [0,1].
3.1 Forward time diffusion process: The signal-to-noise ratio is assumed to decrease strictly with time, formalizing the forward process’s increasing noise.The diffusion coefficients are also assumed smooth, with finite time derivatives.
3.1 Forward time diffusion process: The framework includes variance-preserving and variance-exploding diffusion processes as special cases, while experiments use the variance-preserving version.
3.3 Reverse time generative model: The reverse model samples latents backward from t=1 to t=0 and supports both finite-step and continuous-time formulations.For finite T, time is discretized uniformly into T segments of width τ=1/T.
3.3 Reverse time generative model: The reverse process uses a denoising prediction of x from noisy z_t, with a spherical Gaussian prior at the most-noisy endpoint under sufficiently small endpoint SNR.The reconstruction distribution is factorized over pixels and approximates the true conditional when the initial SNR is sufficiently large.
3.4 Noise prediction model and Fourier features: Fourier features amplify fine-scale changes in noisy inputs and substantially improve likelihood, especially when combined with a learnable SNR function.The features append sinusoidal channels to the input of the noise-prediction model.
3.5 Variational lower bound: Training maximizes the VLB, whose diffusion component depends on the generative model depth T.The model uses a neural-network noise predictor and estimates the prior and reconstruction terms with standard stochastic differentiable techniques.

4 Discrete-time model

The discrete-time model expresses its diffusion loss as a Monte Carlo objective over uniformly sampled timesteps and noisy latents. The resulting VLB can be optimized jointly over model and schedule parameters, with more timesteps improving the bound when denoising is sufficiently accurate.

For finite T, the diffusion loss is formed from KL terms between true and modeled reverse-conditionals across discretized timesteps.
The finite-time diffusion loss simplifies to a general expression valid for any forward diffusion parameters.The latent is sampled as z_t=α_t x+σ_t ε, with the timestep drawn uniformly from {1,...,T}.
The discrete-time model jointly optimizes noise-schedule parameters η and denoising parameters θ by maximizing a Monte Carlo estimate of the VLB.
The simplified objective supports numerically stable implementation in 32-bit or lower-precision floating point through expm1.Earlier discrete-time diffusion implementations required 64-bit floating point for numerical stability.
When the denoising model is sufficiently good, doubling the number of timesteps yields a better VLB: L_2T(x)<L_T(x).

5 Continuous-time model: T →∞

In the continuous-time limit, the diffusion loss becomes an integral over signal-to-noise ratio and depends on the diffusion schedule only through endpoint SNR values. This schedule invariance also yields equivalence among continuous-time diffusion specifications, including variance exploding and variance preserving forms.

Continuous-time limit: Taking T→∞ turns the discrete diffusion loss into a continuous-time loss for a stochastic differential equation.The finite-step model converges to a continuous diffusion process governed by an SDE.
Continuous-time loss: The continuous-time loss can be written by changing variables from time t to signal-to-noise ratio v.The transformed integration limits are SNRmin = SNR(1) and SNRmax = SNR(0).
Schedule invariance: The diffusion loss is invariant to the shape of SNR(t) between its endpoints, so the VLB depends on the schedule only through SNRmin and SNRmax.This result holds for any forward diffusion parameters σt and αt.
Equivalence of diffusion models: Continuous-time diffusion specifications with equal endpoint SNR values define the same data distribution when paired denoising functions satisfy the stated rescaling relationship.The latent distributions are also identical up to a trivial rescaling.
Weighted diffusion loss: The equivalence extends to weighted diffusion objectives, whose weighting function can emphasize noisier data and improve perceptual metrics such as FID and Inception Score.The paper uses w(v) = 1 for its unweighted VLB.
Optimization: Low-discrepancy sampling reduces estimator variance, while schedule optimization can further minimize continuous-time loss-estimator variance.These changes generally improve optimization efficiency.

6 Experiments

Experiments evaluate VDMs on CIFAR-10 and downsampled ImageNet, showing strong likelihood results, faster optimization, variance reduction, and practical compression performance.

Likelihood and samples: VDMs establish new state-of-the-art test-set likelihoods on CIFAR-10 and downsampled ImageNet benchmarks.The CIFAR-10 likelihood-focused model reaches a FID score of 7.41, while weighted diffusion loss improves FID to 4.0.
Likelihood and samples: About 10x faster wall-clock optimization than Sparse Transformer is required for the CIFAR-10 model to surpass the previous best result of 2.80.
Training specifications: As T grows large, continuous-time training performs best, while discrete-time models perform better with small evaluation step counts.Learning the diffusion schedule improves the discrete-time VLB, and variance-minimizing schedule learning helps continuous-time training.
Noise schedule: Learned schedules reduce VLB-estimate variance by spending more time at high SNR(t) and low σ2.The learned schedule substantially reduces variance relative to continuous-time baselines and accelerates training.
Ablations: Fourier features improve likelihood, and learning SNR endpoints is necessary to exploit them fully.With a fixed Ho et al. schedule, test negative likelihood stays above 4 bits per dim; learned maximum log-SNR reaches 13.3 and yields the reported state-of-the-art likelihoods.
Lossless compression: Bits-back coding turns the discrete-time diffusion model into a lossless compression algorithm, achieving state-of-the-art net codelengths.For large Teval, a gap remains to the theoretically optimal negative-VLB codelength, and compression becomes computationally expensive.

7 Conclusion

The paper concludes that VDMs combine state-of-the-art natural-image density modeling with theoretical simplifications about diffusion processes and their variational bounds.

Conclusion: VDMs achieve state-of-the-art natural-image density modeling using a learnable diffusion specification, Fourier features, and other architectural innovations.
Conclusion: In continuous time, the VLB is invariant to the forward diffusion process specification, and several literature diffusion processes are equivalent up to time-dependent data rescaling.
Diffusion process: The forward process has a Markov structure, so q(zs, zt|x) factors into q(zs|x)q(zt|zs) for t > s.
Generative model: The conditional reverse model uses the posterior q(zs|zt, x) with the original data replaced by a denoising prediction x̂θ(zt; t).
Model interpretations: The same model class can be viewed through denoising, noise-prediction, or score-model parameterizations.The paper primarily uses the denoising interpretation theoretically and parameterizes the practical model through noise prediction.
Numerical implementation: The numerically stable expm1 operation supports implementation in 32-bit or lower-precision floating point.

B.1 Model and implementation

The implementation uses a U-Net denoising model with Fourier features and a continuous-time diffusion formulation, while estimating the diffusion loss with unbiased timestep sampling. Fourier features target fine-scale pixel details and improve likelihood by enabling lower noise levels.

B.1 Model and implementation: The denoising model follows a U-Net architecture but removes internal resampling, increases depth, and conditions on rescaled noise levels rather than directly on time.The implementation processes data at its original resolution and uses deeper networks than the cited baseline.
B.1 Model and implementation: Fourier features are added to color-channel inputs so the model can represent high-frequency distributional details at the level of individual pixels.The features are periodic high-frequency projections concatenated with the noisy input before convolutional processing.
B.1 Model and implementation: Higher SNRmax, corresponding to lower noise levels, becomes learnable with Fourier features and produces large likelihood improvements.The paper reports that this improvement was not observed when Fourier features were added to autoregressive models.
B.1 Model and implementation: The continuous-time formulation takes T →∞ and represents the diffusion as a stochastic differential equation with time running backward in the generative model.The diffusion loss is derived as the infinite-step limit of the discrete-time objective.
B.1 Model and implementation: The finite-time diffusion loss is estimated without evaluating every timestep by sampling a timestep uniformly and drawing Gaussian noise, yielding an unbiased Monte Carlo estimator.The noisy latent is reparameterized as zt = αtx + σtϵ.
B.1 Model and implementation: Doubling the number of timesteps lowers the diffusion loss because predicting data from an earlier, less noisy latent is easier.This motivates using the continuous-time VLB corresponding to T →∞.

G Equivalence of diffusion specifications

The paper shows that continuous-time diffusion specifications can be equivalent up to trivial rescaling, with latent information determined by SNR rather than separate scale parameters. It also analyzes weighting functions and variance-reduction strategies for VLB optimization.

G Equivalence of diffusion specifications: Latents from different diffusion specifications are identical up to trivial rescaling, and their information content depends only on the signal-to-noise ratio.This supports equivalence across variance exploding and variance preserving specifications under the stated constraints.
G Equivalence of diffusion specifications: Equal SNRmin and SNRmax endpoints make different specifications produce the same continuous-time diffusion loss.The equality applies to the diffusion loss in continuous time under the stated endpoint assumption.
G Equivalence of diffusion specifications: The corresponding conditional latent distributions and generative models are the same when the denoising functions are transformed consistently between specifications.The resulting reconstruction loss is also preserved under the matching denoising transformation.
G Equivalence of diffusion specifications: Low-discrepancy timestep sampling spreads minibatch times more evenly and reduces variance in the VLB estimate; omitting these methods trades simplicity for slower optimization.The schedule-variance objective can also apply to weighted diffusion losses beyond the VLB.
G Equivalence of diffusion specifications: The continuous-time noise schedule is optimized both for its endpoint SNR values and for the variance of the stochastic VLB estimator.The endpoint parameters are optimized directly with the VLB, while remaining schedule parameters minimize estimator variance.
G Equivalence of diffusion specifications: NCSNv2’s geometric noise schedule yields a constant implied weighting function and is consistent with VLB maximization.The paper contrasts this with other objectives whose weighting functions vary with SNR.
G Equivalence of diffusion specifications: The analysis generalizes the VLB–denoising-score-matching consistency result to dimension-specific noise schedules and arbitrary score multipliers.The paper states that only the special case with equal multipliers is needed for its application.

M Additional samples from our models

The paper includes additional unconditional samples from models trained on CIFAR10, 32x32 ImageNet, and 64x64 ImageNet.

M Additional samples from our models: Additional uncurated random samples are provided for unconditional models trained on CIFAR10, 32x32 ImageNet, and 64x64 ImageNet.The samples appear in Figures 9, 10, and 11.

N Lossless compression

The discrete-time diffusion model can support lossless compression through bits-back coding, with practical codelengths near the negative VLB for small evaluation timestep counts. Larger timestep counts expose implementation gaps and make compression computationally expensive.

N Lossless compression: The discrete-time diffusion model is a hierarchical latent variable model that supports lossless compression through bits-back coding.The negative VLB gives the theoretical expected coding cost when auxiliary random bits are available.
N Lossless compression: BB-ANS implementations achieve net codelengths that closely agree with the negative VLB for small Teval.For large Teval, discrepancies arise from inaccuracies representing discretized Gaussians with small standard deviations and other coding differences.
N Lossless compression: Large Teval leaves a gap from the theoretically optimal codelength and increases computational cost because each timestep requires a neural-network forward pass.Closing this gap requires more efficient bits-back coding for large timestep counts.

O Density estimation on additional data sets

The model was also evaluated at 128 × 128 resolution on CelebA-HQ and LSUN bedrooms, yielding reported test-set likelihoods without comparisons to prior methods because these are not established density-estimation benchmarks.

O Density estimation on additional data sets: 2.14 bits per dim was obtained on CelebA-HQ at 128 × 128 resolution.The reported value is a test-set likelihood.
O Density estimation on additional data sets: 1.44 bits per dim was obtained on LSUN bedrooms at 128 × 128 resolution.The reported value is a test-set likelihood.
O Density estimation on additional data sets: The authors do not compare these results with previous methods because the data sets are not established density-estimation benchmarks and downsampling methods are inconsistent.The higher-resolution model adds a UNet level at 128 × 128 with 16 residual layers and 128 channels.