End-To-End Learning of Gaussian Mixture Priors for Diffusion Sampler
This addresses a key bottleneck in diffusion models for sampling from complex target distributions, offering a more adaptable and expressive prior to improve sample quality and efficiency.
The paper tackled the problem of diffusion models struggling with exploration and mode collapse when using simple Gaussian priors, by proposing end-to-end learnable Gaussian mixture priors (GMPs) and an iterative refinement strategy, resulting in significant performance improvements across diverse benchmarks without extra target evaluations.
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 by adding 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.