Variational Diffusion Models

In plain words

Diffusion models had become known for making convincing images, but on the benchmark that measures how well a model assigns probability to real images, autoregressive models (which predict one pixel after another) had been ahead for years. This paper shows that a diffusion model can be trained on the same kind of score that such benchmarks use, and that it can keep up with them.

A diffusion model learns to undo a gradual noising process. The paper lets the model learn how much noise is added at each stage, together with the network that removes it. The score it trains on depends on that noise schedule only at its two ends, which makes the training signal easier to estimate. The network also gets extra high-frequency inputs, so that it can represent the fine pixel-level detail that the score cares about.

On the standard image benchmarks the authors report the best scores so far. They also show that the same model can drive a lossless compressor, whose code length agrees closely with the bound when the number of steps is small.

TL;DR

• A diffusion model with a learned noise schedule, trained on the variational lower bound (VLB). In continuous time the bound becomes a schedule-weighted MSE on the predicted noise §3.2 eq. 14 eq. 17.
• The continuous-time bound depends on the schedule only through its endpoint signal-to-noise ratios (SNRs), so the interior of the schedule can be chosen to reduce estimator variance §5.1 eq. 18 §5.3.
• Author-reported: 2.65 bits per dim (BPD) on CIFAR-10 without augmentation and 2.49 with it; 3.72 on ImageNet 32x32 and 3.40 on ImageNet 64x64 tab. 1. These are bounds, not exact likelihoods tab. 1 caption.
• Key ingredients: Fourier features on the denoiser input §3.4 app. C, a U-Net denoiser that changes several details of Ho et al.'s architecture app. B.1, and a monotone network for the schedule app. H.
• Official Apache-2.0 JAX code is linked. It differs from the paper in several places, including the schedule used for CIFAR-10 and ImageNet; see "Paper vs code".

Problem

Autoregressive models had long been the best on density-estimation benchmarks such as CIFAR-10 and downsampled ImageNet, because they give a tractable likelihood §1 fig. 1. Diffusion models gave strong samples but had not matched those likelihoods §1.

Earlier diffusion models fixed the forward noising process and parameterised each reverse step §2. This paper asks whether a diffusion model can be trained as a likelihood model, with the noise schedule learned jointly with the rest of the network §1 §3.2.

Method

The forward process adds Gaussian noise to the data x over time t, from t = 0 (least noisy) to t = 1 (most noisy). The generative model inverts it, and the denoiser replaces the unknown x §3.1 §3.3. Training maximises a variational lower bound on log p(x) with three terms: a prior term, a reconstruction term and a diffusion term §3.5 eq. 11. The paper uses variance-preserving noise. The schedule is a monotone function gamma(t) = -log SNR(t) §3.1 §3.2.

Steps, with the paper's notation:

1. Forward process. q(z_t | x) = N(alpha_t x, sigma_t^2 I) eq. 1, with SNR(t) = alpha_t^2 / sigma_t^2 eq. 2 strictly decreasing in t. Smoothness is assumed. The variance-preserving choice alpha_t = sqrt(1 - sigma_t^2) is the one used in experiments. The variance-exploding case alpha_t = 1 is also covered §3.1.
2. Learned schedule. sigma_t^2 = sigmoid(gamma(t)) eq. 3, alpha_t^2 = sigmoid(-gamma(t)) eq. 4 and SNR(t) = exp(-gamma(t)) eq. 5. gamma is a monotone neural network with positive weights §3.2 app. H.
3. Generative model. The prior is p(z_1) = N(0, I), an approximation that holds when SNR(1) is small enough §3.3 eq. 7. The reconstruction p(x | z_0) factorises over the pixels. Each pixel's term is q(z_{0,i} | x_i) normalised over the 256 integer values §3.3 eq. 8. The reverse step is p(z_s | z_t) = q(z_s | z_t, x = x_hat(z_t; t)) §3.3 eq. 9.
4. Denoiser. The network predicts the noise eps_hat, and x_hat = (z_t - sigma_t eps_hat) / alpha_t §3.4 eq. 10. Its input is the latent with extra channels sin(2^n pi z) and cos(2^n pi z) for integer n in a range §3.4 app. C. Time enters as gamma(t) rescaled to about [0, 1] app. B.1.
5. Bound. -log p(x) <= KL(q(z_1|x) || p(z_1)) + E[-log p(x|z_0)] + L_T eq. 11. The paper says the prior and reconstruction terms can be estimated with standard techniques §3.5. The released code evaluates the Gaussian prior KL in closed form model_vdm.py.
6. Discrete time with T steps, t(i) = i/T and s(i) = (i-1)/T §3.3. The diffusion term L_T is a sum of KL terms eq. 12, which simplifies to T/2 E[(SNR(s) - SNR(t)) ||x - x_hat||^2] with i ~ U{1..T} eq. 13. With the schedule and the noise prediction it becomes T/2 E[(exp(gamma(t) - gamma(s)) - 1) ||eps - eps_hat||^2] eq. 14 §4.
7. Continuous time (T -> infinity). The diffusion term becomes an integral over t eq. 15 eq. 16. With the schedule it simplifies to 1/2 E_{eps, t ~ U(0,1)}[gamma'(t) ||eps - eps_hat||^2] eq. 17, estimated by Monte Carlo for both training and evaluation §5.
8. Schedule invariance. Written in the variable v = SNR(t), the loss is 1/2 E ∫_{SNR_min}^{SNR_max} ||x - x_tilde||^2 dv eq. 18. It depends on the schedule only through SNR_min = SNR(1) and SNR_max = SNR(0) §5.1. Two specifications with the same endpoints give the same p(x) once one denoiser's input is rescaled §5.1 app. G.
9. Weighted loss. A general weight w(v) may multiply the squared error eq. 19. The models in the paper use w(v) = 1, which is the unweighted VLB §5.2. The loss of Ho et al. is a different weighting app. K.
10. Variance reduction. Time samples in a minibatch of size k are t^i = mod(u_0 + i/k, 1) with one u_0 ~ U[0,1]. Each t^i is uniform, but the batch covers [0, 1] evenly §5.3 app. I.1. The interior of the schedule is trained to minimise the variance of the loss estimator. The endpoints are trained on the VLB §5.3 app. I.2.
11. Sampling. An ancestral step from t to s < t is z_s = sqrt(alpha_s^2 / alpha_t^2) (z_t - sigma_t c eps_hat) + sqrt((1 - alpha_s^2) c) eps with c = -expm1(gamma(s) - gamma(t)) app. A.4. The discrete model uses T_eval such steps §6.3.

Training step for the continuous-time model, following the equations. It matches the released code in every term; see "Paper vs code" for the places where it does not:

# Continuous-time VDM loss for one minibatch  [§3.5] [§5] [eq. 17]
# Input:  images (B, H, W, C) integers in 0..255; gamma: learned schedule
# Output: scalar loss in nats per example, averaged over the batch

def vdm_loss(images, gamma, denoiser):
    # Map pixels to bin centres in (-1, 1)  [§3.4]
    x = 2.0 * ((images + 0.5) / 256.0) - 1.0

    # Reconstruction term at time 0: noisy latent, then the pixel
    # likelihood given it  [§3.3] [eq. 8]
    noise_0 = random_normal(x.shape)
    latent_0 = x + exp(gamma(0.0) / 2.0) * noise_0
    loss_recon = -log_p_pixels(images, latent_0, gamma(0.0))

    # Prior term: Gaussian KL from q(z_1 | x) to N(0, I)  [§3.5] [eq. 11]
    var_1 = sigmoid(gamma(1.0))
    loss_prior = 0.5 * sum((1 - var_1) * x**2 + var_1 - log(var_1) - 1)

    # Times for the batch: one uniform draw, spread evenly over [0, 1)
    # [app. I.1]
    u = random_uniform()
    times = mod(u + arange(B) / B, 1.0)                   # (B,)

    # Noisy latent at those times  [§3.1] [eq. 1] [eq. 3]
    var_t = sigmoid(gamma(times))                         # (B,)
    noise = random_normal(x.shape)
    latent_t = sqrt(1 - var_t) * x + sqrt(var_t) * noise  # (B, H, W, C)

    # Predicted noise and its summed squared error  [§3.4] [eq. 10]
    noise_hat = denoiser(latent_t, gamma(times))
    squared_error = sum((noise - noise_hat)**2)           # over H, W, C

    # Continuous-time weight 1/2 * gamma'(t), the derivative of the schedule
    # [eq. 17]
    loss_diff = 0.5 * gamma_prime(times) * squared_error

    return mean(loss_recon + loss_prior + loss_diff)

The ancestral sampler from app. A.4:

# One ancestral sampling step from time t down to time s < t  [app. A.4]
# Input:  latent_t (B, H, W, C) at time t; scalar times s < t
# Output: latent_s (B, H, W, C) at time s

def sample_step(latent_t, t, s, gamma, denoiser):
    # Predicted noise at the current time  [§3.4] [eq. 10]
    noise_hat = denoiser(latent_t, gamma(t))

    # Signal and noise factors at both times  [§3.2] [eq. 3] [eq. 4]
    alpha_sq_s = sigmoid(-gamma(s))
    alpha_sq_t = sigmoid(-gamma(t))
    sigma_t = sqrt(sigmoid(gamma(t)))

    # Variance factor for the step, written with expm1  [app. A.4] [app. J]
    c = -expm1(gamma(s) - gamma(t))

    # Mean of the step, then the fresh noise added to it  [app. A.4]
    fresh_noise = random_normal(latent_t.shape)
    mean = sqrt(alpha_sq_s / alpha_sq_t) * (latent_t - sigma_t * c * noise_hat)
    return mean + sqrt((1 - alpha_sq_s) * c) * fresh_noise

Hyperparameters

Common settings. The paper's values are in app. B.1 unless a row says otherwise.

Per-dataset settings:

Results

Author-reported. Setup: test-set bits per dimension (BPD) on CIFAR-10 and on downsampled ImageNet at 32x32 and 64x64, with each baseline's numbers copied from its paper §6 tab. 1. Lower is better. The paper states that the VDM numbers are variational bounds and that importance sampling or a continuous normalising flow could give tighter estimates tab. 1 caption.

Augmentation codes (from the table caption): (A) extensive, (B) small translations, (C) horizontal flips tab. 1.

• The paper says the VDM CIFAR-10 model without augmentation passes the previous best of 2.80 about ten times faster in wall-clock time than the Sparse Transformer §6.1.
• FID for the likelihood-tuned CIFAR-10 model is 7.41. With the weighted loss of Ho et al. the FID improves to 4.0. The authors did not tune further for FID §6.1 §5.2.
• Non-cherry-picked unconditional ImageNet 64x64 samples are shown in fig. 3, from a model trained in continuous time and sampled with T = 1000 fig. 3 caption.
• Ablation on the schedule and Fourier features: with Fourier features the likelihood is much better fig. 5. A fixed Ho et al. schedule keeps the maximum log-SNR near 8 and stays above 4 BPD on the test set. Learning the endpoints lets the maximum log-SNR reach 13.3 §6.2.
• Variance of the per-example VLB estimate on the CIFAR-10 test set (right panel of fig. 4, conditional on the data, excluding minibatch sampling noise): learned schedule 0.53; log-SNR-linear 6.35; beta-linear 31.6; alpha-cosine 31.1 fig. 4. All schedules were given the same endpoints, and each gives a VLB estimate of 2.66 fig. 4 caption.
• Discrete versus continuous training and evaluation on CIFAR-10 tab. 2:

The bits-back column is blank for the discrete rows in the source tab. 2. The compression is done with BB-ANS in a Bit-Swap arrangement, with uniform latent discretisation and an XOR encryption of the ANS stream §6.3. Without encryption, a batch of 100 examples at T_eval = 250 costs 2.74 BPD rather than 2.68 §6.3.

• Density estimation on CelebA-HQ and LSUN bedrooms at 128x128 gives 2.14 and 1.44 BPD. The authors do not compare these to other methods app. O.

Paper vs code

Official code: google-research/vdm at dc27b98a554f65cdc654b800da5aa1846545d41b, Apache-2.0. The repository is linked from the abstract. It is a JAX/Flax implementation.

Discrepancies found by reading the code:

• Schedule form for three of the released configs. The paper uses a monotone network for gamma throughout §3.2 app. H. The CIFAR-10 (no augmentation), ImageNet 32 and ImageNet 64 configs set gamma_type to learnable_scalar, which computes gamma(t) = b + |w| t with two learned scalars starting at gamma_min = -13.3 and gamma_max = 5 cifar10.py model_vdm.py. Only the augmented CIFAR-10 config uses learnable_nnet cifar10_aug.py. The linear-in-log-SNR form matches the initialisation used in the paper's schedule figure fig. 4.
• Monotone network structure. The paper gives l1(t) + l3(phi(l2(l1(t)))) with sigmoid phi app. H. The code rescales the input to [-1, 1] as 2(t - 0.5) and feeds that to l2, not l1(t). It squashes the hidden layer with 2(sigmoid - 0.5) and divides the l3 output by 1024. l3 has no bias. Kernels are made positive with abs model_vdm.py model_vdm.py.
• Endpoint postprocessing and variance-minimising training of the schedule are not in the code. app. I.2 rescales the network so that it hits gamma_0 = -log SNR_max and gamma_1 = -log SNR_min exactly, and trains the interior by the gradient of the squared loss. A search of the repository for postprocess, stop_gradient and variance found no such code. The training objective is the plain VLB gradient experiment.py. The released schedule is not clamped to [gamma_min, gamma_max].
• Fourier feature indexing. The paper writes sin(2^n pi z) with n in {7, 8} §3.4 app. C. The code computes sin and cos of 2 pi 2^n z for n in {6, 7}, because Base2FourierFeatures(start=6, stop=8) model_vdm.py model_vdm.py. The frequencies are the same, since 2 pi 2^6 = 2^7 pi and 2 pi 2^7 = 2^8 pi. Use one convention, not both.
• Residual output. The code returns net(z) + z for the noise prediction, with the final convolution zero-initialised model_vdm.py. The paper describes eps_hat as a neural network and does not mention the skip connection §3.4.
• Attention blocks. The paper uses only the middle attention block for CIFAR-10 without augmentation and for ImageNet app. B.1. The released cifar10_aug.py sets with_attention=False while its comment says "True in paper version" cifar10_aug.py. cifar10_aug_paper.py sets it to True cifar10_aug_paper.py. The middle attention block is always built model_vdm.py.
• Open-source augmented CIFAR-10 is not the paper's model. cifar10_aug.py uses 256 channels, dropout 0.05, batch 768 and learning rate 1e-4. The "paper version" uses batch 1024 and learning rate 2e-4 cifar10_aug.py cifar10_aug_paper.py cifar10_aug_paper.py. The README reports eval_bpd = 2.522 for the open model against 2.49 in the paper, and attributes the gap to the smaller batch and to how the model is conditioned README.md.
• ImageNet configs do not match the paper's batch size or length. The paper uses a total batch of 512 and 2M updates for ImageNet app. B.2 app. B.1. imagenet32.py and imagenet64.py inherit a batch of 128 and 10M steps from the CIFAR-10 base config imagenet64.py. The ImageNet 64 config does use 64 ResNet blocks and 256 channels, as the paper states imagenet64.py.
• Augmentation conditioning. The paper says conditioning on the augmentation flag made almost no difference app. B.2. The augmented input pipeline does return the flag as conditioning dataset.py, and the model concatenates it to the time embedding model_vdm.py.
• Warm-up and steps. A 100-step linear warm-up, constant learning rate afterwards, and a substeps setting of 1000 are code-only details cifar10.py experiment.py.
• Time embedding scale. The code maps gamma to [0, 1] using gamma_min and gamma_max model_vdm.py, then multiplies by 1000 before a sinusoidal embedding model_vdm.py. The paper says only that the embedding follows Ho et al. app. B.1. Changing gamma_min or gamma_max in the config therefore changes the network's input scale as well as the schedule's starting point.
• Pixel encoding. The code maps integer pixels to bin centres 2((x + 0.5)/256) - 1 model_vdm.py. The paper says the data are scaled to [-1, 1] §3.4. The reconstruction term uses the same bin centres model_vdm.py.
• Bits per dimension. The code reports (recon + latent + diff) / (D ln 2) with one sampled time per example experiment_vdm.py. This matches the paper's VLB in bits.
• Numerical forms. The discrete loss uses expm1 model_vdm.py, and so does the sampler model_vdm.py. Both match app. J and eq. 14.
• Sampler. The sampler in model_vdm.py matches app. A.4 term by term model_vdm.py. No discrepancy found.
• Reported number for CIFAR-10 without augmentation. The README says the released checkpoint gives eval_bpd = 2.637 and calls it a match for the paper's 2.64. The paper's Table 1 says 2.65 README.md tab. 1.

Pitfalls

• These are bounds. The test BPD numbers are variational upper bounds on negative log-likelihood, and the paper expects tighter values from importance sampling tab. 1 caption. Do not report them as exact likelihoods.
• Equation and figure numbers were counted from the LaTeX source. They follow the order in which the source numbers equations and floats. Check them against the compiled PDF before quoting them elsewhere. Appendix letters follow the source order as well.
• Sum, not mean, over pixels. The squared error in the code is summed over pixels model_vdm.py, and the BPD divides by D ln 2 experiment_vdm.py. A mean instead of a sum changes the loss scale by a factor of D, which in turn changes the effective learning rate.
• Use expm1 for the discrete term. exp(g_t - g_s) - 1 loses precision in 32-bit floats when the difference is small. Write it as expm1 §4 app. J. The paper also gives a stable form of the sampling variance, -expm1(softplus(gamma(s)) - softplus(gamma(t))) app. J. Figure 6 plots the errors of the naive form against expm1 fig. 6.
• Learn the schedule endpoints. With Fourier features, a fixed Ho et al. schedule stays above 4 BPD on the test set. The reported numbers come from learned endpoints (maximum log-SNR 13.3) together with Fourier features §6.2.
• Likelihood and FID pull apart. The likelihood-tuned CIFAR-10 model has FID 7.41. Weighting the loss as in Ho et al. improves FID to 4.0 but changes the objective §5.2 §6.1.
• Compute. The CIFAR-10 model without augmentation reaches 2.80 BPD within about 2.5 hours. Reaching 2.65 BPD takes 10M updates, about 9 days on 8 TPUv3 chips app. B.2. The augmented model trained for slightly over a week on 128 chips app. B.2.
• Few-step evaluation of a continuous model is poor. A model trained in continuous time and evaluated with T_eval = 10 gives 7.54 BPD tab. 2. Evaluation with many steps is needed for the reported numbers.
• Sampling and compression cost one network call per step. Compression with large T_eval has a gap to the theoretical limit of up to 0.05 bits per dimension and is slow §6.3.
• Naive bits-back needs encryption. Without XOR-encrypting the ANS stream before each decode, a batch of 100 at T_eval = 250 costs 2.74 bits per dimension instead of 2.68 §6.3.
• Variance minimisation is in the paper but not the released code. Reproducing the variance results in fig. 4 requires implementing app. I.2 yourself. The released code trains the schedule on the plain VLB gradient experiment.py.
• Fourier features help the diffusion model only. Adding them to PixelCNN++ gave no benefit §6.2, and the paper did not observe improvements for autoregressive models in general §3.4.
• Augmentation. Each of the three transforms has probability 0.5, so one in eight examples is unaugmented app. B.2. The released pipeline also returns the augmentation flag as a conditioning input dataset.py.
• Denoiser input scale depends on gamma_min and gamma_max. See the time-embedding note in "Paper vs code". Keep these fixed when reproducing the released checkpoints.

Lineage

• Builds on: the diffusion probabilistic model of Sohl-Dickstein et al. (2015), the denoising diffusion models of Ho et al. (2020), the score-based models of Song et al. (2019, 2020), and the improved diffusion models of Nichol and Dhariwal (2021) §2 app. K. The paper changes the forward process to be learned and reparameterises the marginals q(z_t | x) directly §2.
• Concurrent work: Song et al. (2021), Huang et al. (2021) and Vahdat et al. (2021), which also derive variational bounds for continuous-time diffusion. Song et al. and Vahdat et al. use stochastic differential equations; Huang et al. consider both the VAE and the SDE view §2 tab. 1.
• Compared against: the autoregressive, flow and VAE baselines and the diffusion models listed in tab. 1.