Improved DDIM Sampling with Moment Matching Gaussian Mixtures

We sincerely thank you for taking the time to review our paper. We address the questions and concerns below. > **Motivation for a GMM kernel** Our original motivation for the GMM kernel was to explore latent spaces ($x_t$) better than a Gaussian kernel, as well as a prior work [1] on demonstrating improvements over DDIM without strictly matching the marginals. Concurrent to our work [2] and also a recent work [3] show that the true denoiser conditional distributions ($q(x_{t-1}|x_t)$) are multimodal when taking large steps between $t$ and $t-1$ during sampling. This provides a justification for our work because we prove that the proposed DDIM-GMM kernel yields a more expressive multimodal GMM as the denoiser conditional distribution than DDIM (see Appendix A.7). In addition to providing more flexibility for modeling complex $q(x_{t-1}|x_t)$ than DDIM, our sampler does not require any training with specialized loss functions, unlike [2, 3]. And it can further benefit from fine-tuning with parameter exploration and learning based approaches [1] with the potential to improve over the results reported in the paper. We have made this new perspective and motivation clear in the updated version of the paper (Section 1 and Appendix A.7). > **Comparison with other samplers and ODE perspective** Although the main scope of this paper is to introduce a multimodal kernel in the DDIM framework and demonstrate improvements relative to it, we followed up with an additional experiment comparing with the DPM-Solver sampler. Specifically, we conducted experiments on the class-conditional ImageNet model with classifier-free guidance. The results are added to Appendix A.11 in the paper. Also, we have mentioned connections to ODE-solver perspective of samplers in the Related Works Section A.1, to the best of our knowledge. > **Computational cost** Thank you for pointing this out. Although we briefly mentioned the additional computational overhead towards the end of Sections 3.1 and 3.1.3, we realize that it might have been easy to miss. The proposed approach introduces additional computational overhead in an "initialization" phase prior to sampling. All the GMM mean and variance offsets are precomputed and saved in memory before sampling. We can choose to precompute a single set of GMM parameters per batch or the entire sample set. We experimented with both the options and did not see a significant difference in metrics. So it is computationally more efficient to precompute offsets once and fix them. We also experimented with choosing different GMM parameters for different sampling steps $t$ and found no significant difference with setting them the same across all $t$ in our experiments. Due to saving the additional GMM parameters, there is some memory overhead relative to DDIM but it is negligible, especially in the scenario of choosing a single set of offsets across all samples and time steps. In the scenario of using different offsets across subsets (batches) of samples or sampling steps, the overhead scales linearly along the sample and number of step dimensions. The dimensionality of the latent spaces also determine the computational and memory requirements of the GMM offset parameters. For instance, it might be infeasible to compute the outer products of centered offsets (Eq. 11) if the dimensionality of the latent spaces is high, e.g. high-resolution image space diffusion models. In such cases, the DDIM-GMM-ORTHO-VUB is more feasible as it provides an upper bound for the variance offsets without explicitly computing them. We have made a note of these in a new Section in Appendix A.10. > **Weights in the ELBO Loss Upper Bound** We agree with the reviewer that the weights in the ELBO loss for training DDPM has a non-trivial impact on the resulting model. The purpose of deriving an upper bound (Appendix A.6) is to show that there is an equivalent weighted squared-error loss used in DDPM training that serves as an upper bound for the ELBO loss using the DDIM-GMM inference process. This provides a rough justification that the DDIM-GMM sampler can exploit pretrained diffusion models that are trained with equal weighting. As suggested by the reviewer, it would be ideal to use the derived weights during DDPM training but our experiments suggest empirically that DDPM training with uniform weights still works with the DDIM-GMM sampler. It would be an interesting future work to use the DDIM-GMM ELBO upper bound as a surrogate loss function to train DDPM models and see if it leads to better sample quality with the proposed approach. > **References** 1. Daniel Watson et al., Learning fast samplers for diffusion models by differentiating through sample quality, ICLR 2021. 2. Hanzhong Guo et al., Gaussian Mixture Solvers for Diffusion Models, NeurIPS 2023. 3. Zhisheng Xiao et al., Tackling the Generative Learning Trilemma with Denoising Diffusion GANs, ICLR 2022.

Our sincere thanks for taking time to review our paper and provide feedback. Please find our responses to your concerns below. > **Novelty** Although our approach appears to be incremental over DDIM, we believe it addresses a fundamental problem in accelerated sampling from pre-trained diffusion models. Recent works [1, 2] have provided evidence that the denoising conditional distribution ($q(x_{t-1}|x_t)$) is multimodal when the step size is large. By introducing a multimodal transition kernel, it addresses this issue by providing a way to model multimodality in the denoiser better than DDIM (see Appendix A.7). Our work provides theoretical justification for the choice of the GMM parameters by introducing moment matching constraints. We describe three different ways to select GMM parameters that satisfy the constraints. Although being simple, these proposed variants already demonstrate improvements over DDIM based on the metrics reported in the Experiments section. Further, as noted in the paper, it is possible to learn the GMM parameters using a suitable approach such as the DDSS framework of [3] to better fit the multimodal denoiser distributions. Note that unlike [1, 2], our approach can be applied to pre-trained diffusion models without any training or specialized loss functions. We hope that this leads to further research on more sophisticated approaches for learning GMM parameters in the proposed kernel, with more potential to improve upon DDIM and possibly other accelerated sampling approaches. > **Marginal improvements relative to DDIM** We agree with the reviewer that the improvements seem marginal in many cases relative to DDIM. Although this may be subjective, as noted in Appendix A.13, we consider a relative improvement of 1 point in FID or IS metric as significant. The definition of significance is tied to the main scope of this paper, which is to propose a multimodal GMM kernel in the DDIM framework, with theoretical justification, and provide empirical evidence that it brings improvements relative to DDIM. To this end, we provide three simple training-free schemes to choose GMM parameters and show empirically that they provide improvements relative to DDIM, in most fast sampling scenarios (fewer steps). It is possible to design alternative schemes or learn the GMM parameters, which have the potential to further improve performance. We leave this exploration for future work. Another potential reason for marginal improvements, especially with large number of sampling steps, is that the DDIM unimodal kernel models the denoiser distribution [1] sufficiently well. In other words, the proposed schemes for selecting GMM parameters, with limited parameter exploration (see Section 4.1 and Appendix A.2) are close to DDIM in terms of modeling complexity of the true denoiser (Appendix A.7). With a more fine-grained parameter exploration and learning based schemes for the GMM parameters, there is potential for improvement. We leave this for future work. > **Ablations on GMM parameters** Thank you for pointing this out. We have included ablations on the number of mixture components and the offset scaling factor $s$ in Appendix A.9. > **Application to Text-to-Image Generation** Our approach is applicable to any pre-trained diffusion model including text-to-image generation models. To this end, we conducted experiments with the publicly available Stable Diffusion v2.1 model and reported results in the Experiments Section 4.3. The results clearly demonstrate the benefits of the proposed approach over DDIM based on FID and IS metrics. > **How to improve performance at larger number of steps?** As described above, as the number of sampling steps increases or the step size decreases, the DDIM Gaussian kernel yields an expressive enough denoiser to model the true step distributions ($q(x_{t-1}|x_t)$) [1]. We believe this could be the main reason for the diminishing gains from using a multimodal kernel in this scenario. Also, we propose three simple training-free schemes for choosing GMM parameters and do a limited exploration of the parameter space (number of mixture components, mixture weights, offset scales). An exhaustive exploration of these parameters is one possible way to potentially improve the performance with larger number of steps. Another possibility would be to apply learning based schemes, e.g. [3], to fine-tune GMM parameters to get the most out of the proposed approach. We leave these explorations for future work. > **References** 1. Hanzhong Guo et. al., Gaussian Mixture Solvers for Diffusion Models, NeurIPS 2023. 2. Zhisheng Xiao et. al., Tackling the Generative Learning Trilemma with Denoising Diffusion GANs, ICLR 2022. 3. Daniel Watson et. al., Learning fast samplers for diffusion models by differentiating through sample quality, ICLR 2021.

Thank you for your review. Please see our comments below, specifically addressing the weaknesses. > **Writing and notation** Thank you for your suggestions on improving the writing and notation. We have tried to incorporate these as much as possible. Specifically, the notation in the updated version uses bold symbols for all vectors and matrices. We have removed ambiguity involving matrix-vector operations (e.g., Eq. 9). It is our hope that this version improves readability of the paper. > **Motivation for the proposed GMM parameter computation approaches** Regarding the motivation for the proposed methods for computing GMM parameters, we agree that there is prior work including Expectation Maximization (EM) and method of moments. We believe that the main scope of this paper is to introduce a multimodal kernel into the DDIM sampling framework and demonstrate its benefits over using a unimodal kernel. To this end, we focus on the theoretical justifications of such an approach and suggest a few simple heuristics to come up with a set of GMM parameters efficiently. It is possible to use alternative approaches for estimating GMM parameters such as EM or other learning based approaches [1] specific to DDIM samplers using training data. We chose to use the heuristically driven approaches to keep the samplers training-free, in the spirit of original DDIM. Our empirical results suggest that we see benefits over DDIM even using these simple approaches without searching for their parameters exhaustively. It would be an interesting future direction to explore the possibility of using more complex GMM parameter estimation methods to further improve upon the proposed method. > **Experimental results and organization of the paper** We agree with the reviewer that the experimental results on FID and IS metrics are concisely presented in a tabular form. Our goal here is to illustrate the relative performance between DDIM and DDIM-GMM under different $\eta$ settings as well as sampling steps. This is to help the reader understand benefits and trade-offs under different scenarios due to the interaction between $\eta$ and the offsets ($\delta_t^k$'s and $\Delta_t^k$'s) of the three proposed variants of the DDIM-GMM sampler. We found that graphs are more concise, e.g., comparing the space occupied by Fig. 1 vs Table 5. However, following your suggestion on conveying more information within the available space, we have changed contents in the Experiments section. We show more practical utility of the proposed approach by replacing the class conditional ImageNet results with new experiments on a text-to-image diffusion model. Hopefully the newer version looks more organized and informative. > **References** 1. Daniel Watson et. al., Learning fast samplers for diffusion models by differentiating through sample quality, ICLR 2021.

We sincerely thank you for your review. Please find below our responses to each of your concerns. > **Comparison to DDIM at $\eta=0$ vs $\eta \gt 0$** We agree that DDIM performs best at $\eta=0$. The proposed GMM kernel offsets both the mean and the covariance (variance) components of the original DDIM kernel. At $\eta=0$, the offsets affect only the mean value due to $\sigma_t$'s being zero. In this case, we set the variance of all the components in DDIM-GMM kernel explicitly to zero. While this is does not strictly ensure moment matching, empirically we see that the mean offsets ($\delta_t^k$'s) still provide some improvement relative to DDIM. As $\eta$ increases from 0 to 1, the covariance offsets ($\Delta_t^k$'s) reduce the original DDIM variances ($\sigma_t$), in a subset of the dimensions, to match the moments. We hypothesize that the proposed DDIM-GMM kernels allow better exploration of the latent spaces at different steps $t$, especially when the jumps between successive denoising steps are large. As argued by other works [3, 4], these latent spaces are multimodal and exploring with a multimodal kernel is beneficial relative to a unimodal kernel (see Appendix A.7 in the updated version). Increasing values of $\eta$ encourages this further with the DDIM-GMM kernel, which explains the larger improvements in FID at $\eta=1$ vs $\eta=0$. As the number of sampling steps increases, the denoising distribution $q(x_{t-1}|x_t)$ is effectively explored by a unimodal kernel (DDIM) as observed in [3] explaining the smaller difference between DDIM and DDIM-GMM. > **Value of $\eta$ in Figures** Thank you for pointing this out. We have added the $\eta$ values used in generating the visualizations to figure captions in the updated version. Note that we did not cherry-pick images for visualization. It is perhaps subjective, but we believe that there are qualitative differences in DDIM-GMM vs DDIM outputs. For e.g., in Fig. 12, the semantic concept emerges faster (10 steps) with DDIM-GMM relative to DDIM--the images on the left look blurrier than the ones on the right. Objectively, this is corroborated by the FID/IS values (Fig. 12 => Tables 7 & 8). At larger number of sampling steps (100), we agree that it might be difficult to tell the differences qualitatively. So we have eliminated them. > **Results with the DDIM [1] datasets** We did additional experiments with the same settings as DDIM [1] on LSUN datasets. The results are included in Appendix A.12. Unlike LSUN datasets, we were not able to reproduce the numbers on CIFAR10 in [1] using their codebase. So we list the results of our experiments below. ``` ----------------------------------------------- Steps | 10 | 20 | 50 | 100 | 1000 ----------------------------------------------- DDIM ----------------------------------------------- $\eta=0.0$ | 18.53 | 10.95 | 7.02 | 5.55 | 4.09 $\eta=0.2$ | 19.18 | 11.25 | 7.35 | 5.72 | 4.08 $\eta=0.5$ | 22.09 | 13.10 | 8.36 | 6.42 | 4.21 $\eta=1.0$ | 42.31 | 24.78 | 14.63 | 10.51 | 4.74 ----------------------------------------------- DDIM-GMM-ORTHO-VUB ----------------------------------------------- $\eta=0.0$ | 18.19 | 10.56 | 6.70 | 5.25 | 3.70 $\eta=0.2$ | 18.70 | 10.83 | 7.05 | 5.46 | 3.78 $\eta=0.5$ | 21.24 | 12.53 | 7.91 | 6.10 | 3.97 $\eta=1.0$ | 41.00 | 23.83 | 14.27 | 9.99 | 4.52 ----------------------------------------------- ``` > **Excluding face samples** It is more of a subjective decision not to include face images in our paper as we believe that including them raises potential privacy concerns. > **Comparison with generalized version of DDIM [2]** We compare the performance of DDIM-GMM-ORHTO-VUB ($\eta=0$) with DDSS [2] on CIFAR10 below (FID/IS). ``` ---------------------------------------------------------- Steps | 5 | 10 | 15 | 25 ---------------------------------------------------------- DDSS [2]| 13.77/8.52 | 8.22/8.90 | 4.72/9.26 | 4.25/9.19 DDIM-GMM| 43.54/6.81 | 18.19/8.80 | 10.56/9.59 | 9.31/9.68 ---------------------------------------------------------- ``` We expect DDSS to perform better as it is trained on CIFAR10 samples. Our results look worse due to two possible reasons: 1)it is training-free & 2) the baseline DDIM [1] CIFAR10 results look worse in our experiments (see above). However, our approach can be trained within the DDSS framework to further improve the results. We have noted this in the paper in Sections 4.1 and 5. > **References** 1. Jiaming Song et. al., Denoising diffusion implicit models, ICLR 2021. 2. Daniel Watson et. al., Learning fast samplers for diffusion models by differentiating through sample quality, ICLR 2021. 3. Hanzhong Guo et. al., Gaussian Mixture Solvers for Diffusion Models, NeurIPS 2023. 4. Zhisheng Xiao et. al., Tackling the Generative Learning Trilemma with Denoising Diffusion GANs, ICLR 2022.