Estimating the Optimal Covariance with Imperfect Mean in Diffusion Probabilistic Models

Introduction

Recently, diffusion probabilistic models (DPMs) (Sohl-Dickstein et al., 2015; Ho et al., 2020; Song et al., 2020b) have shown great promise for generative modeling. Such models smoothly inject noise to the data distribution, which forms a diffusion process. By learning to reverse the process using a Markov model, DPMs are able to generate high quality images (Ho et al., 2020; Song et al., 2020b; Dhariwal & Nichol, 2021) and audios (Chen et al., 2020; Kong et al., 2020), which are comparable or even superior than the current state-of-the-art generative models (Brock et al., 2018; Wu et al., 2019; Karras et al., 2020; Bińkowski et al., 2019; Kalchbrenner et al., 2018).

However, the iterative generation over the full timesteps of DPMs makes them much less efficient than generative models such as GANs (Goodfellow et al., 2014). Thus, the generation performance on a subset of timesteps is crucial. In this case, the transition of the reversed diffusion process becomes more complex (Xiao et al., 2021), and the covariance design in DPMs matters (Nichol & Dhariwal, 2021; Bao et al., 2022). Most prior works (Ho et al., 2020; Song et al., 2020a; Bao et al., 2022) use an isotropic covariance that only depends on the timestep without considering the state. A notable recent progress is Analytic-DPM (Bao et al., 2022), which estimates the optimal isotropic covariance (in the sense of maximum likelihood) instead of using handcrafted values (Ho et al., 2020; Song et al., 2020a) and shows significant improvement on likelihood estimation and sampling efficiency. However, the isotropic covariance sacrifices the expressive power of DPMs for simplicity. Moreover, the optimality of the covariance estimate in Bao et al. (2022) holds by assuming that the optimal mean is known, which is not the case in practice.

To overcome the aforementioned limitations and further improve DPMs for likelihood estimation and sampling efficiency, we consider diagonal and full covariances to improve the expressive power of DPMs. We derive the optimal mean and covariance from the perspective of the maximum likelihood. We also correct the optimal covariance given an imperfect mean (i.e., considering the approximation and optimization errors) in terms of maximum likelihood. Both the optimal and the corrected ones can be decomposed into terms of conditional expectations over functions of noise, which can be estimated by minimizing mean squared error (MSE) losses. Although our theory applies to the full covariance case, in our implementation, we consider the diagonal covariance for computational efficiency. Besides, we adopt a parameter sharing scheme for inference efficiency and a two-stage training process motivated by our theory.

Our method is applicable to DPMs with both discrete (Ho et al., 2020) and continuous (Song et al., 2020b) timesteps. In our experiments, we directly compare our method with a variety of baselines (Ho et al., 2020; Song et al., 2020a; Bao et al., 2022; Song et al., 2020b) in terms of sample quality and likelihood estimation in DPMs with both discrete and continuous timesteps. At nearly the same computation cost, our method consistently outperforms these baselines on likelihood estimation. Besides, our method also outperforms these baselines on the sample quality in most cases, especially when the number of timesteps is small.

Background

Diffusion probabilistic models (DPMs) are special Markov models with Gaussian transitions:

where $\overline{\alpha}_{1},\overline{\alpha}_{2},\cdots,\overline{\alpha}_{N}\in(0,1)$ is a strictly decreasing sequence, $\overline{\beta}_{n}\coloneqq 1-\overline{\alpha}_{n}$ and ${\bm{I}}$ is the identity matrix. Observing that $q({\bm{x}}_{n}|{\bm{x}}_{0})={\mathcal{N}}({\bm{x}}_{n}|\sqrt{\overline{\alpha}_{n}}{\bm{x}}_{0},\overline{\beta}_{n}{\bm{I}})$, we can quickly sample ${\bm{x}}_{n}$ given ${\bm{x}}_{0}$ by

Note that problem (3) optimizes the mean and the covariance jointly. Thus, we term it joint optimization.

where ${\bm{\epsilon}}_{n}=\frac{{\bm{x}}_{n}-\sqrt{\overline{\alpha}_{n}}{\bm{x}}_{0}}{\sqrt{\overline{\beta}_{n}}}$ is the noise used to generate ${\bm{x}}_{n}$ from ${\bm{x}}_{0}$, $d$ is the dimension of the data ${\bm{x}}_{0}$ and $\gamma_{n}=\sqrt{\overline{\alpha}_{n-1}}-\sqrt{\overline{\beta}_{n-1}-\lambda_{n}^{2}}\sqrt{\frac{\overline{\alpha}_{n}}{\overline{\beta}_{n}}}$. Ho et al. (2020) estimate ${\bm{\mu}}_{n}^{*}({\bm{x}}_{n})$ by a noise prediction network $\hat{{\bm{\epsilon}}}_{n}({\bm{x}}_{n})$

The optimal variance in Eq. (5) can also be estimated using $\hat{{\bm{\epsilon}}}_{n}({\bm{x}}_{n})$ as following (Bao et al., 2022)

Method

We improve the covariance estimate of DPMs building upon the recent progress (Bao et al., 2022). First, we consider diagonal and full covariances instead of using an isotropic one to improve the expressive power of DPMs and obtain the expression of the optimal solution, as detailed in Section 3.1. Second, we correct the optimal covariance given an imperfect mean (i.e., with approximation and optimization errors), as presented in Section 3.2. For clarity, we refer the readers to Appendix A for all proofs.

Although our theoretical results are applicable to the full covariance case (see Appendix B), it is often time consuming to obtain samples from Gaussian transitions with full covariances (e.g., via Cholesky decomposition). To balance the flexibility and the time cost, we focus on diagonal covariances throughout the main text including the experiments.

Instead of using an isotropic covariance that only depends on the timestep without considering the state, we consider a diagonal covariance to improve the expressive power of DPMs. To be clear and rigorous, we first construct an example where the data distribution is a mixture of Gaussian. Even in such a simple case, we can prove that the optimal ELBO (in Eq. (3)) with a diagonal covariance is strictly greater than that with an isotropic covariance. See a formal description and a detailed proof in Proposition A.5.

where ${\bm{\epsilon}}_{n}=\frac{{\bm{x}}_{n}-\sqrt{\overline{\alpha}_{n}}{\bm{x}}_{0}}{\sqrt{\overline{\beta}_{n}}}$ is the noise used to generate ${\bm{x}}_{n}$ from ${\bm{x}}_{0}$, $(\cdot)^{2}$ is the element-wise square, ${\bm{1}}$ is the vector of ones and $\gamma_{n}=\sqrt{\overline{\alpha}_{n-1}}-\sqrt{\overline{\beta}_{n-1}-\lambda_{n}^{2}}\sqrt{\frac{\overline{\alpha}_{n}}{\overline{\beta}_{n}}}$.

The proof idea is similar to Theorem 1 in Bao et al. (2022), which only considers the isotropic covariance case.

In the paper, we refer to the above estimate as SN-DPM.

2 The Optimal Covariance with an Imperfect Mean

It is plausible to obtain the estimates $\hat{{\bm{\mu}}}_{n}({\bm{x}}_{n})$ and $\hat{{\bm{\sigma}}}_{n}({\bm{x}}_{n})^{2}$, according to Theorem 3.1. However, we argue that the exact optimal solution ${\bm{\mu}}^{*}_{n}({\bm{x}}_{n}),{\bm{\sigma}}^{*}_{n}({\bm{x}}_{n})^{2}$ in Theorem 3.1 cannot be achieved due to the approximation and optimization errors, making it possible to further improve the estimates considering the errors. In fact, below we present how to correct the optimal covariance given a potentially imperfect mean $\hat{{\bm{\mu}}}_{n}$.

We note that, in Theorem 3.1, the optimal mean is irrelevant to the covariance, while the optimal covariance is expressed by the optimal mean. Naturally, we can first optimize the mean solely to obtain the mean estimate $\hat{{\bm{\mu}}}_{n}$ and then optimize the covariance solely given $\hat{{\bm{\mu}}}_{n}$. Such a two-stage approach is at least not worse than the estimate in Section 3.1 under the maximum likelihood, because we have

where the equality does not hold in general.

To formalize the idea, we obtain the optimal solutions to the following two problems

as summarized in Theorem 3.2 and Theorem 3.3.

(Optimal solution to optimization solely w.r.t. mean) For any covariance ${\bm{\Sigma}}_{n}$, the optimal mean to problem (12) is always Eq. (4), i.e.,

which is irrelevant to ${\bm{\Sigma}}_{n}$.

Theorem 3.2 reveals that the optimal mean doesn’t depend on the covariance and therefore we can learn the mean on Eq. (7) without knowing information about the covariance.

In contrast to the mean, the optimal covariance to problem (13) depends on the mean, as shown in Theorem 3.3. See the result for full covariances in Appendix B.

where ${\bm{\sigma}}_{n}^{*}({\bm{x}}_{n})^{2}$ is the optimal covariance to the joint optimization problem in Theorem 3.1, ${\bm{\epsilon}}_{n}=\frac{{\bm{x}}_{n}-\sqrt{\overline{\alpha}_{n}}{\bm{x}}_{0}}{\sqrt{\overline{\beta}_{n}}}$ is the noise used to generate ${\bm{x}}_{n}$ from ${\bm{x}}_{0}$, $(\cdot)^{2}$ is the element-wise square, ${\bm{1}}$ is the vector of ones and $\gamma_{n}=\sqrt{\overline{\alpha}_{n-1}}-\sqrt{\overline{\beta}_{n-1}-\lambda_{n}^{2}}\sqrt{\frac{\overline{\alpha}_{n}}{\overline{\beta}_{n}}}$.

The proof of Theorem 3.3 mainly builds on this fact: minimizing the KL divergence between a target density and a Gaussian density conditioned on a fixed mean is equivalent to firstly obtaining the second moment of the target density, and then correcting it by the difference between the first moments of the two densities. We term it conditioned moment matching (see Lemma A.1).

In the paper, we refer to the above estimate as NPR-DPM.

Implementation

We provide more details of implementing SN-DPM and NPR-DPM. Specifically, we design a parameter sharing scheme for inference efficiency in Section 4.1. Motivated by Theorem 3.2 & 3.3, we employ a two-stage training process for both SN-DPM and NPR-DPM in Section 4.2.

During inference, both SN-DPM and NPR-DPM need to evaluate two networks. To reduce the computation cost, we let the two neural networks share parameters as follows

In comparison to original DPMs with the same UNet, both SN-DPM and NPR-DPM require negligible additional memory cost and at most $10\%$ more computation time (see details in Appendix F.3).

2 Two-Stage Learning and Pretraining

Recall that the optimal mean to problem (12) is irrelevant to the covariance. Thus, the learning of the mean and the covariance can naturally be divided into two stages. In the first stage, we learn the mean by training a noise prediction network on Eq. (7), or just use a pretrained one from prior works. In the second stage, we fix the parameter of the noise prediction network, take the UNet as a pretrained model, and only learn the covariance by tuning the parameter ${\bm{\phi}}_{2}$ (in Eq. (17)) of the NPR (or SN) prediction network using the gradient of Eq. (15) (or Eq. (9)):

Extension to DPMs with Continuous Timesteps

where $q({\bm{x}}_{t})$ is the marginal distribution of ${\bm{x}}_{t}$. Similarly to the optimal solution in Eq. (4) & (5), Bao et al. (2022) derive the optimal mean

Then, we extend Theorem 3.3 to DPMs with continuous timesteps. Similarly to problem (12) & (13), we consider the following two optimizations w.r.t. the mean or the covariance solely:

The optimal mean to problem (21) is also Eq. (19), which is irrelevant to ${\bm{\Sigma}}_{s|t}$. The optimal covariance to problem (22) is given in Proposition 5.2.

Experiments

We evaluate SN-DPM and NPR-DPM in DPMs with both discrete and continuous timestepsWe provide our implementation at https://github.com/baofff/Extended-Analytic-DPM..

As for DPMs with continuous timesteps (see Section 5), we consider the VP SDE (Song et al., 2020b) as the forward process. We compare our method with the following baselines: (i) the Euler-Maruyama solver (Song et al., 2020b), which firstly reverses the VP SDE and then discretizes the reverse one using the Euler-Maruyama solver, (ii) the ancestral sampling (Song et al., 2020b), which designs a Markov chain similar to Eq. (1) and samples from the Markov chain, (iii) the probability flow (Song et al., 2020b), which firstly derives an equivalent ODE to the SDE and then discretizes the ODE, (iv) the Analytic-DPM (Bao et al., 2022), as introduced in Section 2, and (v) the “Gotta Go Fast” SDE solver (Jolicoeur-Martineau et al., 2021), which uses adaptive step sizes.

Since the performance on a subset of timesteps is crucial, we compare our method and baselines constrained on trajectories $1\leq\tau_{1}<\cdots<\tau_{K}=N$ with different number of timesteps $K$ (Song et al., 2020a; Bao et al., 2022) (see Appendix D for details of how to apply our method to trajectories). Following Bao et al. (2022), we consider two kinds of trajectories. The first one is the even trajectory (ET) (Nichol & Dhariwal, 2021), where the timesteps are evenly spaced. The second one is the optimal trajectory (OT) (Watson et al., 2021), where the timesteps are determined by dynamic programming that maximizes ELBO.

We evaluate our method on six pretrained noise prediction networks provided by prior works (Ho et al., 2020; Song et al., 2020b, a; Nichol & Dhariwal, 2021; Bao et al., 2022). Three of them are trained on CIFAR10 (Krizhevsky et al., 2009) with the linear schedule (LS) of $\beta_{n}$ (Ho et al., 2020), the cosine schedule (CS) of $\beta_{n}$ (Nichol & Dhariwal, 2021) and the VP SDE (Song et al., 2020b) respectively. We refer to the first two settings as CIFAR10 (LS) and CIFAR10 (CS) respectively, which have discrete timesteps, and refer to the last setting as CIFAR10 (VP SDE), which has continuous timesteps. The others are trained with discrete timesteps on CelebA 64x64 (Liu et al., 2015), ImageNet 64x64 (Deng et al., 2009) and LSUN Bedroom (Yu et al., 2015) respectively. We train NPR (or SN) prediction networks for all of them, following the implementation in Section 4. See more experimental details in Appendix F.

In this part, we compare the sample quality quantitatively, as measured by the commonly used FID score (Heusel et al., 2017). We evaluate DPMs with both the discrete and continuous timesteps. As for DPMs with discrete timesteps, we report results under both the DDPM and DDIM forward processes. Since Watson et al. (2021); Bao et al. (2022) report that a DPM with the optimal trajectory is likely to have a worse FID score, we only include results under the even trajectory.

As shown in Table 2, both NPR-DPM and SN-DPM outperform baselines in most cases. In particular, SN-DPM improves the sample quality remarkably on CIFAR10 and CelebA 64x64 when the number of timesteps is small. Moreover, our methods perform well in DPMs with both discrete and continuous timesteps.

Following Bao et al. (2022), we also compare the least number of timesteps required to achieve a FID around 6, as shown in Table 3. Remarkably, on CelebA 64x64, SN-DPM requires less than half the number of timesteps compared to baselines. In Appendix G.3, we show FID comparison after normalizing for the extra cost, as well as the comparison with the “Gotta Go Fast” SDE solver. In Appendix G.4, we show generated samples of our methods.

2 Likelihood Results

As shown in Table 4, our NPR-DPM consistently outperforms all baselines on all datasets, all number of timesteps and both kinds of trajectories.

In initial experiments, we find that SN-DPM doesn’t perform well on the likelihood results, potentially due to not considering the imperfect mean and the amplified error of $\hat{{\bm{\epsilon}}}_{n}({\bm{x}}_{n})^{2}$ as mentioned in Section 3.1. The inconsistency with its FID results in Table 2 roots in the different natures of the two metrics, which has been observed and explored extensively in 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; Bao et al., 2022). We also note that while NPR-DPM performs better than SN-DPM w.r.t. likelihood, it does not perform better w.r.t. FID. Please see Appendix E.2 for a discussion.

Related Work

DPMs and their variants. The idea that a data generation process can be constructed by reversing a diffusion process is initially introduced by Sohl-Dickstein et al. (2015). Specifically, Sohl-Dickstein et al. (2015) reverse the diffusion process using a Markov chain with discrete timesteps, which is learned on the ELBO objective. Ho et al. (2020) propose to parameterize the mean of the Markov chain by the noise prediction network, which share parameters between different timesteps and is learned on a MSE loss. Song et al. (2020b) consider DPMs with infinitesimal timesteps, which can be represented by stochastic differential equations (SDEs). Recently, some variants of DPMs are proposed. Kingma et al. (2021) additionally learn the forward diffusion process, Vahdat et al. (2021) learn DPMs in latent space and Dockhorn et al. (2021) introduce auxiliary velocity variables to the diffusion process.

Based on the powerful generation performance, DPMs have shown great promise in a wide variety of applications, such as controllable generation (Choi et al., 2021; Meng et al., 2021; Sinha et al., 2021; Nichol et al., 2021), voice conversion (Popov et al., 2021), image super-resolution (Saharia et al., 2021; Li et al., 2021), image-to-image translation (Sasaki et al., 2021), shape generation (Zhou et al., 2021), 3d point cloud generation (Luo & Hu, 2021) and time series forecasting (Rasul et al., 2021).

Faster DPMs. In addition to exploring covariance design in DPMs, there are also other attempts toward faster DPMs. Several works explore to find short trajectories by applying searching algorithms, e.g., grid search (Chen et al., 2020), dynamic programming (Watson et al., 2021) and differentiable search (Anonymous, 2022a). These searching algorithms can be further combined with our method to get a better performance. For example, the combination of the dynamic programming and our method leads to a better likelihood performance, as shown by the OT results in Table 4. Several works change the model family in the reverse process, replacing the Gaussian model. Luhman & Luhman (2021); Anonymous (2022b) distill the original reverse process to a model of a single or a few timesteps, and Xiao et al. (2021) model the reverse process as a conditional generator and propose to train with an adversarial loss. However, their likelihood estimation is non-trivial due to the change of the model family. Several works design faster solvers (Jolicoeur-Martineau et al., 2021; Popov et al., 2021) for DPMs with continuous timesteps. However, these solvers only consider isotropic covariances as in Ho et al. (2020); Bao et al. (2022), which are less flexible compared to ours.

Conclusion

We consider diagonal and full covariances to improve the expressive power of DPMs. We derive the optimal mean and covariance in terms of the maximum likelihood, and also correct it given an imperfect mean. Both the optimal and the corrected ones can be decomposed into terms of conditional expectations over functions of noise, which can be estimated by minimizing MSE losses. We consider diagonal covariances in our implementation for computational efficiency. Besides, we adopt a parameter sharing scheme for inference efficiency and a two-stage training process motivated by our theory. Our method is applicable to DPMs with both discrete and continuous timesteps. Empirically, our method consistently outperforms a wide variety of baselines on likelihood estimation, and improves the sample quality especially on a small number of timesteps.

Acknowledgements

This work was supported by NSF of China Projects (Nos. 62061136001, 61620106010, 62076145, U19B2034, U1811461, U19A2081, 6197222); Beijing NSF Project (No. JQ19016); Beijing Outstanding Young Scientist Program NO. BJJWZYJH012019100020098; a grant from Tsinghua Institute for Guo Qiang; the NVIDIA NVAIL Program with GPU/DGX Acceleration; the High Performance Computing Center, Tsinghua University; and Major Innovation & Planning Interdisciplinary Platform for the “Double-First Class” Initiative, Renmin University of China.

References

Appendix A Proof

which admits an optimal solution ${\bm{\mu}}^{*}={\bm{\mu}}_{q}$.

By taking gradients, we know $\sigma_{i}^{-2}\left[({\bm{\Sigma}}_{q})_{i,i}+({\bm{\mu}}_{i}-({\bm{\mu}}_{q})_{i})^{2}\right]+\log\sigma_{i}^{2}$ as a function of $\sigma_{i}^{2}$ gets its minimum at

Suppose $q({\bm{x}}_{0:N})$ is defined as Eq. (2). Then we have

where ${\bm{\epsilon}}_{n}=\frac{{\bm{x}}_{n}-\sqrt{\overline{\alpha}_{n}}{\bm{x}}_{0}}{\sqrt{\overline{\beta}_{n}}}$ is the noise used to generate ${\bm{x}}_{n}$ from ${\bm{x}}_{0}$ and $\gamma_{n}=\sqrt{\overline{\alpha}_{n-1}}-\sqrt{\overline{\beta}_{n-1}-\lambda_{n}^{2}}\sqrt{\frac{\overline{\alpha}_{n}}{\overline{\beta}_{n}}}$.

Bao et al. (2022) show that the KL divergence between a joint probability density function $q({\bm{x}}_{0:N})$ and a Markov chain $p({\bm{x}}_{0:N})$ can be written as

where $c$ is a constant only related to $q$ (see Lemma 8 in Bao et al. (2022)). As a result, the optimization of the KL divergence is decomposed into $n$ independent optimization sub-problems:

Similarly to the proof in Theorem 3.1, these optimizations are equivalent

According to Lemma A.1 and Lemma A.2, we know the optimal ${\bm{\mu}}_{n}^{*}({\bm{x}}_{n})$ is

Similarly to the proof in Theorem 3.1, these optimizations are equivalent

According to Lemma A.2 and the parameterization of ${\bm{\mu}}_{n}({\bm{x}}_{n})$ in Eq. (6), we have

A.2 Proof of Proposition 5.1 and Proposition 5.2

The conditional distribution of ${\bm{x}}_{s}$ given ${\bm{x}}_{t}$ and ${\bm{x}}_{0}$ determined by the SDE is

According to Lemma A.1 and Lemma A.3, the optimal solution is

A.3 Strictly Improved Expressive Power when the Data Distribution is a Mixture of Gaussian

Suppose $q({\bm{x}},{\bm{y}})$ is a probability density function, s.t., $q({\bm{x}})=\sum\limits_{j=1}^{J}\gamma_{j}{\mathcal{N}}({\bm{\mu}}_{j},c{\bm{I}})$ with $J\geq 2$ and pairwise distinct $\{{\bm{\mu}}_{j}\}_{j=1}^{J}$, and $q({\bm{y}}|{\bm{x}})={\mathcal{N}}({\bm{y}}|\sqrt{\alpha}{\bm{x}},\beta{\bm{I}})$ for some $\alpha,\beta>0$. Suppose $p({\bm{x}}|{\bm{y}})={\mathcal{N}}({\bm{x}}|{\bm{\mu}}({\bm{y}}),{\bm{\Sigma}}({\bm{y}}))$ is a Gaussian distribution conditioned on ${\bm{y}}$. Then we have,

Since $q({\bm{y}})$ is fully supported, it implies that

which is also a mixture of Gaussian. According to the property of mixture of Gaussian, the diagonal of its covariance is

$\sum\limits_{j=1}^{J}\eta_{j}({\bm{y}})\|{\bm{\mu}}_{j}\|_{2}^{2}-\|\overline{{\bm{\mu}}}({\bm{y}})\|_{2}^{2}$ is also irrelevant to ${\bm{y}}$. Let ${\bm{y}}=t{\bm{y}}_{0}$, $j_{0}=\operatorname*{arg\,max}\limits_{j}{\bm{\mu}}_{j}^{\top}{\bm{y}}_{0}$. Then $\lim\limits_{t\rightarrow\infty}\eta_{j_{0}}(t{\bm{y}}_{0})=1$, and

Thus, $\sum\limits_{j=1}^{J}\eta_{j}({\bm{y}})\|{\bm{\mu}}_{j}\|_{2}^{2}-\|\overline{{\bm{\mu}}}({\bm{y}})\|_{2}^{2}=0$ for all ${\bm{y}}$. This means

Thus, ${\bm{\mu}}_{1}={\bm{\mu}}_{2}=\cdots={\bm{\mu}}_{J}=\overline{{\bm{\mu}}}({\bm{y}})$. This contradicts the assumption that $\{{\bm{\mu}}_{j}\}_{j=1}^{J}$ are pairwise distinct. Therefore, Eq. (25) must hold. ∎

If the data distribution $q({\bm{x}}_{0})=\sum\limits_{j=1}^{J}\gamma_{j}{\mathcal{N}}({\bm{\mu}}_{j},c{\bm{I}})$ is a mixture of $J\geq 2$ Gaussian, and $\{{\bm{\mu}}_{j}\}_{j=1}^{J}$ are pairwise distinct, then we have

where $c$ is a constant only related to $q$. Thus, we only need to prove

Appendix B Results for Full Covariances

In this part, we derive results on full covariances. In Proposition B.2, we derive the optimal solution to the joint optimization problem (3). In Proposition B.3, we derive the optimal covariance to the optimization problem (13) w.r.t. the covariance solely. Their proof requires Lemma B.1, which is very similar to Lemma A.1.

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

Suppose ${\bm{\Sigma}}_{n}({\bm{x}}_{n})$ is a full covariance. Then the optimal mean to problem (3) is ${\bm{\mu}}_{n}^{*}({\bm{x}}_{n})$ as in Eq. (4), and the optimal covariance to problem (3) is

Bao et al. (2022) show that the KL divergence between a joint probability density function $q({\bm{x}}_{0:N})$ and a Markov chain $p({\bm{x}}_{0:N})$ can be written as

where $c$ is a constant only related to $q$ (see Lemma 8 in Bao et al. (2022)). As a result, the optimization of the KL divergence is decomposed into $n$ independent optimization sub-problems:

According to Lemma A.2 and the parameterization of ${\bm{\mu}}_{n}({\bm{x}}_{n})$ in Eq. (6), we have

However, obtaining samples from Gaussian transitions with full covariances requires a decomposition of the covariance matrix, e.g., Cholesky decomposition, which has a ${\mathcal{O}}(d^{3})$ time complexity. In practice, $d$ can be large (e.g., high resolution images), and such a decomposition can be time consuming. Therefore, in practical settings, it is more suitable to consider diagonal covariances as in the main text, which do not involve a matrix decomposition.

Appendix C Details on Extension to DPMs with Continuous Timesteps

where $t$ is uniformly sampled from $[0,T]$. Then we obtain the estimate of ${\bm{\sigma}}_{s|t}^{*}({\bm{x}}_{t})^{2}$:

Appendix D Inference on Trajectories

To speed up the inference, we can reverse a shorter forward process $q({\bm{x}}_{\tau_{1}},\cdots,{\bm{x}}_{\tau_{k}}|{\bm{x}}_{0})$ constrained on a trajectory $1\leq\tau_{1}<\cdots<\tau_{K}=N$ of $K$ timesteps (Song et al., 2020a; Bao et al., 2022). Following the notation of Bao et al. (2022), the shorter forward process is defined as

Similarly to Eq. (1), the shorter Markov model to reverse the forward process is defined as

where ${\bm{\epsilon}}_{\tau_{k}}=\frac{{\bm{x}}_{\tau_{k}}-\sqrt{\overline{\alpha}_{\tau_{k}}}{\bm{x}}_{0}}{\sqrt{\overline{\beta}_{\tau_{k}}}}$ is the noise used to generate ${\bm{x}}_{\tau_{k}}$ from ${\bm{x}}_{0}$ and $\gamma_{\tau_{k-1}|\tau_{k}}=\sqrt{\overline{\alpha}_{\tau_{k-1}}}-\sqrt{\overline{\beta}_{\tau_{k-1}}-\lambda_{\tau_{k-1}|\tau_{k}}^{2}}\sqrt{\frac{\overline{\alpha}_{\tau_{k}}}{\overline{\beta}_{\tau_{k}}}}$. Thus, we can reuse the SN prediction network ${\bm{h}}_{n}({\bm{x}}_{n})$ trained on the loss in Eq. (9), and estimate ${\bm{\sigma}}_{\tau_{k-1}|\tau_{k}}^{*}({\bm{x}}_{\tau_{k}})^{2}$ as follows:

We emphasize that the inference on a shorter trajectory doesn’t need to train a new NPR (or SN) prediction network on the corresponding shorter process, since the shorter process and the original one share the same marginal distribution. We reuse the one trained on the full timesteps.

Appendix E Additional Discussion

E.2 NPR-DPM is Not Always Better Than SN-DPM

We note that SN-DPM is better w.r.t. FID and NPR-DPM is better w.r.t. likelihood. The performance of NPR-DPM on likelihood is guaranteed, since it aims to optimize the log-likelihood with the actual imperfect mean. Note the sample quality and likelihood are not necessarily consistent (see Section 3.2 in Theis et al. (2015)), and NPR-DPM does not directly optimize the sample quality. Thus, the performance of NPR-DPM on FID is not guaranteed.

In practice, the NPR-DPM is useful in applications such as lossless compression (see Section 7.3 in Kingma et al. (2021)), where the likelihood performance and the computation efficiency play a central role.

Appendix F Experimental Details

In Table 5, we list details of pretrained noise prediction networks used in our experiments. The ImageNet 64x64 pretrained model includes a noise prediction network and additionally a covariance network, and we only use the former one.

F.2 Structure Details of Prediction Networks

F.3 Details of Memory and Time Cost

In Table 7, we list the memory and time cost of models (with the corresponding methods) used in our experiments. The extra memory cost of the NPR (or SN) prediction network is negligible. The extra time cost of the NPR (or SN) prediction network is negligible on CIFAR10, CelebA 64x64, at most 4.5% on ImageNet 64x64, and at most 10% on LSUN Bedroom.

F.4 Training Details

We use a similar training setting to that of the noise prediction network in Bao et al. (2022). On all datasets, we use the AdamW optimizer (Loshchilov & Hutter, 2017) with a learning rate of 0.0001; we train 500K iterations; we use an exponential moving average (EMA) with a rate of 0.9999. We use a batch size of 64 on LSUN Bedroom, and 128 on other datasets. We save a checkpoint every 10K iterations and select the one with the best FID on 1000 generated samples. By default, these samples are generated with full timesteps, except for LSUN Bedroom. On LSUN Bedroom, these samples are generated with 100 timesteps for an acceptable time cost (the time cost on LSUN Bedroom is much larger than other datasets as shown in Table 7).

Training a NPR (or SN) network on CIFAR10 takes about 32 hours on one GeForce RTX 2080 Ti. Training on CelebA 64x64 takes about 72 hours on two GeForce RTX 2080 Ti. Training on ImageNet 64x64 takes about 83 hours on two GeForce RTX 2080 Ti. Training on LSUN Bedroom takes about 171 hours on four GeForce RTX 2080 Ti.

F.5 Log-Likelihood and Sampling

The image data is a vector of integers in $\{0,1,\cdots,255\}$ and we linearly scale it to $$ following Ho et al. (2020); Bao et al. (2022).

Log-likelihood. Following Ho et al. (2020); Bao et al. (2022), we discretize the last Markov transition $p({\bm{x}}_{0}|{\bm{x}}_{1})$ to get discrete NLL and its upper bound. The likelihood results are evaluated on the whole test dataset by default.

F.6 Number of Monte Carlo Samples to Calculate the Optimal Trajectory

F.7 Experimental Details of Table 2, 3 & 4

In Table 2 & 4, the results of baselines and our methods are based on the same noise prediction networks (i.e., those listed in Table 5). The baseline results in DPMs with discrete timesteps are provides by Bao et al. (2022). As for DPMs with continuous timesteps, we get results of the Euler-Maruyama solver and the probability flow by running the pytorch official code (https://github.com/yang-song/score_sde_pytorch) of Song et al. (2020b), and get results of the ancestral sampling and Analytic-DPM using our implementation.

In Table 3, the results of DDPM, DDIM, Analytic-DPM and our methods are based on the same noise prediction networks (i.e., those listed in Table 5). The time cost in a single timestep is the same for all baselines (DDPM, DDIM, Improved DDPM and Analytic-DPM), based on the statistics reported in Table 4 in Bao et al. (2022). The ratio of the time cost in a single timestep in NPR-DPM and SN-DPM to that in baselines is based on Table 7. The result of baselines are provided by Bao et al. (2022). The results of NPR-DPM and SN-DPM are based on the DDPM forward process on LSUN Bedroom, and based on the DDIM forward process on other datasets. These choices require smaller number of timesteps than their alternatives.

Appendix G Additional Experiments

G.2 Variance of Likelihood Results

We note that under the full timesteps, the likelihood results of Analytic-DDPM and NPR-DDPM in Table 4 are close. To reduce the effect of randomness, we report the standard deviation of the likelihood results under the full timesteps in Table 9.

G.3 Comparison After Normalizing for the Extra Cost

We present results after normalizing for the extra inference cost. As shown in Table 10, our method still outperforms the strong baseline Analytic-DPM. We also compare with the “Gotta Go Fast” SDE solver (Jolicoeur-Martineau et al., 2021), and our method performs better when the NFE is small (e.g., $\leq 49$), and thereby is more efficient. When the NFE is large (e.g., $\geq 147$), different methods have similar performance.

We also present results after normalizing for the extra training cost. On CIFAR10 (LS), we train a noise prediction network with 300K iterations and train a NPR (or SN) prediction network with 200K iterations. This ensures that the total training cost is strictly less than that of the original CIFAR10 (LS) model provided by Bao et al. (2022), which is trained with 500K iterations. With 24 and 49 timesteps respectively, NPR-DDPM gets a FID of 10.56 and 6.24; SN-DDPM gets a FID of 7.51 and 4.64. The results are not affected much by the normalization for the extra training cost, and our method still outperforms the strong baseline Analytic-DPM.

G.4 Samples

Recall that Table 3 reports the least number of timesteps required to achieve a FID around 6 on CIFAR10, CelebA 64x64 and LSUN Bedroom. In Figure 1-3, we show generated samples of NPR-DPM and SN-DPM under these number of timesteps on these datasets. Here we use $K$ to denote the number of timesteps. In Figure 4, we also show generated samples on ImageNet 64x64 under $K=25$ timesteps and the DDPM forward process.

In Figure 6-9, we show more generated samples from both NPR-DPM and SN-DPM under both the DDPM and DDIM forward processes, as well as the VP SDE.