Diffusion-based supervised learning of generative models for efficient sampling of multimodal distributions

TL;DR

The paper tackles high-dimensional multimodal sampling by decomposing the problem into unimodal components. It identifies modes via multi-start optimization, partitions the domain with C-SVC, locally samples each mode with Langevin dynamics, and generates labeled data with training-free diffusion steps to train simple neural transport maps. Mixing weights are estimated with Gaussian bridge sampling to assemble a global generator, enabling near-instantaneous samples from complex targets and enabling application to Bayesian inverse problems. The approach demonstrates accurate recovery up to 100 dimensions and shows advantages over traditional MCMC baselines in KL and runtime metrics, with demonstrated utility in PDE-informed inverse problems. This modular framework offers scalable, parallelizable sampling without requiring intricate architectures or extensive training data, and it can be extended to a variety of multimodal settings.

Abstract

We propose a hybrid generative model for efficient sampling of high-dimensional, multimodal probability distributions for Bayesian inference. Traditional Monte Carlo methods, such as the Metropolis-Hastings and Langevin Monte Carlo sampling methods, are effective for sampling from single-mode distributions in high-dimensional spaces. However, these methods struggle to produce samples with the correct proportions for each mode in multimodal distributions, especially for distributions with well separated modes. To address the challenges posed by multimodality, we adopt a divide-and-conquer strategy. We start by minimizing the energy function with initial guesses uniformly distributed within the prior domain to identify all the modes of the energy function. Then, we train a classifier to segment the domain corresponding to each mode. After the domain decomposition, we train a diffusion-model-assisted generative model for each identified mode within its support. Once each mode is characterized, we employ bridge sampling to estimate the normalizing constant, allowing us to directly adjust the ratios between the modes. Our numerical examples demonstrate that the proposed framework can effectively handle multimodal distributions with varying mode shapes in up to 100 dimensions. An application to Bayesian inverse problem for partial differential equations is also provided.

Figures (16)

• Figure 1: Mixture models consisting of two Gaussian modes at different distances, defined as in \ref{['eq:GM_mixture_2d']}. Note that the $x$-axis and $y$-axis are scaled differently.
• Figure 2: Peaks of target PDF modes found by multi-start gradient descent. From left to right: well separated modes, weakly connected modes, and completely overlapped modes.
• Figure 3: The segmentation of PDF domain into separate subdomains: (left) well-separated modes ($a=-6$); (right) weakly connected modes ($a=2$). The case of completely overlapped modes ($a=6$) is omitted here since no segmentation is necessary. The shaded colors represent the decision function values, with each region (reddish vs. bluish) corresponding to a class. Darker regions indicate areas where the classifier is more confident in its classification, and the curve where colors change from blue to red represents the decision boundary.
• Figure 4: The generation of samples for target unimodal components using Langevin dynamics. Each unimodal component is sampled independently. Top: Input samples. Bottom: Output samples obtained from the Langevin scheme \ref{['eq:lgv']}.
• Figure 5: The output labeled data generated by solving the time-reverse ODEs \ref{['DM:RSDE']}. Samples for each unimodal component, distinguished here by different colors, are generated independently from different diffusion models.