Analytic-DPM: an Analytic Estimate of the Optimal Reverse Variance in Diffusion Probabilistic Models

Introduction

A diffusion process gradually adds noise to a data distribution over a series of timesteps. By learning to reverse it, diffusion probabilistic models (DPMs) (Sohl-Dickstein et al., 2015; Ho et al., 2020; Song et al., 2020b) define a data generative process. Recently, it is shown that DPMs are able to produce high-quality samples (Ho et al., 2020; Nichol & Dhariwal, 2021; Song et al., 2020b; Dhariwal & Nichol, 2021), which are comparable or even superior to the current state-of-the-art GAN models (Goodfellow et al., 2014; Brock et al., 2018; Wu et al., 2019; Karras et al., 2020b).

Despite their success, the inference of DPMs (e.g., sampling and density evaluation) often requires to iterate over thousands of timesteps, which is two or three orders of magnitude slower (Song et al., 2020a) than other generative models such as GANs. A key problem in the inference is to estimate the variance in each timestep of the reverse process. Most of the prior works use a handcrafted value for all timesteps, which usually run a long chain to obtain a reasonable sample and density value (Nichol & Dhariwal, 2021). Nichol & Dhariwal (2021) attempt to improve the efficiency of sampling by learning a variance network in the reverse process. However, it still needs a relatively long trajectory to get a reasonable log-likelihood (see Appendix E in Nichol & Dhariwal (2021)).

In this work, we present a surprising result that both the optimal reverse variance and the corresponding optimal KL divergence of a DPM have analytic forms w.r.t. its score function (i.e., the gradient of a log density). Building upon it, we propose Analytic-DPM, a training-free inference framework to improve the efficiency of a pretrained DPM while achieving comparable or even superior performance. Analytic-DPM estimates the analytic forms of the variance and KL divergence using the Monte Carlo method and the score-based model in the pretrained DPM. The corresponding trajectory is calculated via a dynamic programming algorithm (Watson et al., 2021). Further, to correct the potential bias caused by the score-based model, we derive both lower and upper bounds of the optimal variance and clip its estimate for a better result. Finally, we reveal an interesting relationship between the score function and the data covariance matrix.

Analytic-DPM is applicable to a variety of DPMs (Ho et al., 2020; Song et al., 2020a; Nichol & Dhariwal, 2021) in a plug-and-play manner. Empirically, Analytic-DPM consistently improves the log-likelihood of these DPMs and meanwhile enjoys a $20\times$ to $40\times$ speed up. Besides, Analytic-DPM also consistently improves the sample quality of DDIMs (Song et al., 2020a) and requires up to 50 timesteps (which is a $20\times$ to $80\times$ speed up compared to the full timesteps) to achieve a comparable FID to the corresponding baseline.

Background

Diffusion probabilistic models (DPMs) firstly construct a forward process $q({\bm{x}}_{1:N}|{\bm{x}}_{0})$ that injects noise to a data distribution $q({\bm{x}}_{0})$, and then reverse the forward process to recover it. Given a forward noise schedule $\beta_{n}\in(0,1),n=1,\cdots,N$, denoising diffusion probabilistic models (DDPMs) (Ho et al., 2020) consider a Markov forward process:

The reverse process for Eq. (2) is defined as a Markov process aimed to approximate $q({\bm{x}}_{0})$ by gradually denoising from the standard Gaussian distribution $p({\bm{x}}_{N})={\mathcal{N}}({\bm{x}}_{N}|{\bm{0}},{\bm{I}})$:

which is equivalent to optimizing the KL divergence between the forward and the reverse process:

where $n$ is uniform between $1$ and $N$, $q_{n}({\bm{x}}_{n})$ is the marginal distribution of the forward process at timestep $n$, ${\bm{\epsilon}}$ is a standard Gaussian noise, ${\bm{x}}_{n}$ on the right-hand side is reparameterized by ${\bm{x}}_{n}=\sqrt{\overline{\alpha}_{n}}{\bm{x}}_{0}+\sqrt{\overline{\beta}_{n}}{\bm{\epsilon}}$ and $c$ is a constant only related to $q$. Indeed, Eq. (5) is exactly a weighted sum of score matching objectives (Song & Ermon, 2019), which admits an optimal solution ${\bm{s}}_{n}^{*}({\bm{x}}_{n})=\nabla_{{\bm{x}}_{n}}\log q_{n}({\bm{x}}_{n})$ for all $n\in\{1,2\cdots,N\}$.

Analytic Estimate of the Optimal Reverse Variance

For a DPM, we first show that both the optimal mean ${\bm{\mu}}_{n}^{*}({\bm{x}}_{n})$ and the optimal variance $\sigma_{n}^{*2}$ to Eq. (4) have analytic forms w.r.t. the score function, which is summarized in the following Theorem 1.

(Score representation of the optimal solution to Eq. (4), proof in Appendix A.2)

The optimal solution ${\bm{\mu}}_{n}^{*}({\bm{x}}_{n})$ and $\sigma_{n}^{*2}$ to Eq. (4) are

where $q_{n}({\bm{x}}_{n})$ is the marginal distribution of the forward process at the timestep $n$ and $d$ is the dimension of the data.

The proof of Theorem 1 consists of three key steps:

The first step (see Lemma 9) is known as the moment matching (Minka, 2013), which states that approximating arbitrary density by a Gaussian density under the KL divergence is equivalent to setting the first two moments of the two densities as the same. To our knowledge, the connection of moment matching and DPMs has not been revealed before.

In the second step (see Lemma 13), we carefully use the law of total variance conditioned on ${\bm{x}}_{0}$ and convert the second moment of $q({\bm{x}}_{n-1}|{\bm{x}}_{n})$ to that of $q({\bm{x}}_{0}|{\bm{x}}_{n})$.

In the third step (see Lemma 11), we surprisingly find that the second moment of $q({\bm{x}}_{0}|{\bm{x}}_{n})$ can be represented by the score function, and we plug the score representation into the second moment of $q({\bm{x}}_{n-1}|{\bm{x}}_{n})$ to get the final results in Theorem 1.

In contrast to the handcrafted strategies used in (Ho et al., 2020; Song et al., 2020a), Theorem 1 shows that the optimal reverse variance $\sigma_{n}^{*2}$ can also be estimated without any extra training process given a pretrained score-based model ${\bm{s}}_{n}({\bm{x}}_{n})$. In fact, we first estimate the expected mean squared norm of $\nabla_{{\bm{x}}_{n}}\log q_{n}({\bm{x}}_{n})$ by $\Gamma=(\Gamma_{1},...,\Gamma_{N})$, where

$M$ is the number of Monte Carlo samples. We only need to calculate $\Gamma$ once for a pretrained model and reuse it in downstream computations (see Appendix H.1 for a detailed discussion of the computation cost of $\Gamma$). Then, according to Eq. (7), we estimate $\sigma_{n}^{*2}$ as follows:

According to Eq. (7) and Eq. (9), the bias of the analytic estimate $\hat{\sigma}_{n}^{2}$ is

Our estimate of the variance employs a score-based model ${\bm{s}}_{n}({\bm{x}}_{n})$ to approximate the true score function $\nabla_{{\bm{x}}_{n}}\log q_{n}({\bm{x}}_{n})$. Thus, the approximation error in Eq. (10) is irreducible given a pretrained model. Meanwhile, the coefficient in Eq. (10) can be large if we use a shorter trajectory to sample (see details in Section 4), potentially resulting in a large bias.

To reduce the bias, we derive bounds of the optimal reverse variance $\sigma_{n}^{*2}$ and clip our estimate based on the bounds. Importantly, these bounds are unrelated to the data distribution $q({\bm{x}}_{0})$ and hence can be efficiently calculated. We firstly derive both upper and lower bounds of $\sigma_{n}^{*2}$ without any assumption about the data. Then we show another upper bound of $\sigma_{n}^{*2}$ if the data distribution is bounded. We formalize these bounds in Theorem 2.

(Bounds of the optimal reverse variance, proof in Appendix A.3)

$\sigma_{n}^{*2}$ has the following lower and upper bounds:

If we further assume $q({\bm{x}}_{0})$ is a bounded distribution in $[a,b]^{d}$, where $d$ is the dimension of data, then $\sigma_{n}^{*2}$ can be further upper bounded by

Analytic Estimation of the Optimal Trajectory

The number of full timesteps $N$ can be large, making the inference slow in practice. Thereby, we can construct a shorter forward process $q({\bm{x}}_{\tau_{1}},\cdots,{\bm{x}}_{\tau_{K}}|{\bm{x}}_{0})$ constrained on a trajectory $1=\tau_{1}<\cdots<\tau_{K}=N$ of $K$ timesteps (Song et al., 2020a; Nichol & Dhariwal, 2021; Watson et al., 2021), and $K$ can be much smaller than $N$ to speed up the inference. Formally, the shorter process is defined as $q({\bm{x}}_{\tau_{1}},\cdots,{\bm{x}}_{\tau_{K}}|{\bm{x}}_{0})=q({\bm{x}}_{\tau_{K}}|{\bm{x}}_{0})\prod_{k=2}^{K}q({\bm{x}}_{\tau_{k-1}}|{\bm{x}}_{\tau_{k}},{\bm{x}}_{0})$, where

The corresponding reverse process is $p({\bm{x}}_{0},{\bm{x}}_{\tau_{1}},\cdots,{\bm{x}}_{\tau_{K}})=p({\bm{x}}_{\tau_{K}})\prod_{k=1}^{K}p({\bm{x}}_{\tau_{k-1}}|{\bm{x}}_{\tau_{k}})$, where

According to Theorem 1, the mean and variance of the optimal $p^{*}({\bm{x}}_{\tau_{k-1}}|{\bm{x}}_{\tau_{k}})$ in the sense of KL minimization is

where $\alpha_{\tau_{k}|\tau_{k-1}}\coloneqq\overline{\alpha}_{\tau_{k}}/\overline{\alpha}_{\tau_{k-1}}$. According to Theorem 2, we can derive similar bounds for $\sigma_{\tau_{k-1}|\tau_{k}}^{*2}$ (see details in Appendix C). Similarly to Eq. (9), the estimate of $\sigma_{\tau_{k-1}|\tau_{k}}^{*2}$ is

where $\Gamma$ is defined in Eq. (8) and can be shared across different selections of trajectories. Based on the optimal reverse process $p^{*}$ above, we further optimize the trajectory:

where $J(\tau_{k-1},\tau_{k})=\log(\sigma_{\tau_{k-1}|\tau_{k}}^{*2}/\lambda_{\tau_{k-1}|\tau_{k}}^{2})$ and $c$ is a constant unrelated to the trajectory $\tau$ (see proof in Appendix A.4). The KL in Eq. (14) can be decomposed into $K-1$ terms and each term has an analytic form w.r.t. the score function. We view each term as a cost function $J$ evaluated at $(\tau_{k-1},\tau_{k})$, and it can be efficiently estimated by $J(\tau_{k-1},\tau_{k})\approx\log(\hat{\sigma}_{\tau_{k-1}|\tau_{k}}^{2}/\lambda_{\tau_{k-1}|\tau_{k}}^{2})$, which doesn’t require any neural network computation once $\Gamma$ is given. While the logarithmic function causes bias even when the correct score function is known, it can be reduced by increasing $M$.

As a result, Eq. (14) is reduced to a canonical least-cost-path problem (Watson et al., 2021) on a directed graph, where the nodes are $\{1,2,\cdots,N\}$ and the edge from $s$ to $t$ has cost $J(s,t)$. We want to find a least-cost path of $K$ nodes starting from $1$ and terminating at $N$. This problem can be solved by the dynamic programming (DP) algorithm introduced by Watson et al. (2021). We present this algorithm in Appendix B. Besides, we can also extend Eq. (14) to DPMs with continuous timesteps (Song et al., 2020b; Kingma et al., 2021), where their corresponding optimal KL divergences are also decomposed to terms determined by score functions. Thereby, the DP algorithm is also applicable. See Appendix E.2 for the extension.

Relationship between The Score Function and the Data Covariance Matrix

Experiments

We apply our Analytic-DPM to three pretrained score-based models provided by prior works (Ho et al., 2020; Song et al., 2020a; Nichol & Dhariwal, 2021), as well as two score-based models trained by ourselves. The pretrained score-based models are trained on CelebA 64x64 (Liu et al., 2015), ImageNet 64x64 (Deng et al., 2009) and LSUN Bedroom (Yu et al., 2015) respectively. Our score-based models are trained on CIFAR10 (Krizhevsky et al., 2009) with two different forward noise schedules: the linear schedule (LS) (Ho et al., 2020) and the cosine schedule (CS) (Nichol & Dhariwal, 2021). We denote them as CIFAR10 (LS) and CIFAR10 (CS) respectively. The number of the full timesteps $N$ is 4000 for ImageNet 64x64 and 1000 for other datasets. During sampling, we only display the mean of $p({\bm{x}}_{0}|{\bm{x}}_{1})$ and discard the noise following Ho et al. (2020), and we additionally clip the noise scale $\sigma_{2}$ of $p({\bm{x}}_{1}|{\bm{x}}_{2})$ for all methods compared in Table 2 (see details in Appendix F.2 and its ablation study in Appendix G.4). See more experimental details in Appendix F.

We conduct extensive experiments to demonstrate that analytic-DPM can consistently improve the inference efficiency of a pretrained DPM while achieving a comparable or even superior performance. Specifically, Section 6.1 and Section 6.2 present the likelihood and sample quality results respectively. Additional experiments such as ablation studies can be found in Appendix G.

Although we mainly focus on learning-free strategies of choosing the reverse variance, we also compare to another strong baseline that predicts the variance by a neural network (Nichol & Dhariwal, 2021). With full timesteps, Analytic-DPM achieves a NLL of 3.61 on ImageNet 64x64, which is very close to 3.57 reported in Nichol & Dhariwal (2021). Besides, while Nichol & Dhariwal (2021) report that the ET drastically reduces the log-likelihood performance of their neural-network-parameterized variance, Analytic-DPM performs well with the ET. See details in Appendix G.6.

2 Sample Quality

As for the sample quality, we consider the commonly used FID score (Heusel et al., 2017), where a lower value indicates a better sample quality. As shown in Table 2, under trajectories of different $K$, our Analytic-DDIM consistently improves the sample quality of the original DDIM. This allows us to generate high-quality samples with less than 50 timesteps, which results in a $20\times$ to $80\times$ speed up compared to the full timesteps. Indeed, in most cases, Analytic-DDIM only requires up to 50 timesteps to get a similar performance to the baselines. Besides, Analytic-DDPM also improves the sample quality of the original DDPM in most cases. For fairness, we use the ET implementation in Nichol & Dhariwal (2021) for all results in Table 2. We also report the results on CelebA 64x64 using a slightly different implementation of the ET following Song et al. (2020a) in Appendix G.7, and our Analytic-DPM is still effective. We show generated samples in Appendix G.9.

We observe that Analytic-DDPM does not always outperform the baseline under the FID metric, which is inconsistent with the likelihood results in Table 1. Such a behavior essentially roots in the different natures of the two metrics and has been investigated in extensive prior works (Theis et al., 2015; Ho et al., 2020; Nichol & Dhariwal, 2021; Song et al., 2021; Vahdat et al., 2021; Watson et al., 2021; Kingma et al., 2021). Similarly, using more timesteps doesn’t necessarily yield a better FID. For instance, see the Analytic-DDPM results on CIFAR10 (LS) and the DDIM results on ImageNet 64x64 in Table 2. A similar phenomenon is observed in Figure 8 in Nichol & Dhariwal (2021). Moreover, a DPM (including Analytic-DPM) with OT does not necessarily lead to a better FID score (Watson et al., 2021) (see Appendix G.5 for a comparison of ET and OT in Analytic-DPM).

We summarize the efficiency of different methods in Table 3, where we consider the least number of timesteps required to achieve a FID around 6 as the metric for a more direct comparison.

Related Work

Faster DPMs. Several works attempt to find short trajectories while maintaining the DPM performance. Chen et al. (2020) find an effective trajectory of only six timesteps by the grid search. However, the grid search is only applicable to very short trajectories due to its exponentially growing time complexity. Watson et al. (2021) model the trajectory searching as a least-cost-path problem and introduce a dynamic programming (DP) algorithm to solve this problem. Our work uses this DP algorithm, where the cost is defined as a term of the optimal KL divergence. In addition to these trajectory searching techniques, Luhman & Luhman (2021) compress the reverse denoising process into a single step model; San-Roman et al. (2021) dynamically adjust the trajectory during inference. Both of them need extra training after getting a pretrained DPM. As for DPMs with continuous timesteps (Song et al., 2020b), Song et al. (2020b) introduce an ordinary differential equation (ODE), which improves sampling efficiency and enables exact likelihood computation. However, the likelihood computation involves a stochastic trace estimator, which requires a multiple number of runs for accurate computation. Jolicoeur-Martineau et al. (2021) introduce an advanced SDE solver to simulate the reverse process in a more efficient way. However, the log-likelihood computation based on this solver is not specified.

Variance Learning in DPMs. In addition to the reverse variance, there are also works on learning the forward noise schedule (i.e., the forward variance). Kingma et al. (2021) propose variational diffusion models (VDMs) on continuous timesteps, which use a signal-to-noise ratio function to parameterize the forward variance and directly optimize the variational bound objective for a better log-likelihood. While we primarily apply our method to DDPMs and DDIMs, estimating the optimal reverse variance can also be applied to VDMs (see Appendix E).

Conclusion

We present that both the optimal reverse variance and the corresponding optimal KL divergence of a DPM have analytic forms w.r.t. its score function. Building upon it, we propose Analytic-DPM, a training-free inference framework that estimates the analytic forms of the variance and KL divergence using the Monte Carlo method and a pretrained score-based model. We derive bounds of the optimal variance to correct potential bias and reveal a relationship between the score function and the data covariance matrix. Empirically, our analytic-DPM improves both the efficiency and performance of likelihood results, and generates high-quality samples efficiently in various DPMs.

Acknowledgments

This work was supported by NSF of China Projects (Nos. 62061136001, 61620106010, 62076145), Beijing NSF Project (No. JQ19016), Beijing Outstanding Young Scientist Program NO. BJJWZYJH012019100020098, Beijing Academy of Artificial Intelligence (BAAI), Tsinghua-Huawei Joint Research Program, a grant from Tsinghua Institute for Guo Qiang, and the NVIDIA NVAIL Program with GPU/DGX Acceleration, Major Innovation & Planning Interdisciplinary Platform for the “Double-First Class” Initiative, Renmin University of China.

Ethics Statement

This work proposes an analytic estimate of the optimal variance in the reverse process of diffusion probabilistic models. As a fundamental research in machine learning, the negative consequences are not obvious. Though in theory any technique can be misused, it is not likely to happen at the current stage.

Reproducibility Statement

We provide our codes and links to pretrained models in https://github.com/baofff/Analytic-DPM. We provide details of these pretrained models in Appendix F.1. We provide details of data processing, log-likelihood evaluation, sampling and FID computation in Appendix F.2. We provide complete proofs of all theoretical results in Appendix A.

References

Appendix A Proofs and Derivations

(Cross-entropy to Gaussian) Suppose $q({\bm{x}})$ is a probability density function with mean ${\bm{\mu}}_{q}$ and covariance matrix ${\bm{\Sigma}}_{q}$ and $p({\bm{x}})={\mathcal{N}}({\bm{x}}|{\bm{\mu}},{\bm{\Sigma}})$ is a Gaussian distribution, then the cross-entropy between $q$ and $p$ is equal to the cross-entropy between ${\mathcal{N}}({\bm{x}}|{\bm{\mu}}_{q},{\bm{\Sigma}}_{q})$ and $p$, i.e.,

(KL to Gaussian) Suppose $q({\bm{x}})$ is a probability density function with mean ${\bm{\mu}}_{q}$ and covariance matrix ${\bm{\Sigma}}_{q}$ and $p({\bm{x}})={\mathcal{N}}({\bm{x}}|{\bm{\mu}},{\bm{\Sigma}})$ is a Gaussian distribution, then

where $H(\cdot)$ denotes the entropy of a distribution.

According to Lemma 1, we have $H(q,p)=H({\mathcal{N}}({\bm{x}}|{\bm{\mu}}_{q},{\bm{\Sigma}}_{q}),p)$. Thereby,

(Equivalence between the forward and reverse Markov property) Suppose $q({\bm{x}}_{0:N})=q({\bm{x}}_{0})\prod\limits_{n=1}^{N}q({\bm{x}}_{n}|{\bm{x}}_{n-1})$ is a Markov chain, then $q$ is also a Markov chain in the reverse direction, i.e., $q({\bm{x}}_{0:N})=q({\bm{x}}_{N})\prod\limits_{n=1}^{N}q({\bm{x}}_{n-1}|{\bm{x}}_{n})$.

Thereby, $q({\bm{x}}_{0:N})=q({\bm{x}}_{N})\prod\limits_{n=1}^{N}q({\bm{x}}_{n-1}|{\bm{x}}_{n})$. ∎

(Entropy of a Markov chain) Suppose $q({\bm{x}}_{0:N})$ is a Markov chain, then

(Entropy of a DDPM forward process) Suppose $q({\bm{x}}_{0:N})$ is a Markov chain and $q({\bm{x}}_{n}|{\bm{x}}_{n-1})={\mathcal{N}}({\bm{x}}_{n}|\sqrt{\alpha_{n}}{\bm{x}}_{n-1},\beta_{n}{\bm{I}})$, then

(Entropy of a conditional Markov chain) Suppose $q({\bm{x}}_{1:N}|{\bm{x}}_{0})$ is Markov, then

(Entropy of a generalized DDPM forward process) Suppose $q({\bm{x}}_{1:N}|{\bm{x}}_{0})$ is Markov, $q({\bm{x}}_{N}|{\bm{x}}_{0})$ is Gaussian with covariance $\overline{\beta}_{N}{\bm{I}}$ and $q({\bm{x}}_{n-1}|{\bm{x}}_{n},{\bm{x}}_{0})$ is Gaussian with covariance $\lambda_{n}^{2}{\bm{I}}$, then

(KL to a Markov chain) Suppose $q({\bm{x}}_{0:N})$ is a probability distribution and $p({\bm{x}}_{0:N})=p({\bm{x}}_{N})\prod\limits_{n=1}^{N}p({\bm{x}}_{n-1}|{\bm{x}}_{n})$ is a Markov chain, then we have

If $q({\bm{x}}_{0:N})$ is also a Markov chain, according to Lemma 4, we have $c=0$. ∎

(The optimal Markov reverse process with Gaussian transitions is equivalent to moment matching) Suppose $q({\bm{x}}_{0:N})$ is probability density function and $p({\bm{x}}_{0:N})=\prod\limits_{n=1}^{N}p({\bm{x}}_{n-1}|{\bm{x}}_{n})p({\bm{x}}_{N})$ is a Gaussian Markov chain with $p({\bm{x}}_{n-1}|{\bm{x}}_{n})={\mathcal{N}}({\bm{x}}_{n-1}|{\bm{\mu}}_{n}({\bm{x}}_{n}),\sigma_{n}^{2}{\bm{I}})$, then the joint KL optimization

which match the first two moments of $q({\bm{x}}_{n-1}|{\bm{x}}_{n})$. The corresponding optimal KL is

Lemma 9 doesn’t assume the form of $q({\bm{x}}_{0:N})$, thereby it can be applied to more general Gaussian models, such as multi-layer VAEs with Gaussian decoders (Rezende et al., 2014; Burda et al., 2015). In this case, $q({\bm{x}}_{1:N}|{\bm{x}}_{0})$ is the hierarchical encoders of multi-layer VAEs.

In the optimal case, ${\mathcal{F}}(\sigma_{n}^{*2})=\frac{d}{2}(1+\log\sigma_{n}^{*2})$ and

(Marginal score function) Suppose $q({\bm{v}},{\bm{w}})$ is a probability distribution, then

(Score representation of conditional expectation and covariance) Suppose $q({\bm{v}},{\bm{w}})=q({\bm{v}})q({\bm{w}}|{\bm{v}})$, where $q({\bm{w}}|{\bm{v}})={\mathcal{N}}({\bm{w}}|\sqrt{\alpha}{\bm{v}},\beta{\bm{I}})$, then

(Convert the moments of $q({\bm{x}}_{n-1}|{\bm{x}}_{n})$ to moments of $q({\bm{x}}_{0}|{\bm{x}}_{n})$) The optimal solution ${\bm{\mu}}_{n}^{*}({\bm{x}}_{n})$ and $\sigma_{n}^{*2}$ to Eq. (4) can be represented by the first two moments of $q({\bm{x}}_{0}|{\bm{x}}_{n})$

where $q_{n}({\bm{x}}_{n})$ is the marginal distribution of the forward process at timestep $n$ and $d$ is the dimension of the data.

According to Lemma 9, the optimal ${\bm{\mu}}_{n}^{*}$ and $\sigma_{n}^{*2}$ under KL minimization is

Then we consider $\sigma_{n}^{*2}$. According to the law of total variance, we have

A.2 Proof of Theorem 1

According to Lemma 11 and Lemma 13, we have

A.3 Proof of Theorem 2

According to Lemma 13 and Theorem 1, we have

If we further $q({\bm{x}}_{0})$ assume is a bounded distribution in $[a,b]^{d}$, then $q({\bm{x}}_{0}|{\bm{x}}_{n})$ is also a bounded distribution in $[a,b]^{d}$. According to Lemma 12, we have

A.4 Proof of the Decomposed Optimal KL

(Decomposed optimal KL, proof in Appendix A.4)

The KL divergence between the shorter forward process and its optimal reverse process is

where $J(\tau_{k-1},\tau_{k})=\log\frac{\sigma_{\tau_{k-1}|\tau_{k}}^{*2}}{\lambda_{\tau_{k-1}|\tau_{k}}^{2}}$ and $c$ is a constant unrelated to the trajectory $\tau$.

According to Lemma 7 and Lemma 9, we have

A.5 The Formal Result for Section 5 and its Proof

Here we present the formal result of the relationship between the score function and the data covariance matrix mentioned in Section 5.

(Proof in Appendix A.5) The expected conditional covariance matrix of the data distribution is determined by the score function $\nabla_{{\bm{x}}_{n}}\log q_{n}({\bm{x}}_{n})$ as follows:

which contributes to the data covariance matrix according to the law of total variance

Since $q({\bm{x}}_{n}|{\bm{x}}_{0})={\mathcal{N}}({\bm{x}}_{n}|\sqrt{\overline{\alpha}_{n}}{\bm{x}}_{0},\overline{\beta}_{n}{\bm{I}})$, according to Lemma 11, we have

The law of total variance is a classical result in statistics. Here we prove it for completeness:

Appendix B The DP Algorithm for the Least-Cost-Path Problem

Given a cost function $J(s,t)$ with $1\leq s<t$ and $k,n\geq 1$, we want to find a trajectory $1=\tau_{1}<\cdots<\tau_{k}=n$ of $k$ nodes starting from $1$ and terminating at $n$, s.t., the total cost $J(\tau_{1},\tau_{2})+J(\tau_{2},\tau_{3})+\cdots+J(\tau_{k-1},\tau_{k})$ is minimized. Such a problem can be solved by the DP algorithm proposed by Watson et al. (2021). Let $C[k,n]$ be the minimized cost of the optimal trajectory, and $D[k,n]$ be the $\tau_{k-1}$ of the optimal trajectory. For simplicity, we also let $J(s,t)=\infty$ for $s\geq t\geq 1$. Then for $k=1$, we have C[1,n]=\left\{\begin{array}[]{ll}0&n=1\\ \infty&N\geq n>1\end{array}\right. and $D[1,n]=-1$ (here $\infty$ and $-1$ represent undefined values for simplicity). For $N\geq k\geq 2$, we have

As long as $D$ is calculated, we can get the optimal trajectory recursively by $\tau_{K}=N$ and $\tau_{k-1}=D[k,\tau_{k}]$. We summarize the algorithm in Algorithm 1.

Appendix C The Bounds of the Optimal Reverse Variance Constrained on a Trajectory

In Section 4, we derive the optimal reverse variance constrained on a trajectory. Indeed, the optimal reverse variance can also be bounded similar to Theorem 2. We formalize it in Corollary 1.

(Bounds of the optimal reverse variance constrained on a trajectory)

$\sigma_{\tau_{k-1}|\tau_{k}}^{*2}$ has the following lower and upper bounds:

If we further assume $q({\bm{x}}_{0})$ is a bounded distribution in $[a,b]^{d}$, where $d$ is the dimension of data, then $\sigma_{n}^{*2}$ can be further upper bounded by

Appendix D Simplified Results for the DDPM Forward Process

(Simplified score representation of the optimal solution)

(Simplified bounds of the optimal reverse variance)

If we further assume $q({\bm{x}}_{0})$ is a bounded distribution in $[a,b]^{d}$, where $d$ is the dimension of data, then $\sigma_{n}^{*2}$ can be further upper bounded by

Besides, Theorem 3 can also be simplified for DDPMs. We list the simplified result in Corollary 4.

and $c$ is a constant unrelated to the trajectory $\tau$.

Appendix E Extension to Diffusion Process with Continuous Timesteps

where $\overline{\alpha}_{t}$ and $\overline{\beta}_{t}$ are scalar-valued functions satisfying some regular conditions (Kingma et al., 2021) with domain $t\in$. Such a parameterization induces a diffusion process on continuous timesteps $q({\bm{x}}_{0},{\bm{z}}_{})$, s.t.,

where $\alpha_{t|s}\coloneqq\overline{\alpha}_{t}/\overline{\alpha}_{s}$ and $\beta_{t|s}\coloneqq\overline{\beta}_{t}-\alpha_{t|s}\overline{\beta}_{s}$.

Kingma et al. (2021) introduce $p({\bm{z}}_{s}|{\bm{z}}_{t})={\mathcal{N}}({\bm{z}}_{s}|{\bm{\mu}}_{s|t}({\bm{z}}_{t}),\sigma_{s|t}^{2})$ ($s<t$) to reverse from timestep $t$ to timestep $s$, where $\sigma_{s|t}^{2}$ is fixed to $\frac{\overline{\beta}_{s}}{\overline{\beta}_{t}}\beta_{s|t}$. In contrast, we show that $\sigma_{s|t}^{2}$ also has an optimal solution in an analytic form of the score function under the sense of KL minimization. According to Lemma 9 and Lemma 11, we have

Thereby, both the optimal mean and variance have a closed form expression w.r.t. the score function. In this case, we first estimate the expected mean squared norm of the score function by $\Gamma_{t}$ for $t\in$, where

Notice that there are infinite timesteps in $$. In practice, we can only choose a finite number of timesteps$0=t_{1}<\cdots$and calculate$\Gamma_{t_{n}}$. For a timestep$t$between$t_{n-1}$and$t_{n}$, we can use a linear interpolation between$\Gamma_{t_{n-1}}$and$\Gamma_{t_{n}}$. Then, we can estimate$\sigma_{s|t}^{*2}$ by

E.2 Analytic Estimation of the Optimal Reverse Trajectory

Now we consider optimize the trajectory $0=\tau_{1}<\cdots<\tau_{K}=1$ in the sense of KL minimization

Appendix F Experimental Details

The CelebA 64x64 pretrained score-based model is provided in the official code (https://github.com/ermongroup/ddim) of Song et al. (2020a). The LSUN Bedroom pretrained score-based model is provided in the official code (https://github.com/hojonathanho/diffusion) of Ho et al. (2020). Both of them have a total of $N=1000$ timesteps and use the linear schedule (Ho et al., 2020) as the forward noise schedule.

The CIFAR10 score-based models are trained by ourselves. They have a total of $N=1000$ timesteps and are trained with the linear forward noise schedule and the cosine forward noise schedule respectively. We use the same U-Net model architecture to Nichol & Dhariwal (2021). Following Nichol & Dhariwal (2021), we train 500K iterations with a batch size of 128, use a learning rate of 0.0001 with the AdamW optimizer (Loshchilov & Hutter, 2017) and use an exponential moving average (EMA) with a rate of 0.9999. We save a checkpoint every 10K iterations and select the checkpoint according to the FID results on 1000 samples generated under the reverse variance $\sigma_{n}^{2}=\beta_{n}$ and full timesteps.

F.2 Log-likelihood and Sampling

Following Ho et al. (2020), we linearly scale the image data consisting of integers in $\{0,1,\cdots,255\}$ to $$, and discretize the last reverse Markov transition$p({\bm{x}}_{0}|{\bm{x}}_{1})$ to obtain discrete log-likelihoods for image data.

We use the official implementation of FID to pytorch (https://github.com/mseitzer/pytorch-fid). We calculate the FID score on 50K generated samples on all datasets. Following Nichol & Dhariwal (2021), the reference distribution statistics are computed on the full training set for CIFAR10 and ImageNet 64x64. For CelebA 64x64 and LSUN Bedroom, the reference distribution statistics is computed on 50K training samples.

F.3 Choice of the Number of Monte Carlo Samples and Calculation of ΓΓ\Gamma

We use a maximal $M$ without introducing too much computation. Specifically, we set $M=50000$ on CIFAR10, $M=10000$ on CelebA 64x64 and ImageNet 64x64 and $M=1000$ on LSUN Bedroom by default without a sweep. All of the samples are from the training dataset. We use the default settings of $M$ for all results in Table 1, Table 2 and Table 3.

We only calculate $\Gamma$ in Eq. (8) once for a pretrained model, and $\Gamma$ is reused during inference under different settings (e.g., trajectories of smaller $K$) in Table 1, Table 2 and Table 3.

F.4 Implementation of the Even Trajectory

F.5 Experimental Details of Table 3

The DDPM and DDIM results on CIFAR10 are based on the quadratic trajectory following Song et al. (2020a), which gets better FID than the even trajectory. The Analytic-DPM result is based on the DDPM forward process on LSUN Bedroom, and based on the DDIM forward process on other datasets. These choices achieve better efficiency than their alternatives.

Appendix G Additional Results

G.2 Ablation Study on the Number of Monte Carlo Samples

We show that only a small number of Monte Carlo (MC) samples $M$ in Eq. (8) is enough for a small MC variance. As shown in Figure 3, the values of $\Gamma_{n}$ with $M=100$ and $M=50000$ Monte Carlo samples are almost the same in a single trial. To explicitly see the variance, in Figure 4 and Figure 5, we plot the mean, the standard deviation and the relative standard deviation (RSD) (i.e., the ratio of the standard deviation to the mean) of a single Monte Carlo sample $\frac{||{\bm{s}}_{n}({\bm{x}}_{n})||^{2}}{d}$, ${\bm{x}}_{n}\sim q_{n}({\bm{x}}_{n})$ and $\Gamma_{n}$ with different $M$ respectively on CIFAR10 (LS). In all cases, the RSD decays fast as $n$ increases. When $n$ is small (e.g., $n=1$), using $M=10$ Monte Carlo samples can ensure that the RSD of $\Gamma_{n}$ is below 0.1, and using $M=100$ Monte Carlo samples can ensure that the RSD of $\Gamma_{n}$ is about 0.025. When $n>100$, the RSD of a single Monte Carlo sample is below 0.05, and using only $M=1$ Monte Carlo sample can ensure the RSD of $\Gamma_{n}$ is below 0.05. Overall, a small $M$ like 10 and 100 is sufficient for a small Monte Carlo variance.

Furthermore, we show that Analytic-DPM with a small $M$ like 10 and 100 has a similar performance to that with a large $M$. As shown in Figure 6 (a), using $M=100$ or $M=50000$ almost does not affect the likelihood results on CIFAR10 (LS). In Table 5 (a), we show results with even smaller $M$ (e.g., $M=1,3,10$). Under both the NLL and FID metrics, $M=10$ achieves a similar result to that of $M=50000$. The results are similar on ImageNet 64x64, as shown in Figure 6 (b) and Table 5 (b). Notably, the expected performance of FID is almost not influenced by the choice of $M$.

As a result, Analytic-DPM consistently improves the baselines using a much smaller $M$ (e.g., $M=10$), as shown in Table 6.

G.3 Tightness of the Bounds

In Section 3.1 and Appendix C, we derive upper and lower bounds of the optimal reverse variance. In this section, we show these bounds are tight numerically in practice. In Figure 7, we plot the combined upper bound (i.e., the minimum of the upper bounds in Eq. (11) and Eq. (12)) and the lower bound on CIFAR10. As shown in Figure 7 (a,c), the two bounds almost overlap under the full-timesteps ($K$=$N$) trajectory. When the trajectory has a smaller number of timesteps (e.g., $K$=100), the two bounds also overlap when the timestep $\tau_{k}$ is large. These results empirically validate that our bounds are tight, especially when the timestep is large.

In Figure 8, we also plot the two upper bounds in Eq. (11) and Eq. (12) individually. The upper bound in Eq. (11) is tighter when the timestep is small and the other one is tighter when the timestep is large. Thereby, both upper bounds contribute to the combined upper bound.

To see how these bounds work in practice, in Figure 9, we plot the probability that $\hat{\sigma}_{n}^{2}$ is clipped by the bounds in Theorem 2 with different number of Monte Carlo samples $M$ on CIFAR10 (LS). For all $M$, the curves of ratio v.s. $n$ are similar and the estimate is clipped more frequently when $n$ is large. This is as expected because when $n$ is large, the gap between the upper bound in Eq. (12) and the lower bound in Eq. (11) tends to zero. The results also agree with the plot of the bounds in Figure 7. Besides, the similarity of results between different $M$ implies that the clipping by bounds occurs mainly due to the error of the score-based model ${\bm{s}}_{n}({\bm{x}}_{n})$, instead of the randomness in Monte Carlo methods.

This section validates the argument in Appendix F.2 that properly clipping the noise scale $\sigma_{2}$ in $p({\bm{x}}_{1}|{\bm{x}}_{2})$ leads to a better sample quality. As shown in Figure 10 and Figure 11, it greatly improves the sample quality of our analytic estimate. The curves of clipping and no clipping overlap as $K$ increases, since $\sigma_{2}$ is below the threshold for a large $K$.

Indeed, as shown in Table 7, the clipping threshold designed for sampling in Appendix F.2 is 1 to 3 orders of magnitude smaller than the combined upper bound in Theorem 2 (i.e., the minimum of the upper bounds in Eq. (11) and Eq. (12)) when $K$ is small.

G.5 Sample Quality Comparison between Different Trajectories

While the optimal trajectory (OT) significantly improves the likelihood results, it doesn’t lead to better FID results. As shown in Figure 13, the even trajectory (ET) has better FID results. Such a behavior essentially roots in the different natures of the two metrics and has been investigated in extensive prior works (Ho et al., 2020; Nichol & Dhariwal, 2021; Song et al., 2021; Vahdat et al., 2021; Watson et al., 2021; Kingma et al., 2021).

G.6 Additional Likelihood Comparison

We compare our Analytic-DPM to Improved DDPM (Nichol & Dhariwal, 2021) that predicts the reverse variance by a neural network. The comparison is based on the ImageNet 64x64 model described in Appendix F.1. As shown in Table 8, with full timesteps, Analytic-DPM achieves a NLL of 3.61, which is very close to 3.57 achieved by predicting the reverse variance in Improved DDPM. Besides, we also notice that the ET reduces the log-likelihood performance of Improved DDPM when $K$ is small, and this is consistent with what Nichol & Dhariwal (2021) report. In contrast, our Analytic-DPM performs well with the ET.

G.7 CelebA 64x64 Results with a Slightly Different Implementation of the Even Trajectory

G.8 Comparison to other Classes of Generative Models

While DPMs and their variants serve as the most direct baselines to validate the effectiveness of our method, we also compare with other classes of generative models in Table 10. Analytic-DPM achieves competitive sample quality results among various generative models, and meanwhile significantly reduces the efficiency gap between DPMs and other models.

G.9 Samples

In Figure 14-17, we show Analytic-DDIM constrained on a short trajectory of $K=50$ timesteps can generate samples comparable to these under the best FID setting.

In Figure 19-21, we also show samples of both Analytic-DDPM and Analytic-DDIM constrained on trajectories of different number of timesteps $K$.

Appendix H Additional Discussion

The extra cost of the Monte Carlo estimate $\Gamma$ is small compared to the whole inference cost. In fact, the Monte Carlo estimate requires $MN$ additional model function evaluations. During inference, suppose we generate $M_{1}$ samples or calculate the log-likelihood of $M_{1}$ samples with $K$ timesteps. Both DPMs and Analytic-DPMs need $M_{1}K$ model function evaluations. Employing the same score-based models, the relative additional cost of Analytic-DPM is $\frac{MN}{M_{1}K}$. As shown in Appendix G.2, a very small $M$ (e.g., $M=10,100$) is sufficient for Analytic-DPM, making the relative additional cost small if not negligible. For instance, on CIFAR10, let $M=10$, $N=1000$, $M_{1}=50000$ and $K\geq 10$, we obtain $\frac{MN}{M_{1}K}\leq 0.02$ and Analytic-DPM still consistently improves the baselines as presented in Table 6.

Further, the additional calculation of the Monte Carlo estimate occurs only once given a pretrained model and training dataset, since we can save the results of $\Gamma=(\Gamma_{1},\cdots,\Gamma_{N})$ in Eq.(8) and reuse it among different inference settings (e.g., trajectories of various $K$). The reuse is valid, because the marginal distribution of a shorter forward process $q({\bm{x}}_{0},{\bm{x}}_{\tau_{1}},\cdots,{\bm{x}}_{\tau_{K}})$ at timestep $\tau_{k}$ is the same as that of the full-timesteps forward process $q({\bm{x}}_{0:N})$ at timestep $n=\tau_{k}$. Indeed, in our experiments (e.g., Table 1,2), $\Gamma$ is shared across different selections of $K$, trajectories and forward processes. Moreover, in practice, $\Gamma$ can be calculated offline and deployed together with the pretrained model and the online inference cost of Analytic-DPM is exactly the same as DPM.

H.2 The Stochasticity of the Variational Bound after Plugging the Analytic Estimate

we only need to study the convexity of $L_{n}$ w.r.t. $\sigma_{n}^{2}$.

H.3 Comparison to other Gaussian models and their results

The reverse process of DPMs is a Markov process with Gaussian transitions. Thereby, it is interesting to compare it with other Gaussian models, e.g., the expectation propagation (EP) with the Gaussian process (GP) (Kim & Ghahramani, 2006).

In EP with GP (Kim & Ghahramani, 2006), $p_{target}$ is the product of a single likelihood factor and all other approximate factors for tractability. In fact, the form of the likelihood factor is chosen such that the first two moments of $p_{target}$ can be easily computed or approximated. For instance, the original EP (Minka, 2001) considers Gaussian mixture likelihood (or Bernoulli likelihood for classification) and the moments can be directed obtained by the properties of Gaussian (or integration by parts). Besides, at the cost of the tractability, there is no converge guarantee of EP in general.

In contrast, $p_{target}$ in this paper is the conditional distribution $q({\bm{x}}_{n-1}|{\bm{x}}_{n})$ of the corresponding joint distribution $q({\bm{x}}_{0:N})$ defined by the forward process. Note that the moments of $q({\bm{x}}_{n-1}|{\bm{x}}_{n})$ are nontrivial to calculate because it involves an unknown and potentially complicated data distribution. Technically, in Lemma 13, we carefully use the law of total variance conditioned on ${\bm{x}}_{0}$ and convert the second moment of $q({\bm{x}}_{n-1}|{\bm{x}}_{n})$ to that of $q({\bm{x}}_{0}|{\bm{x}}_{n})$, which surprisingly can be expressed as the score function as proven in Lemma 11.

H.4 Future Works

In our work, we mainly focus on image data. It would be interesting to apply Analytic-DPM to other data modalities, e.g. speech data (Chen et al., 2020). As presented in Appendix E, our method can be applied to continuous DPMs, e.g., variational diffusion models (Kingma et al., 2021) that learn the forward noise schedule. It is appealing to see how Analytic-DPM works on these continuous DPMs. Finally, it is also interesting to incorporate the optimal reverse variance in the training process of DPMs.