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

Authors: Denis Blessing, Xiaogang Jia, Gerhard Neumann
Primary area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)
Keywords: Variational Inference, Sampling, Diffusion Models, Mixture Models

Abstract:
Diffusion models optimized via variational inference (VI) have emerged as a promising tool for generating samples from unnormalized target densities. These models create samples by simulating a stochastic differential equation, starting from a simple, tractable prior, typically a Gaussian distribution. However, when the support of this prior differs greatly from that of the target distribution, diffusion models often struggle to explore effectively or suffer from large discretization errors. Moreover, learning the prior distribution can lead to mode-collapse, exacerbated by the mode-seeking nature of reverse Kullback-Leibler divergence commonly used in VI.
To address these challenges, we propose end-to-end learnable Gaussian mixture priors (GMPs). GMPs offer improved control over exploration, adaptability to target support, and increased expressiveness to counteract mode collapse. We further leverage the structure of mixture models by proposing a strategy to iteratively refine the model through the addition of mixture components during training. Our experimental results demonstrate significant performance improvements across a diverse range of real-world and synthetic benchmark problems when using GMPs without requiring additional target evaluations.

Decision: accept (poster)
Reviewer scores: 6, 8, 6, 6

Reviews and discussion:
<thread>
creation date:
1730195711653

thread_type:
official_review

<thread_object>
creator: reviewer_3ski

<content>
summary:
In this paper, the authors propose an innovative sampling approach using the framework of controlled diffusions. They address current limitations in diffusion samplers, such as poor exploration, mode collapse, and numerical instabilities, which hinder model effectiveness. To counter these issues, they suggest parametrizing the prior of the backward diffusion process as a Gaussian Mixture Model (GMM). The authors detail how this adjustment can be applied to both denoising diffusions and Langevin diffusions. They introduce an iterative learning procedure that jointly optimizes the control networks and the GMM prior parameters. Through extensive experiments, the authors demonstrate that their approach outperforms traditional diffusion samplers across a wide range of tasks.

strengths:
- The paper is well written
- The idea of parametrizing the prior seems novel

weaknesses:
- Lack of multimodal experiments (while it was the main motivation according to Fig. 1)
- Competing algorithms seems to have been wrongly implemented as the control networks lack the $\nabla \log \pi$ information (according to L927-928)
- Sec 6 only considers metrics known to be unsuited for multimodal distributions (and especially to detect mode collapse C3) as per [1] (cited by the authors)
- The computational procedure is very heavy

questions:
- Why not using the log-variance objective which as proven to be very effective ? [2] (Note that this objective wouldn't require the reparameterization trick on the prior)
- 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 ?
- How does your algorithm scale with the number of modes $K$ as GMM notoriously showcase numerical instabilities in this scenario ?
- 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 ?
- 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 ?
- 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 ?
- Could you compare X-GMP against X-GP but adding the target informed parametrization of the controlled networks ? It seems like it could make a big difference.
- The problem of overparameterisation of the GMM (when the GMM as more mode than the target distribution) has two different interpretations between L479-483 and Fig 9. Could you clarify it ?
- Could you compute mode-collapse aware metrics such as the EUBO (and derivations) or the EMC from [1] ?

[1] Blessing, D., Jia, X., Esslinger, J., Vargas, F., & Neumann, G. (2024). Beyond ELBOs: A Large-Scale Evaluation of Variational Methods for Sampling. In Proceedings of the 41st International Conference on Machine Learning (pp. 4205–4229). PMLR.

[2] Nüsken, N., Richter, L. Solving high-dimensional Hamilton–Jacobi–Bellman PDEs using neural networks: perspectives from the theory of controlled diffusions and measures on path space. Partial Differ. Equ. Appl. 2, 48 (2021). https://doi.org/10.1007/s42985-021-00102-x
</content>
</thread_object>

</thread>
<thread>
creation date:
1730645275871

thread_type:
official_review

<thread_object>
creator: reviewer_paqt

<content>
summary:
This paper proposes an end-to-end learnable Gaussian mixture priors (GMPs) to address challenges in diffusion-based sampling methods including poor exploration, non-linear dynamics and mode collapse. Generally speaking, the paper contains some interesting ideas, and some experimental results are provided to demonstrate the performance improvement of the proposed method.

strengths:
The paper presents an approach for improving diffusion-based sampling techniques by introducing end-to-end learnable Gaussian Mixture Priors (GMPs). The authors identify three key challenges in diffusion models: lack of exploration capabilities (C1), non-linear dynamics requiring many diffusion steps (C2), and mode collapse due to reverse KL minimization (C3). To address these challenges, the paper proposes GMPs, which offer improved control over exploration, adaptability to target support, and increased expressiveness to counteract mode collapse. Experimental results are provided to demonstrate performance improvement of the proposed methods.

weaknesses:
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. It is well known that diffusion models are widely used for generating samples with complicated priors, such as high-quality real-world images. Yet, from the numeral results provided in the paper, the reviewer cannot see any potential of the proposed method in this respect. As such, it is not convincing to claim that the proposed methodology can indeed solve the problems (C1)-(C3) of the existing diffusion models.

questions:
1. 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?

2. 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?

3. 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?

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

5. 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?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730772683935

thread_type:
official_review

<thread_object>
creator: reviewer_d8id

<content>
summary:
The paper proposes using learned Gaussian mixture priors in diffusion samplers instead of the standard fixed, unimodal Gaussian. The authors experimentally demonstrate the advantages of having a multi-modal prior; better exploration, fewer diffusion steps, and avoiding mode collapse. To effectively train the learned Gaussian mixture prior model, the authors introduce an iterative refinement approach that gradually adds new modes to the prior distribution.

strengths:
- **Novelty**: The idea of learning a mixture prior in a diffusion process is underexplored and could have significant impact on future works. To the best of my knowledge, there have been few papers that explore a learned prior in diffusion models. Additionally, the idea of gradually adding complexity to the prior is very interesting.
- **Evaluation**: The authors provide adequate comparisons between the proposed method and existing approaches that demonstrate the overall improvements when using the proposed model. The results show that their method better captures the different modes of the distributions learned.
- **Presentation**: The writing is succinct and helps the reader understand the necessary preliminaries before diving into the details of the proposed method. Even with limited prior knowledge of the different diffusion samplers out there, the paper is easy to follow.

weaknesses:
- 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.

questions:
- 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?
</content>
</thread_object>

</thread>
<thread>
creation date:
1731153709113

thread_type:
official_review

<thread_object>
creator: reviewer_jr3v

<content>
summary:
This paper introduces a novel approach to diffusion-based sampling that incorporates end-to-end learnable Gaussian Mixture Priors (GMP), offering a more expressive alternative to conventional Gaussian priors. The flexibility of GMP mitigates mode collapse, enhances exploration capabilities, and reduces the need for extensive diffusion steps by alleviating non-linear dynamics. The authors also propose an Iterative Model Refinement (IMR) strategy, progressively increasing model complexity to optimize performance on high-dimensional, multi-modal targets. Extensive experiments demonstrate that using GMP in place of Gaussian priors in existing diffusion-based sampling methods achieves improved sampling performance and outperforms previous methods across synthetic and real-world datasets.

strengths:
- The paper is well-structured, first identifying the limitations of Gaussian priors in diffusion-based sampling and then logically progressing to propose end-to-end learnable Gaussian Mixture Priors (GMP) as a solution. Each step in this argument is clearly justified, supported by sufficient background information and helpful illustrations, making it easy for readers to follow.

- The experimental design effectively supports the theoretical claims, providing empirical evidence of the advantages offered by GMP over traditional Gaussian priors. Additionally, the comprehensive review of related work and background offers readers a useful overview of recent developments in this field.

- The proposed method enables end-to-end learning of the prior distribution, allowing existing diffusion-based sampling methods, which previously relied on fixed naive priors, to benefit from a more optimized starting prior. This adaptability enhances sampling performance and broadens the applicability of the approach across various diffusion-based sampling methods.

weaknesses:
- Although GMP can achieve competitive performance with a single Gaussian prior using fewer steps, high-dimensional data often require more complex dynamics and an increased number of components to handle multiple modes effectively. As both the number of components $K$ and steps $N$ grow, the computational cost increases substantially, raising concerns about the scalability of this approach for even higher-dimensional applications in the future.
- As acknowledged in the paper, there is no clear criterion for determining the appropriate number of mixture components needed for optimal performance. Furthermore, IMR relies on heuristic rules, such as a fixed number of iterations, to add new components, introducing additional hyperparameter tuning requirements for model optimization.
- Some minor inconsistencies and certain claims in the paper are not fully supported by the experimental results, which would benefit from further clarification. For instance, additional explanation may be needed for the DIS results in Figure 4.

questions:
- In Section 5, it seems that the positioning of challenges C1 and C2 may have been reversed. Specifically, in lines 308–310, reducing the complexity of dynamics to minimize the required diffusion steps seems more directly related to the issue identified in the Introduction as C2 (lines 51–53). Conversely, enhancing the exploration capability of diffusion-based sampling methods to cover diverse regions seems more relevant to addressing C1.
- 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)
- How does estimated memory consumption vary with changes in $K$ and $N$?
- 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? Furthermore, while DIS-GMP + IMR with progressively learned components shows better Sinkhorn distance and qualitative sampling results, its ELBO and $\triangle$ log Z performance is lower than that of DIS-GP and DIS-GMP. This observation suggests that ELBO may have limitations in evaluating sampling performance for multi-modal target distributions (as also seen in Figure 9).
- Table 2 presents only ELBO results; does the Sinkhorn distance follow a similar trend to what was observed in the Funnel case?
- The IMR approach uses a criterion of adding components every 500 training iterations. Are there any ablation study results regarding this criterion for IMR?

Minor:
- In Equation 31 of Appendix A.1, the term $(-\frac{2}{\sigma^2})^2$ seems to have a sign inconsistency with Equation 32, where the third term should perhaps be positive.
</content>
</thread_object>

</thread>
<thread>
creation date:
1733311884274

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General response

comment:
We would like to extend our sincere gratitude to all the reviewers for their time, effort, and valuable feedback during the review process. We are glad that all reviewers have expressed their support for acceptance. Below, we provide a brief summary of the enhancements and additional studies included during the rebuttal phase in response to the reviewers' insightful comments:

- Two-dimensional Gaussian mixture target (GMM): We introduced an additional GMM target where we performed several analyses to empirically substantiate the claims made in the paper.
- Lattice $\phi^4$ theory: Moreover, we applied our method to simulate a statistical lattice field theory near, at, and beyond the phase transition at various problem dimensionalities. We provided qualitative and quantitative results that show that Gaussian mixture priors improve free energy estimates and counteract mode collapse.
- Comparison with a long-run Sequential Monte Carlo (SMC) sampler: We additionally compare diffusion samplers with a learned Gaussian mixture prior to SMC with a high number of discretization steps (up to 4096). We find that the diffusion sampler are on par or outperform SMC in most cases.
- Kullback-Leibler vs. Log-Variance divergence: We further compare using the KL divergence to using the log-variance divergence for optimizing the control functions. We find that the KL divergence typically performs better than the log-variance divergence.
- Influence of the control architecture: We further evaluate the performance using an
architecture that additionally incorporates the score of the target. Our results indicate that including the target score often leads to slightly worse results.
- Iterations for IMR: We additionally conducted an ablation study which considers different numbers of iterations for the iterative model refinement scheme at which new components are added. The results indicate that the performance remains stable for different choices of the hyperparameter.
- Variation of dynamics: We additionally compared the variability in the dynamics of the
diffusion models via the time-integrated spectral norm of the learned control function. The results show that Gaussian mixture priors indeed simplify the underlying dynamics compared to using a fixed Gaussian prior distribution.
- Wallclock-time: We provided an analysis showing that the relative cost of learned Gaussian mixture priors in terms of wallclock time is minimal, particularly in cases where the diffusion sampler necessitates the evaluation of the target density at each diffusion step.

Aside from the experiments and ablation studies, we extended the result discussion and included an additional performance criterion, the entropic mode coverage, which is well-suited for quantifying mode collapse.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper proposes a Gaussian mixture prior-based diffusion model for sampling from unknown distributions. The mixture modeling could alleviate the mode-collapse issue in learning the samplers. The effectiveness of the proposed method is shown on both synthetic and real-world data. The reviewers had concerns on the experiments and the complexity of the proposed method and these concerns were addressed in the rebuttal and discussion. The final rating of all reviewers are positive. I agree with the reviewers and recommend acceptance of this paper.