End-to-end Learning of Gaussian Mixture Priors for Diffusion Sampler

We thank the reviewer for taking the time to review our work and the many helpful comments and suggestions. We hope the following replies address the questions and concerns raised. Additionally, we uploaded a revised version of the manuscript that includes several changes as outlined in the sequel. The changes are highlighted using the color magenta.

Why not using the log-variance objective which has proven to be very effective? [2] (Note that this objective wouldn't require the reparameterization trick on the prior)

The log-variance loss is indeed a very interesting alternative to the KL divergence. However, we found the performance to be overall worse compared to KL and in some cases unstable. We included additional results in Appendix D, “Ablation: Kullback-Leibler vs. Log-Variance divergence”.

Moreover, we thank the reviewer for the very interesting fact about the reparameterization trick. We mentioned in the main part of the paper that one can optimize arbitrary divergences and added a footnote that the reparameterization trick is only necessary if the optimization requires expectations with respect to samples from the generative process $\mathcal{P}^{\theta}$.

Could you provide an ablation study in a controlled environment (typically sampling Mixture of Gaussian) to actually challenge the procedure and assess that the claims were addressed?

We thank the reviewer for the great suggestion. We added an ablation study on a two-dimensional Gaussian mixture target in Appendix D, “Ablation on Gaussian Mixture Target”.

How does your algorithm scale with the number of modes $K$ as GMM notoriously showcase numerical instabilities in this scenario?

We found that end-to-end training of the GMM parameters works surprisingly well and didn't observe any numerical instabilities even when using $K>10$ components (see e.g. Figure 7 & 8).

Fig 4 and Fig 9 clearly show that the model mode collapsed as most of the generated images are almost identical. How do you explain this behavior?

DIS-GMP in Fig. 4 and 9 used a ‘naively’ initialized GMM with zero mean and unit variance components. Note that this is not well-suited for tackling challenging exploration problems such as high-dimensional targets with well-separated modes. For this reason, we additionally introduced the IMR strategy to facilitate mode discovery which yields better samples at the cost of additional design choices.

According to Fig 6, the Funnel sampling procedure (for DIS) took over 5 hours. Is it actually competitive against MCMC samplers like MALA or HMC under the same computational budget ?

MCMC algorithms (which are not model-based), such as MALA or HMC, are only guaranteed to converge to the target distribution $\pi$ in infinite time. In contrast, learned samplers such as DIS are guaranteed to sample from the target in a finite time horizon $T$ (given that the learned model has learned the optimal control).

If time permits we will add another experiment that compares diffusion-based sampler to traditional MCMC methods under the same computational budget. However, we want to note that these methods are known to have problems mixing for complicated targets such as with the Funnel target or multimodal densities [1].

It seems like SMC and CRAFT have been unfairly implemented. Indeed, as those algorithms use a different annealing scheme, there are no reasons to keep the same number of intermediate distributions. Moreover, as they are cheaper, this number should be increased. Could you re-run the experiments with increased budget ?

To the best of our knowledge, SMC, i.e. AIS with resampling uses the same annealing as MCD with the difference that MCD uses a learned control network to counteract the suboptimality of the backward Markov kernels used by SMC [2, Section 2.2 & 3]. Moreover, CRAFT converges to the SDE described by CMCD in the limit of $N \rightarrow \infty$ (see [3], Section 5.2). The major difference between CRAFT and CMCD is the parameterization for the transition densities, where CRAFT uses normalizing flows and CMCD a control network. In light of this, we do believe that the comparison is fair, in the sense, that all algorithms use the same number of discretization steps.

Nevertheless, we included a comparison to a long-run SMC sampler which uses up to $N=4096$ discretization steps. Please note that using a high number of discretization steps is typically not possible for CRAFT since one has to maintain a normalizing flow model for every step in memory. We included the additional results in Appendix D, “Comparison to long-run Sequential Monte Carlo”.

As shown in the left side of Figure 4, the proposed IMR strategy seems not work well in all metrics (DIS-GMP, DIS-GMP-IMR), how could this be explained?

We thank the reviewer for the question. We attribute the slightly worse ELBO and log Z values to the following two reasons:

1. Limitations of ELBO and $\Delta \log Z$ as Metrics for Multi-Modal Targets: For multi-modal targets with well-separated modes, ELBO and $\Delta \log Z$ values often fail to reflect sample quality accurately. In such cases, these metrics tend to favor models that fit a single mode perfectly which was also observed in [1]. In contrast, the Sinkhorn distance, which leverages ground-truth samples, provides a more reliable criterion for evaluating sample quality. For the Sinkhorn distance, we observe consistent improvements as $K$ increases, which aligns with the qualitative results shown in Figures 4 and 9.
2. Increased Complexity of Control Functions for Larger $K$: With higher $K$, the diffusion model has to learn more complex control functions, as it needs to operate over the support of the entire Gaussian Mixture Model (GMM) rather than a single Gaussian. This added complexity can introduce more opportunities for approximation errors, which may negatively impact ELBO and $\Delta \log Z$ values compared to using a single learnable Gaussian. Nevertheless, the resulting samples from DIS-GMP are closer to the target distribution in terms of optimal transport (as indicated by the Sinkhorn distance). Importantly, these errors remain significantly smaller than those observed with non-learnable priors, as demonstrated in Figure 4.

We thank the reviewer again and agree that these results require further clarification. To address this, we have expanded the discussion of these results in Section 6.2 of the revised manuscript.

As introduced in Section 5, the IMR strategy requires to optimize parameters until a predefined criterion is met, how to choose the criterion and how does the criterion affect the training result?

We chose to use 500 training iterations as an ad-hoc decision, as it worked well in practice; through visual inspection of model samples, we observed that newly added components focused on different modes. Consequently, we kept this value without exploring alternative values. To provide further insight, we included an ablation study in Appendix D, “Ablation: Iterations for IMR”, where we find that the model performs robustly across various choices for the iteration at which new components are added. Nevertheless, the concept of iteratively adding components is important, such that when a new component is added, the initialization is informed by the already existing mixture model (note that the heuristic in Eq. 22 depends on the likelihood of the GMP) to prevent initializing components at the same location.

What are the potential limitations or drawbacks of using Gaussian mixture priors compared to other types of learnable priors, and how could these be addressed?

To the best of our knowledge, we are the first to introduce learnable priors in the context of diffusion-based sampling. As mentioned in an earlier response, Gaussian Mixture Priors (GMPs) offer notable flexibility but can also complicate the training of control functions. This complexity arises because the control functions must operate over the support of the entire Gaussian Mixture Model (GMM), rather than a single Gaussian, which can lead to approximation errors.

We believe these errors could be significantly mitigated by employing more sophisticated architectures, as opposed to the two-layer MLPs used in our current work. Exploring more advanced architectures remains an exciting direction for future research.

We thank the reviewer for their feedback and hope that our responses have successfully addressed the questions and concerns raised. Should the reviewer have any additional questions or further concerns, we would be more than happy to provide further clarifications or answers

We thank the reviewer for taking the time to review our work and for the many helpful questions. We hope the following replies address the questions and concerns raised. Additionally, we uploaded a revised version of the manuscript that includes several changes as outlined in the sequel. All changes are highlighted using the color magenta.

The main concern of the paper is that the considered numerical experiments are oversimplified. […] The authors only considered logistic regression models and the fashion target on the MNIST fashion dataset in experiments.

We appreciate the reviewer’s concern. However, we would like to highlight that our work includes the majority of the target densities presented in a recent sampling benchmark study [1]. Moreover, different from generative modeling, sampling problems with $d=100$ are already considered medium to high dimensional. See e.g. [2, 3] for recent works that only consider problems with $d \leq 50$. As discussed in [1, Section 7.1], the $d=784$ dimensional Fashion MNIST target is recognized as a highly challenging problem, where alternative approaches struggle to achieve satisfactory sample quality and diversity.

It is important to clarify that the Fashion target is only loosely related to the Fashion MNIST dataset. Specifically, the dataset is utilized solely for training a normalizing flow, which subsequently serves as the target density. The actual dataset itself is not employed for training the diffusion sampler. This somewhat artificial problem introduces several significant challenges, with the most notable being the presence of well-separated modes in a high-dimensional space.

It is well known that diffusion models are widely used for generating samples with complicated priors, such as high-quality real-world images.

We believe there may be a slight misunderstanding regarding the specific context of our paper. In posterior sampling, the target distribution can be expressed as $\pi(x) = p(x|y) = \frac{p(y|x)p(x)}{Z}$, where $y$ is a given measurement and $p(x)$ and $p(y|x)$ are the prior and likelihood, respectively. For certain problems, the prior can also be more complex, such as those encountered in inverse problems involving image, audio, or video distributions [4,5,6,7]. In settings that differ from ours, where samples from the prior $p(x)$ are available, recent methods leverage diffusion priors, i.e., diffusion models pre-trained, to facilitate sampling from $\pi$. This is, however, different from our (more general) setting where minimal assumptions about the specific form of the target are made. Moreover, the term `prior', different from Bayesian inference, is commonly used in diffusion-based sampling (see e.g. [3]) to refer to the starting distribution of the stochastic differential equation which transports samples to the target.

How can the proposed iterative model refinement (IMR) strategy be further improved to optimize the selection criteria for adding new components and dynamically adjust the number of components during training?

This is indeed an interesting question we already highlighted in the future work section. One interesting avenue for future research is the improvement of the generation of candidate samples. We believe that an interesting direction is to use Markov chains where the invariant distribution is the target on higher temperatures $T$ for $\pi(x,T) \propto \exp \frac{-E(x)}{T}$. This would yield biased samples but facilitates the mixing of the Markov chain and thus encourages exploration of different regions of the target support which can help to find promising candidates for initializing new means of the Gaussian mixture prior.

The mean of a newly added Gaussian component is chosen according to (22), how about the choice of variance?

We thank the reviewer for the question and apologize for missing this detail in the original submission. The Gaussian components were initialized using unit variance. We added the corresponding details in Appendix C.2 in the revised version.

I sincerely thank the authors for the detailed reply. I still have a remaining concern regarding the performance comparison of the proposed method with the SOTA in the field. Particularly, if possible, please provide comparisons with score-based plug-and-play methods, such as [R1], [R2].

[R1] Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in neural information processing systems, 32, 2019.

[R2] Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13456, 2020.

We thank the reviewer for taking the time to review our work and for the helpful questions and suggestions. We uploaded a revised version of the manuscript that includes several changes outlined in the sequel. All changes are highlighted in magenta.

For the exploration and mode collapse issues discussed, is it not possible to transform the target data distribution such that its support is covered by a single unit Gaussian, given that you use the knowledge of the support of the true distribution when adding new modes? How much worse is the base DIS compared to the learned and learned+mixture variants under this experiment setting?

We completely agree with the reviewer that ensuring a fair comparison requires the base version of DIS to utilize the same information as DIS-GMP+IMR. For the Fashion target, we indeed transformed the base version of DIS such that the prior covers the support of the target by setting the variance of the initial distribution to the same value that was used for initializing the candidate samples (see Appendix C.2).

We sincerely apologize for missing this detail in the initial submission and added further details in Section 6.2 (Iterative Model Refinement (IMR)) and Appendix C.2 (Gaussian Priors (GP) and Gaussian Mixture Priors (GMP)).

Some more ablations could help with understanding the specific advantages of the proposed method better. One of the issues of the unimodal Gaussian prior that is mentioned is the increased complexity in the learned dynamics. In some examples of Fig. 5 the EES only marginally improves with more modes, making the number of diffusion steps the more impactful hyperparameter. This hints towards the dynamics not being simpler with more modes, as we would expect fewer diffusion steps to also yield good EES in those cases. A more in-depth experiment that demonstrates the differences in dynamics could help ground the overall argument for a mixture prior.

We thank the reviewer for the comment. Fig. 5 only compares between models where the prior is learned, ($K=1$ uses a learned Gaussian prior), in which case the variation in dynamics is typically not as pronounced as for a non-learned prior, see e.g. Fig. 1, C2 vs. C3.

We, therefore, added a more in-depth experiment where we quantify the variation of the dynamics over time using the spectral norm of the Jacobian of the learned forward control, i.e.,

$

S = \mathbb{E} \left[\int_{0}^T|\frac{\partial \sigma u^{\theta}(x,t)}{\partial x}|_2 d t\right],

$

on various tasks to compare the dynamics between DIS and DIS-GMP. The results are included in Appendix D (Ablation: Variation of dynamics) which indicate that DIS-GMP indeed has less variation in the dynamics. These findings are also in line with those in Fig. 3 and Table 2 where DIS-GMP has significantly higher ELBO values compared to DIS without learned prior.

Furthermore, we provide an additional experiment in Appendix D (Ablation on Gaussian Mixture Target) where we study a two-dimensional Gaussian mixture model in order to provide further insights and comparisons.

We thank the reviewer for their feedback and hope that our responses have successfully addressed the questions and concerns raised. Should the reviewer have any additional questions or further concerns, we would be more than happy to provide further clarifications or answers.

We thank the reviewer for taking the time to review our work and for the many helpful comments and suggestions. We hope the following replies address the questions and concerns raised. Additionally, we uploaded a revised version of the manuscript that includes several changes outlined in the sequel. All changes are highlighted in magenta.

In Section 5, it seems that the positioning of challenges C1 and C2 may have been reversed.

We thank the reviewer for bringing this to our attention. In the revised version, we swapped the positions accordingly.

While Figure 6 presents wall-clock time for relatively low-dimensional cases, it would be helpful to understand how computational demands scale in higher-dimensional scenarios. Could you provide an analysis of the computational cost associated with varying the number of mixture components  $K$ and diffusion steps $N$ for a more formal comparison? (e.g. Big-O notation)

Formally, we have the time-complexity $\mathcal{O}(NK)$. Here, the $K$ stems from the additional requirement of computing the likelihood of the GMP in every diffusion step. However, in practice, the evaluation of the likelihood for the GMP can be parallelized across its components, which substantially reduces the computational overhead. Moreover, most diffusion-based methods additionally require evaluating the target density at every step. For experiments where it is expensive to evaluate the target, e.g. if the target is a big neural network or requires simulating molecular dynamics, the relative cost of evaluating the GMM likelihood becomes less pronounced or even negligible.

The results reported in the paper are for DIS which only requires evaluating the GMM likelihood and not the target. We additionally included a comparison between the wallclock time for DIS and CMCD on the high-dimensional Fashion target in Table 5. For CMCD, which requires target evaluations in every diffusion step, the relative cost added by the GMP is minor.

Moreover, we added additional discussions in Appendix B (”Time complexity”) and Appendix D (”Wallclock time”).

How does estimated memory consumption vary with changes in $K$ and $N$?

Using the standard ‘discrete-then-optimize’ approach for optimizing the KL divergence in Eq. 11, which requires differentiation through the SDE, the memory consumption is linear in $K$ and $N$. In contrast, using e.g. the stochastic adjoint method for KL optimization [1], the memory consumption is constant. We used the former approach for our experiments due to its simplicity. However, if one uses a large number of components / diffusion steps, the latter option might be the preferred choice. On a sidenote, it is also possible to obtain constant memory consumption by using other loss functions such as the log-variance loss [2] or moment-loss [3].

We thank the reviewer for the question. We included an additional discussion on memory consumption in Appendix B (”Memory consumption.”)

In Figure 4, DIS-GMP performs worse than DIS-GP across all metrics of sampling quality. This outcome seems contrary to what the paper suggests; how might this be explained?

We thank the reviewer for the question. We attribute the slightly worse ELBO and log Z values to the following two reasons:

1. Limitations of ELBO and $\Delta \log Z$ as Metrics for Multi-Modal Targets: For multi-modal targets with well-separated modes, ELBO and $\Delta \log Z$ values often fail to reflect sample quality accurately which was also observed in [4]. In such cases, these metrics tend to favor models that fit a single mode perfectly. In contrast, the Sinkhorn distance, which leverages ground-truth samples, provides a more reliable criterion for evaluating sample quality. For the Sinkhorn distance, we observe consistent improvements as $K$ increases, which aligns with the qualitative results shown in Figures 4 and 9.
2. Increased Complexity of Control Functions for Larger $K$: With higher $K$, the diffusion model has to learn more complex control functions, as it needs to operate over the support of the entire Gaussian Mixture Model (GMM) rather than a single Gaussian. This added complexity can introduce more opportunities for approximation errors, which may negatively impact ELBO and $\Delta \log Z$ values compared to using a single learnable Gaussian. Nevertheless, the resulting samples from DIS-GMP are closer to the target distribution in terms of optimal transport (as indicated by the Sinkhorn distance). Importantly, these errors remain significantly smaller than those observed with non-learnable priors, as demonstrated in Figure 4.

We thank the reviewer again and agree that these results require further clarification. To address this, we have expanded the discussion of these results in Section 6.2 of the revised manuscript.

Thank you for providing additional experiments and explanations in response to my questions. I sincerely appreciate the significant improvements made to the paper during the rebuttal period, as well as the effort you have put into enhancing the experimental results and explanations.

However, I still have a few unresolved concerns:

Limitations of $\text{ELBO}$ and $\Delta\text{ log }Z$ as Metrics for Multi-Modal Targets: For multi-modal targets with well-separated modes, $\text{ELBO}$ and $\Delta\text{ log }Z$ values often fail to reflect sample quality accurately which was also observed in [4].

Computing the Sinkhorn distance requires samples from the target density which are not available for the 'real-world' tasks presented in Table 2. Consequently, for these tasks, we are limited to reporting ELBO and ESS values. We added further details Section 6.1 and Appendix C.4 to make this requirement more explicit.

The authors discuss the limitations of $\text{ELBO}$ and $\Delta\text{ log }Z$ as metrics for multi-modal targets. For real-world datasets, where the true density is unknown, $\text{ELBO}$ and $\text{ESS}$ are used as the primary metrics in this paper. Specifically, in the synthetic density case of FMNIST, where $\text{ELBO}$ has known limitations, the authors demonstrated the effectiveness of GMP and IMR using Sinkhorn distance or $\text{EMC}$ metrics.

However, this raises the question: does the superior $\text{ELBO}$ performance shown in Table 2 for real-world datasets necessarily guarantee the learning of better posterior distributions, or is it only ensured for lower-dimensional cases?Demonstrating that the DIS-GMP + IMR model consistently achieves the highest $\text{ESS}$ performance even in cases where $\text{ELBO}$ and Sinkhorn distance yield conflicting results—such as with Fashion data—could help address this concern indirectly.

We thank the reviewer for their question. The value ‘500’ was an ad-hoc decision that we kept without exploring alternative values. However, we agree that this design-choice requires further explanation. In response, we included an ablation study in Appendix D, “Ablation: Iterations for IMR”, where we find that the model performs robustly across various choices for the iteration at which new components are added. Nevertheless, the concept of iteratively adding components is important, such that when a new component is added, the initialization is informed by the already existing mixture model (note that the heuristic in Eq. 22 depends on the likelihood of the GMP) to prevent initializing components at the same location.

In Table 10, Although authors have provided results showing that performance improves as the number of IMR training steps increases, this insight seems to suggest that "more training within available resources and time is better" rather than offering a clear criterion for selecting the appropriate number of iterations. Introducing a less heuristic stopping criterion or conditions to determine the optimal IMR iterations could yield more robust follow-up results and insights.

Based on the results so far, I keep my current score but remain open to further insights or discussions on these points.

P.S. In Appendix A (Proofs), the terms in Eq. 33 are not eliminated to zero. Specifically, $$ -\frac{2}{\sigma^2} $$ should be $$ \frac{2}{\sigma^2} $$ in Eq. 32.