Year: 2021

Abstract:
Denoising Diffusion Probabilistic Models (DDPMs) have emerged as a powerful family of generative models that, yielding high-fidelity samples and competitive log-likelihoods across a range of domains, including image and speech synthesis. Key advantages of DDPMs include ease of training, in contrast to generative adversarial networks, and speed of generation, in contrast to autoregressive models. However, DDPMs typically require hundreds-to-thousands of steps to generate a high fidelity sample, making them prohibitively expensive for high dimensional problems. Fortunately, DDPMs allow trading generation speed for sample quality through adjusting the number of refinement steps during inference. Prior work has been successful in improving generation speed through handcrafting the time schedule through trial and error. We instead view the selection of the inference time schedules as an optimization problem, and show that, with a simple dynamic programming algorithm, one can find the log-likelihood-optimal discrete time schedules for any pre-trained DDPM. Our method exploits the fact that the evidence lower bound (ELBO) can be decomposed into separate KL divergence terms, and given any computation budget, we discover the time schedule that maximizes the training ELBO exactly. Our method is efficient, has no hyper-parameters of its own, and can be applied to any pre-trained DDPM with no retraining. We discover inference time schedules requiring as few as 32 refinement steps, while sacrificing less than 0.1 bits per dimension compared to the default 4,000 steps used on an ImageNet 64x64 model.

heading: 1 INTRODUCTION
text: Denoising Diffusion Probabilistic Models (DDPMs) have emerged as a powerful class of generative models (Sohl-Dickstein et al., 2015; Ho et al., 2020). DDPMs model the data distribution through an iterative denoising process, and have been applied successfully to a variety of applications, including unconditional image generation (Song & Ermon, 2019; Ho et al., 2020; Song et al., 2021; Nichol & Dhariwal, 2021), shape generation (Cai et al., 2020), text-to-speech (Chen et al., 2021; Kong et al., 2020) and single image super-resolution (Saharia et al., 2021; Li et al., 2021).
DDPMs are easy to train, featuring a simple denoising objective (Ho et al., 2020) with noise schedules that successfully transfer across different models and datasets. This contrasts to Generative Adversarial Networks (GANs) (Goodfellow et al., 2014), which require an inner-outer loop optimization procedure that often entails instability and requires careful hyperparameter tuning. DDPMs also admit a simple non-autoregressive inference process; this contrasts to autoregressive models with often prohibitive computational costs on high dimensional data. The DDPM inference process starts with samples from the corresponding prior noise distribution (e.g., standard Gaussian), and iteratively denoises the samples under the fixed noise schedule. However, DDPMs often need hundreds-tothousands of denoising steps (each involving a feedforward pass of a large neural network) to achieve strong results. While this process is still much faster than autoregressive models, this is still often computationally prohibitive, especially when modeling high dimensional data.
There has been much recent work focused on improving the sampling speed of DDPMs. WaveGrad (Chen et al., 2021) introduced a manually crafted schedule requiring only 6 refinement steps; however, this schedule seems to be only applicable to the vocoding task where there is a very strong conditioning signal. Denoising Diffusion Implicit Models (DDIMs) (Song et al., 2020) accelerate sampling from
pre-trained DDPMs by relying on a family of non-Markovian processes. They accelerate the generative process through taking multiple steps in the diffusion process. However, DDIMs sacrifice the ability to compute log-likelihoods. Nichol & Dhariwal (2021) also explored the use of ancestral sampling with a subsequence of the original denoising steps, trying both a uniform stride and other hand-crafted strides. San-Roman et al. (2021) improve few-step sampling further by training a separate model after training a DDPM to estimate the level of noise, and modifying inference to dynamically adjust the noise schedule at every step to match the predicted noise level.
All these fast-sampling techniques rely on a key property of DDPMs – there is a decoupling between the training and inference schedule. The training schedule need not be the same as the inference schedule, e.g., a diffusion model trained to use 1000 steps may actually use only 10 steps during inference. This decoupling characteristic is typically not found in other generative models. In past work, the choice of inference schedule was often considered a hyperpameter selection problem, and often selected via intuition or extensive hyperparmeter exploration (Chen et al., 2021). In this work, we view the choice of the timesteps of the inference schedule (which we just call an inference path) as an independent optimization problem, wherein we attempt to learn the best schedule. Our approach relies on the observation that we can solve this optimization problem with dynamic programming. Given a fixed budget of K refinement steps and a pre-trained DDPM, we find the set of timesteps that maximizes the corresponding evidence lower bound (ELBO). As an optimization objective, the ELBO has a key decomposability property: the total ELBO is the sum of individual KL terms, and for any two inference paths, if the timesteps (s, t) contiguously occur in both, they share a common KL term, therefore admitting memoization (see Section 4.1 for a precise definition).
Our main contributions are the following: • We introduce a method that that finds the likelihood-optimal inference paths with a simple
dynamic programming algorithm for all possible computation budgets of K refinement steps. The algorithm searches over T > K timesteps, only requiring O(T ) neural network forward passes. It only needs to be applied once to a pre-trained DDPM, does not require training or retraining a DDPM, and is applicable to both time-discrete and time-continuous DDPMs.
• We experiment with DDPM models from prior work. On both Lsimple CIFAR10 and Lhybrid ImageNet 64x64, we discover schedules which require only 32 refinement steps, yet sacrifice only 0.1 bits per dimension compared to their original counterparts with 1,000 and 4,000 steps, respectively.
• We show that our method can be applied to any decomposable set of objectives. In particular, optimizing a reweighted ELBO can favourably bias our algorithm towards solutions with better FID scores, as we find that optimizing the exact variational lower bound may lead to worse FID scores, which is consistent with prior work on unconditional image generation.

heading: 2 BACKGROUND ON DENOISING DIFFUSION PROBABILISTIC MODELS
text: Denoising Diffusion Probabilistic Models (DDPMs) (Ho et al., 2020; Sohl-Dickstein et al., 2015) are defined in terms of a forward Markovian diffusion process q and a learned reverse process pθ. The forward diffusion process gradually adds Gaussian noise to a data point x0 through T iterations,
q(x1:T | x0) = ∏T
t=1 q(xt | xt−1) , (1)
q(xt | xt−1) = N (xt | √ αt xt−1, (1− αt)I) , (2)
where the scalar parameters α1:T determine the variance of the noise added at each diffusion step, subject to 0 < αt < 1. The learned reverse process aims to model q(x0) by inverting the forward process, gradually removing noise from signal starting from pure Gaussian noise xT ,
p(xT ) = N (xT | 0, I) (3) pθ(x0:T ) = p(xT ) ∏T
t=1 pθ(xt−1 | xt) (4)
pθ(xt−1 | xt) = N (xt−1 | µθ(xt, t), σ2t I) . (5) The parameters of the reverse process can be optimized by maximizing the following variational lower bound on the training set:
Eq log p(x0) ≥ Eq [ log pθ(x0|x1)−
T∑ t=2 DKL ( q(xt−1|xt,x0)‖pθ(xt−1|xt) ) − LT (x0)
] (6)
where LT (x0) = DKL ( q(xT |x0) ‖ p(xT ) ) . Nichol & Dhariwal (2021) have demonstrated that training DDPMs by maximizing the ELBO yields competitive log-likelihood scores on both CIFAR10 and ImageNet 64×64 achieving 2.94 and 3.53 bits per dimension respectively. Two notable properties of Gaussian diffusion process that help formulate DDPMs tractably and efficiently include:
q(xt | x0) = N (xt | √ γt x0, (1− γt)I) , where γt = ∏t i=1 αi , (7)
q(xt−1 | x0,xt) = N ( xt−1 ∣∣∣ √γt−1 (1− αt)x0 +√αt (1− γt−1)xt 1− γt , (1− γt−1)(1− αt) 1− γt I ) .(8)
Given the marginal distribution of xt given x0 in (7), one can sample from the q(xt | x0) independently for different t and perform SGD on a randomly chosen KL term in (6). Furthermore, given that the posterior distribution of xt−1 given xt and x0 is Gaussian, one can compute each KL term in (6) between two Gaussians in closed form and avoid high variance Monte Carlo estimation.

heading: 3 LINKING DDPMS TO CONTINUOUS TIME AFFINE DIFFUSION PROCESSES
text: Before describing our approach to efficiently sampling from DDPMs, it is helpful to link DDPMs to continuous time affine diffusion processes, as it shows the compatibility of our approach to both time-discrete and time-continuous DDPMs (Song et al., 2021; Kingma et al., 2021). Let x0 ∼ q(x0) denote a data point drawn from the empirical distribution of interest and let q(xt|x0) denote a stochastic process for t ∈ [0, 1] defined through an affine diffusion process through the following stochastic differential equation (SDE):
dXt = fsde(t)Xtdt+ gsde(t)dBt , (9) where fsde, gsde : [0, 1]→ [0, 1] are integrable functions satisfying fsde(0) = 1 and gsde(0) = 0. Following Särkkä & Solin (2019) (section 6.1), we can compute the exact marginals q(xt|xs) for any 0 ≤ s < t ≤ 1. We get:
q(xt | xs) = N ( xt ∣∣∣ψ(t, s)xs,(∫ t s ψ(t, u)2g(u)2du ) I ) (10)
where ψ(t, s) = exp ∫ t s f(u)du. Since these integrals are difficult to work with, we instead propose (in parallel to Kingma et al. (2021)) to define the marginals directly: q(xt | x0) = N (xt | f(t)x0, g(t)2I) (11)
where f, g : [0, 1] → [0, 1] are differentiable, monotonic functions satisfying f(0) = 1, f(1) = 0, g(0) = 0, g(1) = 1. Then, by implicit differentiation it follows that the corresponding diffusion is
dXt = f ′(t)
f(t) Xtdt+
√ 2g(t) ( g′(t)− f ′(t)g(t)
f(t)
) dBt . (12)
We provide a proof for Equation 12 in the appendix (A.1). To complete our formulation, let fts = f(t) f(s) and gts = √ g(t)2 − f2tsg(s)2. Then, it follows that for any 0 < s < t ≤ 1 we have that
q(xt | xs) = N ( xt | ftsxs, g2tsI ) , (13)
q(xs | xt,x0) = N ( xs ∣∣∣ 1 g2t0 (fs0g 2 tsx0 + ftsg 2 s0xt), g2s0g 2 ts g2t0 I ) , (14)
We include proofs for (13) and (14) in the appendix (A.2). These equations show that we can perform inference with any ancestral sampling path (i.e., the timesteps can attain continuous values) by formulating the reverse process in terms of the posterior distribution as
pθ(xs | xt) = q ( xs | xt, x̂0 = 1ft0 (xt − gt0 θ(xt, t)) ) , (15)
justifying the compatibility of our main approach with time-continuous DDPMs. We note that this reverse process is also mathematically equivalent to a reverse process based on a time-discrete DDPM derived from a subsequence of the original timesteps as done by Song et al. (2020); Nichol & Dhariwal (2021). For the case of s = 0 in the reverse process, we follow the parametrization of Ho et al. (2020) to obtain discretized log likelihoods and compare our log likelihoods fairly with prior work.
Algorithm 1: Given a matrixL ∼ (T+1)×(T+1) of precomputed L(·, ·) terms, find the likelihoodoptimal schedules for all step budgets. def vectorized_dp_all_budgets(L): T = len(L) - 1 D = np.full(L.shape, -1) C = np.full(L.shape, np.inf) C[0, 0] = 0 for k in range(1, T + 1): bpds = C[k - 1, None] + L C[k] = np.amin(bpds, axis=-1) D[k] = np.argmin(bpds, axis=-1)
return D
Algorithm 2: Fetch the shortest path of K steps from the dynamic programming results implicitly returned by Algorithm 1. def fetch_shortest_path(D, K): optpath = [] t = K for k in reversed(range(K)): optpath.append(t) t = D[k, t]
return optpath

heading: 4 LEARNING TO EFFICIENTLY SAMPLE FROM DDPMS
text: We now introduce our dynamic programming (DP) approach. In general, after training a DDPM, one can use a different inference path than the one used during training. Additionally, one can optimize a loss or reward function with respect to the timesteps themselves after the DDPM is trained. In this paper, we use the ELBO as our loss function, however we note that it is possible to directly optimize the timesteps with other objectives.

heading: 4.1 OPTIMIZING THE ELBO
text: In our work, we choose to optimize ELBO as our objective. We rely on one key property of ELBOs, their decomposability. Before defining decomposability, we formally define a K-step inference path as a finite, monotonically increasing sequence of timesteps 0 = t′0 < t ′ 1 < ... < t ′ K−1 < t ′ K = 1. Now, given a set S ⊆ [0, 1], we define a family of lower bounds L of an “ideal” objective Lideal to be S-decomposable if:
1. There is a bijection from L to the set of all inference paths t with all timesteps in S, and 2. There exists a function L : S × S → [0,∞) such that, for all inference paths t with all
timesteps in S, Lideal ≥ ∑|t|−1 i=1 L(ti, ti−1) + C (C a constant).
We now show that DDPM ELBOs are decomposable. As shown by Song et al. (2020); Nichol & Dhariwal (2021) and the equations in Section 3, for any K and any K-step inference path t, there is a corresponding ELBO
− LELBO = EqDKL ( q(x1|x0)‖pθ(x1) ) + K∑ i=1 L(t′i, t ′ i−1) (16)
where
L(t, s) = { −Eq log pθ(xt|x0) s = 0 EqDKL ( q(xs|xt,x0)‖pθ(xs|xt) ) s > 0
(17)
Since all of these are lower bounds of Eq log p(x0), we conclude the family of DDPM evidence lower bounds is decomposable. Specifically, a DDPM trained on a set of timesteps S admits Sdecomposable ELBOs. For DDPMs trained with continuous timesteps, S = [0, 1]. For DDPMs trained on discrete timesteps, S is the set of those timesteps, as there is no guarantee that the behavior of the model won’t be pathological when give timesteps it has never seen during training. Now the question remains, given a fixed budget K steps, what is the optimal inference path?
First, we observe that any two paths that share a (t, s) transition will share a common L(t, s) term. We exploit this property in our dynamic programming algorithm. When given a grid S of plausible inference paths 0 = t0 < t1 < ... < tT−1 < tT = 1 with T ≥ K, it is possible to efficiently find the ELBO-optimal K-step inference path contained in S by memoizing all the individual L(t, s) ELBO terms for s, t ∈ {t0, ..., tT } with s < t. We can then solve the canonical least-cost-path problem on a directed graph where s→ t are nodes and the edge connecting them has cost L(t, s).

heading: 4.2 DYNAMIC PROGRAMMING ALGORITHM
text: We now outline our methodology to solve the least-cost-path problem. Our solution is similar to Dijkstra’s algorithm, but it differs to the classical least-cost-path problem where the latter is typically used, as our problem has additional constraints: we restrict our search to paths of exactly K + 1 nodes, and the start and end nodes are fixed.
Let C and D be (K + 1)× (T + 1) matrices. C[k, t] will be the total cost of the least-cost-path of length k from t to 0. D will be filled with the timesteps corresponding to such paths; i.e., D[k, t] will be the timestep s immediately previous to t for the optimal k-step path (assuming t is also part of such path).
We initialize C[0, 0] = 0 and all the other C[0, ·] to∞ (the D[0, ·] are irrelevant, but for ease of index notation we keep them in this section). Then, for each k from 1 to K, we iteratively set, for each t,
C[k, t] = min s
(C[k − 1, s] + L(t, s))
D[k, t] = argmin s
(C[k − 1, s] + L(t, s))
where L(t, s) is the cost to transition from t to s (see Equation 17). For all s ≥ t, we set L(t, s) =∞ (e.g., we only move backwards in the diffusion process). This procedure captures the shortest path cost in C and the shortest path itself in D. We further observe that running the DP algorithm for each k from 1 to T (instead of K), we can extract the optimal paths for all possible budgets K. Algorithm 1 illustrates a vectorized version of the procedure we have outlined in this section, while Algorithm 2 shows how to explicitly extract the optimal paths from D.

heading: 4.3 EFFICIENT MEMOIZATION
text: A priori, our dynamic programming approach appears to be inefficient because it requires computing O(T 2) terms (recall, as we rely on all the L(t, s) terms which depend on a neural network forward pass). We however observe that a single forward pass of the DDPM can be used to compute all the L(t, ·) terms. This holds true even in the case where the pre-trained DDPM learns the variances. For example, in Nichol & Dhariwal (2021) instead of fixing them to g̃ts = gtsgs0gt0 as we outlined in the previous section, the forward pass itself still only depends on t and not s, and the variance of pθ(xs|xt) is obtained by interpolating the forward pass’s output logits v with exp(v log g2ts + (1− v) log g̃2ts). Thus, computing the table of all the L(t, s) ELBO terms only requires O(T ) forward passes.

heading: 5 EXPERIMENTS
text: We apply our method on a wide variety of pre-trained DDPMs from prior work. This emphasizes the fact that our method is applicable to any pre-trained DDPM model. In particular, we rely the CIFAR10 model checkpoints released by Nichol & Dhariwal (2021) on both their Lhybrid and Lvlb objectives. We also showcase results on CIFAR10 (Krizhevsky et al., 2009) with the exact configuration used by Ho et al. (2020), which we denote as Lsimple, as well as Lhybrid on ImageNet 64x64 (Deng et al.,
2009) following Nichol & Dhariwal (2021), training these last two models from scratch for 800K and 3M steps, respectively, but otherwise using the exact same configurations as the authors.
In our experiments, we always search over a grid that includes all the timesteps used to train the model, i.e., {t/T : t ∈ {1, ..., T − 1}}. For our CIFAR10 results, we computed the memoization tables with Monte Carlo estimates over the full training dataset, while on ImageNet 64x64 we limited the number of datapoints in the Monte Carlo estimates to 16,384 images on the training dataset.
For each pre-trained model, we compare the negative log likelihoods (estimated using the full heldout dataset) of the strides discovered by our dynamic programming algorithm against even and quadratic strides, following Song et al. (2020). We find that our dynamic programming algorithm discovers strides resulting in much better log likelihoods than the hand-crafted strides used in prior work, particularly in the few-step regime. We provide a visualization of the log likelihood curves as a function of computation budget in Figure 1 for Lsimple CIFAR10 and Lhybrid ImageNet 64x64 (Deng et al., 2009), a full list of the scores in the few-step regime in Table 1, and a visualization of the discovered steps themselves in Figure 2.

heading: 5.1 COMPARISON WITH FID
text: We further evaluate our discovered strides by reporting FID scores (Heusel et al., 2017) on 50,000 model samples against the same number of samples from the training dataset, as is standard in the literature. We find that, although our strides are yield much better log likelihoods, such optimization does not necessarily translate to also improving the FID scores. Results are included in Figure 3. This weakened correlation between log-likehoods and FID is consistent with observations in prior work (Ho et al., 2020; Nichol & Dhariwal, 2021).
To remedy this issue, we show that despite our focus in this work being likelihood, we can significantly improve the FID scores discovered by our method simply by optimizing a reweighted ELBO. Recall that, as discussed in Section 4, our proposed method is compatible with any decomposable objective. Moreover, prior work has shown that the choice of ELBO weights has a significant effect on sample quality (Ho et al., 2020; Durkan & Song, 2021), and that choosing weights corresponds to choosing an equally valid variational lower bound of the data for a DDIM (Song et al., 2020). Similarly to prior work in the VAE literature, we thus stumble upon an open problem where different variational lower bounds compatible with the model (even with the same bits/dim) can lead to samples with different qualitative charachteristics (Alemi et al., 2018). As our focus is likelihood, we leave this research question of finding the weights / ELBO that lead to most agreement with FID for future work, but nevertheless identify one such choice that favourably biases our algorithm toward this front. Details about our construction of weights and results are included in the appendix (A.3).

heading: 5.2 MONTE CARLO ABLATION
text: To investigate the feasibility of our approach using minimal computation, we experimented with setting the number of Monte Carlo datapoints used to compute the dynamic programming table of negative log likelihood terms to 128 samples (i.e., easily fit into a single batch of GPU memory). We find that, for CIFAR10, the difference in log likelihoods is negligible, while on ImageNet 64x64 there is a visible yet slight improvement in negative log likelihood when filling the table with more samples. We hypothesize that this is due to the higher diversity of ImageNet. Nevertheless, we highlight that our procedure can be applied very quickly (i.e., with just T forward passes of a neural network when using a single batch, as opposed to a running average over batches), even for large models, to significantly improve log their likelihoods in the few-step regime.

heading: 6 RELATED WORK
text: DDPMs (Ho et al., 2020) have recently shown results that are competitive with GANs (Goodfellow et al., 2014), and they can be traced back to the work of Sohl-Dickstein et al. (2015) as a restricted family of deep latent variable models. Dhariwal & Nichol (2021) have more recently shown that DDPMs can outperform GANs in FID scores (Heusel et al., 2017). Song & Ermon (2019) have also linked DDPMs to denoising score matching (Vincent et al., 2008; 2010), which is crucial to the continuous-time formulation (Song et al., 2021; Kingma et al., 2021). More recent work on the few-step regime of DDPMs (Song et al., 2020; Chen et al., 2021; Nichol & Dhariwal, 2021; San-Roman et al., 2021; Kong & Ping, 2021; Jolicoeur-Martineau et al., 2021) has also guided our research efforts. DDPMs are also very closely related to variational autoencoders (Kingma & Welling,
2013), where more recent work has shown that, with many stochastic layers, they can also attain competitive negative log likelihoods in unconditional image generation (Child, 2020). Also very closely related to DDPMs, there has also been work on non-autoregressive modeling of text sequences that can be regarded as discrete-space DDPMs with a forward process that masks or remove tokens (Lee et al., 2018; Gu et al., 2019; Stern et al., 2019; Chan et al., 2020; Saharia et al., 2020). The UNet architecture (Ronneberger et al., 2015) has been key to the recent success of DDPMs, and as shown by Ho et al. (2020); Nichol & Dhariwal (2021), augmenting UNet with self-attention (Shaw et al., 2018) in scales where attention is computationally feasible has helped bring DDPMs closer to the current state-of-the-art autoregressive generative models (Child et al., 2019; Jun et al., 2020; Roy et al., 2021).

heading: 7 CONCLUSION AND DISCUSSION
text: By regarding the selection of the inference schedule as an optimization problem, we present a novel and efficient approach to discover a likelihood-optimal inference schedule for a pre-trained DDPM with a simple dynamic programming algorithm. Our method need only be applied once to discover the schedule, and does not require re-training the DPPM. In the few-step regime, we discover schedules on Lsimple CIFAR10 and Lhybrid ImageNet 64x64 that require only 32 steps, yet sacrifice≤ 0.1 bits per dimension compared to state-of-the-art DDPMs using hundreds-to-thousands of refinement steps. Our approach only needs forward passes of the DDPM neural network to fill the dynamic programming table of L(t, s) terms, and we show that we can fill the dynamic programming table with just O(T ) forward passes. Moreover, we can estimate the table using only 128 Monte Carlo samples, finding this to be sufficient even for datasets such as ImageNet with high diversity. Our method achieves strong likelihoods with very few refinement steps, outperforming prior work utilizing hand-crafted strides (Ho et al., 2020; Nichol & Dhariwal, 2021).
Despite very strong log-likelihood results, we observe that maximizing the unweighted ELBO can actually lead to higher (worse) FID scores, and on ImageNet 64x64, a decrease in sample quality for the smallest budgets K ∈ {32, 64}; this is consistent with findings in prior work (Ho et al., 2020; Nichol & Dhariwal, 2021). Nevertheless, our method is compatible with any decomposable objective such as reweighted variational lower bounds, and we show that a simple choice of reweighted ELBO (or equivalently a choice of DDIM) can remedy this issue. Developing principled methods to choose variational lower bounds or other decomposable metrics that correlate best with image quality is thus an important direction for future research. Finally, we remark that likelihood optimization itself is useful for specific applications: as well as better compression, in domains such as non-autoregressive text generation where likelihood correlates much better with sample quality and where diffusion models are starting to make progress (Austin et al., 2021), our method has the potential to improve sampling speed with far less cost in generation fidelity and without such adaptations.

heading: REPRODUCIBILITY STATEMENT
text: We will fully open source our work and provide code pointers in the paper in the camera-ready version. Nevertheless, we provide pseudocode with a complete implementation of our proposed algorithm to maximize ease of reproducibility while we work on open-sourcing our work (see Algorithms ?? and ??). Since we experiment with open-sourced datasets and pre-trained models that already have publicly available checkpoints, our work is fully reproducible. We additionally emphasize that our method has no hyperparameters of its own.

heading: ETHICS STATEMENT
text: Innovations in generative models have the potential to enable harmful and unethical applications. In applications where no harm is intended, bias and other failure modes of generative models and datasets used to train them can also lead to issues with fairness, discrimination, and other forms of harm. While our work is focused on making diffusion models more efficient, we believe its public release will not cause any form of immediate harm, as much more efficient generative models for images like GANs can still achieve high sample quality at a fraction of the speed of diffusion models.

heading: A APPENDIX

heading: A.1 PROOF FOR EQUATION 12
text: From Equation 10, we get by implicit differentiation that
f(t) = ψ(t, 0) = exp (∫ t 0 fsde(u)du ) ⇒f ′(t) = exp
(∫ t 0 fsde(u)du ) d dt ∫ t 0 fsde(u)du = f(t)fsde(t)
⇒fsde(t) = f ′(t)
f(t)
Similarly as above and also using the fact that ψ(t, s) = ψ(t,0)ψ(s,0) ,
g(t)2 = ∫ t 0 ψ(t, u)2gsde(u) 2du = ∫ t 0 f(t)2 f(u)2 gsde(u) 2du = f(t)2 ∫ t 0 gsde(u) 2 f(u)2 du
⇒2g(t)g′(t) = 2f(t)f ′(t) g(t) 2
f(t)2 + f(t)2
d
dt ∫ t 0 gsde(u) 2 f(u)2 du = 2fsde(t)g(t) 2 + gsde(t) 2
⇒gsde(t) = √ 2(g(t)g′(t)− fsde(t)g(t)2).

heading: A.2 PROOF FOR EQUATIONS 13 AND 14
text: From Equation 10 and ψ(t, s) = ψ(t,0)ψ(s,0) it is immediate that fts is the mean of q(xt|xs). To show that g2ts is the variance of q(xt|xs), Equation 10 implies that
Var[xt|xs] = ∫ t s ψ(t, u)2gsde(u) 2du
= ∫ t 0 ψ(t, u)2gsde(u) 2du− ∫ s 0 ψ(t, u)2gsde(u) 2du
= g(t)2 − ψ(t, 0)2 ∫ s 0
ψ(s, u)2
ψ(s, u)2ψ(u, 0)2 gsde(u)
2du
= g(t)2 − ψ(t, 0)2 ∫ s 0 ψ(s, u)2 ψ(s, 0)2 gsde(u) 2du = g(t)2 − ψ(t, s)2g(s)2
= g(t)2 − ftsg(s)2.
The mean of q(xs|xt, x0) is given by the Gaussian conjugate prior formula (where all the distributions are conditioned on x0). Let µ = ftsxs, so we have a prior over µ given by
xs|x0 ∼ N (fs0x0, g2s0Id)⇒ µ|x0 ∼ N (fs0ftsx0, f2tsg2s0Id) ∼ N (ft0x0, f2tsg2s0Id),
and a likelihood with mean µ
xt|xs, x0 ∼ xt|xs ∼ N (ftsxs, g2tsId)⇒ xt|µ, x0 ∼ xt|µ ∼ N (µ, g2tsId).
Then it follows by the formula that µ|xt, x0 has variance
Var[µ|xt, x0] = ( 1
f2tsg 2 s0
+ 1
g2ts
)−1 = ( g2ts + f 2 tsg 2 s0
f2tsg 2 s0g 2 ts
)−1 = f2tsg 2 s0g 2 ts
g2ts + f 2 tsg 2 s0
⇒Var[xs|xt, x0] = 1
f2ts Var[µ|xt, x0] =
g2s0g 2 ts
g2ts + f 2 tsg 2 s0
= g2s0g 2 ts
g2t0 = g̃2ts
and mean E[µ|xt, x0] = ( 1
f2tsg 2 s0
+ 1
g2ts )−1( ft0x0 f2tsg 2 s0 + xt g2ts ) = ft0g 2 tsx0 + f 2 tsg 2 s0xt g2ts + f 2 tsg 2 s0 = ft0g 2 tsx0 + f 2 tsg 2 s0xt g2t0
⇒E[xs|xt, x0] = 1
fts E[µ|xt, x0] =
ft0 fts g2tsx0 + ftsg 2 s0xt
g2t0 = fs0g
2 tsx0 + ftsg 2 s0xt
g2t0 = f̃ts(xt, x0).

heading: A.3 REWEIGHTED ELBO RESULTS
text: While the focus of our work is likelihood, we report FID scores for the sake of completeness, as well as to show the adaptability of our method via reweighting to focus on sample quality, as mentioned in Section 5.1.
To choose a reweighting scheme for each L(t, s) term that takes into account both t and s, Kingma et al. (2021) show that choosing discretized terms
Lw(t, s) = [ − ∫ t s w(u)SNR′(u)du ] ‖x0 − x̂0(xt, t)‖2 (18)
ensures that at the limit of infinite timesteps the continuous ELBO becomes − 12Eq ∫ 1 0 w(t)SNR′(t)‖x0 − x̂0(xt, t)‖2dt (where SNR(t) = f 2 t0 g2t0 , and a constant w(t) = 1 yields the unweighted ELBO).
While choosing w(t) = − SNR(t)SNR′(t) (which is Lsimple at the limit) did not work, we find that choosing w(t) = − 1SNR′(t) (which leads to a continuous objective of unweighted mean square errors and Lw(t, s) = (t− s)‖x0 − x̂0(xt, t)‖2) allows our algorithm to outperform DDPM FID scores and achieve similar scores to DDIM (Song et al., 2020). We call this “MSE reweighting” and include FID scores in Table 5, comparing to DDPM but also DDIM(η = 0) which does not admit likelihood computation but has been shown to be one of the strongest FID baselines in the few-step regime. The FID scores were estimated with 50,000 model and training data samples, as is standard in the literature.

heading: A.4 NOTE ON SAMPLING STRATEGIES FOR FAIR COMPARISON
text: Finally, we discuss an alternative approach to that used in the figures of the paper for producing more qualitatively comparable samples. Given a fixed budget K, Figures 4 and 5 are produced to generate “comparable” samples across different strides by fixing all the standard Gaussian vectors in the sampling chain. However, another approach that allows to compare samples across different budgets is to fix a single Brownian motion trajectory on all T steps, and using discretizations (based on any stride) of this single random trajectory to generate the samples. We empirically find, however, that the former approach tends to produce much more similar images (see Figures 7 and 8 below). We suspect that the use of different strides (and hence different random directions from the fixed random trajectory) along with the very chaotic, non-linear behavior of DDPMs is the cause of this behavior.

Review summaries:
This work presents a method to efficiently sample from a pre-trained DDPM by solving a dynamic programming problem that can maximize the log likelihood of the data samples given a fixed computational budget. This is done by defining a least-cost path problem to select a reduced set of time steps among a full grid of potential time steps across different possible step budget sizes, where the ELBO is used as the cost function. The authors show that their method can identify DDPM schedules that can achieve significantly higher log likelihood (i.e. lower bits/dim) than prior DDPM schedules in the regime where about a hundred steps or fewer are used.

Samples are generated from DDPMs by solving an SDE (often in "discrete time", which is used to refer to specifically the Euler--Maruyama discretisation). This necessitates a choice for where to make numerical steps. Each choice of step locations has a corresponding ELBO. This paper demonstrates that (on a pretrained model) the optimal ELBO may be obtained via a dynamic programming algorithm for the location of the steps.