Reverse Diffusion Monte Carlo

Recent success of diffusion models has shown great promise for the the reverse diffusion processes in generating samples from a complex distribution (Song et al., 2020; Rombach et al., 2022). In the existing line of works, one is given samples from the target distribution and aims to generate more samples from the same target. One would diffuse the target distribution into a standard normal one, and use score matching to learn the transitions between the consecutive intermediary distributions (Ho et al., 2020; Song et al., 2020). Reversing the learned diffusion process leads us back to the target distribution. The benefit of the reverse diffusion process lies in the efficiency of convergence from any complex distribution to a normal one (Chen et al., 2022a; Lee et al., 2023). For example, diffusing a target multi-modal distribution into a normal one is not harder than diffusing a single mode. Backtracking the process from the normal distribution directly yields the desired multi-modal target. If one instead adopts the forward diffusion process from a normal distribution to the multi-modal one, there is the classical challenge of mixing among the modes, as illustrated in Figure  1. This observation motivates us to ask:

Can we create an efficient, general purpose Monte Carlo sampler from reverse diffusion processes?

For Monte Carlo sampling, while we have access to an unnormalized density function $p_{*}({\bm{x}})\propto\exp(-f_{*}({\bm{x}}))$, samples from the target distribution are unavailable (Neal, 1993; Jerrum & Sinclair, 1996; Robert et al., 1999). As a result, we need a different and yet efficient method of score estimation to perform the reverse SDE. This leads to the first contribution of this paper. We leverage the fact that the diffusion process from our target distribution $p_{*}$ towards a standard normal one $p_{\infty}$ is an Ornstein-Uhlenbeck (OU) process, which admits explicit solutions. We thereby transform the score matching problem into a non-parametric mean estimation one, without training a parameterized diffusion model. We name this new algorithm as reverse diffusion Monte Carlo (rdMC).

To implement rdMC and solve the aforementioned mean estimation problem, we propose two approaches. One is to sample from a normal distribution $\rho_{t}$—determined at each time $t$ by the transition kernel of the OU process—then compute the mean estimates weighted by the target $p_{*}$. This approach translates all the computational challenge to the sample complexity in the importance sampling estimator. The iteration complexity required to achieve an overall $\epsilon$ TV accuracy is $O(\epsilon^{-2})$, under minimal assumptions. Another approach is to use Unadjusted Langevin Algorithm (ULA) to generate samples from the product distribution of the target density $p_{*}$ and the normal one $\rho_{t}$. This approach greatly reduces sample complexity in high dimensions, and yet is a better conditioned algorithm than ULA over $p_{*}$ due to the multiplication of the normal distribution. In our experiments, we find that a combination of the two approaches excels at distributions with multiple modes.

We then analyze the efficacy and efficiency of our rdMC method. We study the benefits of the reverse diffusion approach for both multi-modal distributions and high dimensional heavy-tail distributions that breaks the isoperimetric properties (Gross, 1975; Poincaré, 1890; Vempala & Wibisono, 2019). In multi-modal distributions, the Langevin algorithm based MCMC approaches suffer from an exponential cost for the need of barrier crossing, which makes mixing time extremely long. The rdMC approach can circumvent the hurdle and solve the problem. For high-dimensional heavy-tail distributions, our rdMC method circumvents the often-required isoperimetric properties of the target distributions, thereby avoiding the curse of dimensionality.

Contributions. We propose a non-parametric sampling algorithm that leverages the reverse SDE of the OU process. Our proposed approach involves estimating the score function by a mean estimation sub-problem for posteriors, enabling the efficient generation of samples through the reverse SDE. We focus on the complexity of the sub-problems and establish the convergence of our algorithm. We found that our approach effectively tackles sampling tasks with ill-conditioned log-Sobolev constants. For example, it excels in challenging scenarios characterized by multi-modal and long-tailed distributions. Our analysis sheds light on the varying complexity of the sub-problems at different time points, providing a fresh perspective for revisiting score estimation in diffusion-based models.

In this section, we begin by introducing related work from the perspectives of Markov Chain Monte Carlo (MCMC) and diffusion models, and we discuss the connection between these works and the present paper. Next, we provide a notation for the reverse process of diffusion models, which specifies the stochastic differential equation (SDE) that particles follow in rdMC.

We first introduce the related works as below.

Langevin-based MCMC. The mainstream gradient-based sampling algorithms are mainly based on the continuous Langevin dynamics (LD) for sampling from a target distribution $p_{*}\propto\exp(-f_{*})$. The Unadjusted Langevin Algorithm (ULA) discretizes LD using the Euler-Maruyama scheme and obtains a biased stationary distribution. Due to its simplicity, ULA is widely used in machine learning. Its convergence has been investigated for different criteria, including Total Variation (TV) distance, Wasserstein distances (Durmus & Moulines, 2019), and Kullback-Leibler (KL) divergence (Dalalyan, 2017; Cheng & Bartlett, 2018; Ma et al., 2019; Freund et al., 2022), in different settings, such as strongly log-concave, log-Sobolev inequality (LSI) (Vempala & Wibisono, 2019), and Poincaré inequality (PI) (Chewi et al., 2021). Several works have achieved acceleration convergence for ULA by decreasing the discretization error with higher-order SDE, e.g., underdamped Langevin dynamics (Cheng et al., 2018; Ma et al., 2021; Mou et al., 2021), and aligning the discretized stationary distribution to the target $p_{*}$, e.g., Metropolis-adjusted Langevin algorithm (Dwivedi et al., 2018) and proximal samplers (Lee et al., 2021b; Chen et al., 2022b; b; Altschuler & Chewi, 2023). Liu & Wang (2016); Dong et al. (2023) also attempt to perform sampling tasks with deterministic algorithms whose limiting ODE is derived from the Langevin dynamics.

Regarding the convergence guarantees, most of these works have landscape assumptions for the target distribution $p_{*}$, e.g., strong log-concavity, LSI, or PI. For more general distributions, Erdogdu & Hosseinzadeh (2021) and Mousavi-Hosseini et al. (2023) consider KL convergence in modified LSI and weak PI. These extensions allow for slower tail-growth of negative log-density $f_{*}$, compared to the quadratic or even linear case. Although these works extend ULA to more general distributions, the computational burden of ill-conditioned isoperimetry still exists, e.g., exponentially dependent on the dimension (Raginsky et al., 2017). In this paper, we introduce another SDE to guide the sampling algorithm, which is non-Markovian and time-inhomogeneous. Our algorithm discretizes such an SDE and can reduce the isoperimetric and dimension dependence wrt TV convergence.

Diffusion Models and Stochastic Localization. In recent years, diffusion models have gained significant attention due to their ability to generate high-quality samples (Ho et al., 2020; Rombach et al., 2022). The core idea of diffusion models is to parameterize the score, i.e., the gradient of the log-density, during the entire forward OU process from the target distribution $p_{*}$. In this condition, the reverse SDE is associated with the inference process in diffusion models to perform unnormalized sampling for intricate target distributions. Apart from conventional MCMC trajectories, the most desirable property of this process is that, if the score function can be well approximated, it can sample from a general distribution Chen et al. (2022a); Lee et al. (2022; 2023). This implies the reverse SDE has a lower dependency on the isoperimetric requirement for the target distribution, which inspires the designs of new sampling algorithms. On the other hand, stochastic localization framework (Eldan, 2013; 2020; 2022; El Alaoui et al., 2022; Montanari, 2023) formalizes the general reverse diffusion and discusses the relationship with current diffusion models. Along this line, Montanari & Wu (2023) consider the unnormalized sampling for the symmetric spiked model. However, for the properties of the general unnormalized sampling, the investigation is limited.

This work employs the backward path of diffusion models to design a sampling strategy which we can prove to be more effective than the forward path of Langevin algorithm under suitable conditions. Several other recent studies (Vargas et al., 2023a; Berner et al., 2022; Zhang et al., 2023; Vargas et al., 2023c; b) have also utilized the diffusion path in their posterior samplers. Instead of the closed form MC sampling approach, the above works learns an approximate score function via a parametrized model, following the standard practice of the diffusion generative models. The resulting approximate distribution following the backward diffusion path is akin to the ones obtained from variational inference, and contains errors that depend on the expressiveness of the parametric models for the score functions. In addition, the convergence of the sampling process using such an approach depends on the generalization error of learned score functions (Tzen & Raginsky, 2019) (See Appendix for more details). It remains open how to bound sample complexity of such approaches, due to the need for bounding the error of the learned score function along the entire trajectory.

In contrast, rdMC proposed in this paper can be directly analyzed. Specifically, our algorithm draws samples from the unnormalized distribution without training data and the algorithm is designed to avoid learning based score function estimation. We are interested in the comparisons between rdMC and conventional MCMC algorithms. Thus, we analyze the complexity to estimate the score by drawing samples from the posterior distribution and found that our proposed algorithm is much more efficient in certain cases when the log-Sobolev constant of the target distribution is small.

Reverse SDE in Diffusion Models. In this part, we introduce the notations and formulation of the reverse SDE associated with the inference process in diffusion models, which will also be commonly used in the following sections. We begin with an OU process starting from $p_{*}$ formulated as

where $(B_{t})_{t\geq 0}$ is Brownian motion and $p_{0}$ is assigned as $p_{*}$. This is similar to the forward process of diffusion models (Song et al., 2020) whose stationary distribution is ${\mathcal{N}}\left({\bm{0}},{\bm{I}}\right)$ and we can track the random variable ${\bm{x}}$ at time $t$ with density function $p_{t}$.

According to Cattiaux et al. (2021), under mild conditions in the following forward process $\mathrm{d}{\bm{x}}_{t}=b_{t}({\bm{x}}_{t})\mathrm{d}t+\sqrt{2}\mathrm{d}B_{t},$ the reverse process also admits an SDE description. If we fix the terminal time $T>0$ and set $\tilde{{\bm{x}}}_{t}={\bm{x}}_{T-t},\ \mathrm{for}\ t\in[0,T],$ the process $(\tilde{{\bm{x}}}_{t})_{t\in[0,T]}$ satisfies the SDE $\mathrm{d}\tilde{{\bm{x}}}_{t}=\tilde{b}_{t}(\tilde{{\bm{x}}}_{t})\mathrm{d}t+\sqrt{2}\mathrm{d}B_{t},$ where the reverse drift satisfies the relation $b_{t}+\tilde{b}_{T-t}=2\nabla\ln p_{t},{\bm{x}}_{t}\sim p_{t}.$ In this condition, the reverse process of Eq. (1) is as $\mathrm{d}\tilde{{\bm{x}}}_{t}=\left(\tilde{{\bm{x}}}_{t}+2\nabla\ln p_{T-t}(\tilde{{\bm{x}}}_{t})\right)\mathrm{d}t+\sqrt{2}\mathrm{d}B_{t},\tilde{{\bm{x}}}_{t}\sim\tilde{p}_{t}.$ Thus, once the score function $\nabla\ln p_{T-t}$ is obtained, the reverse SDE induce a sampling algorithm.

Discretization and realization of the reverse process. To numerically solve the previous SDE, suppose $k\coloneqq\lfloor t/\eta\rfloor$ for any $t\in[0,T]$, we approximate the score function $\nabla\ln p_{T-t}$ with ${\bm{v}}_{k}$ for $t\in[k\eta,(k+1)\eta]$. This modification results in a new SDE, given by

Specifically, when $k=0$, we set $\tilde{{\bm{x}}}_{0}$ to be sampled from $\tilde{p}_{0}$, which can approximate $p_{T}$.To find the closed form of the solution by setting an auxiliary random variable as ${\mathbf{r}}_{i}\left(\tilde{{\bm{x}}}_{t},t\right)\coloneqq\tilde{{\bm{x}}}_{t,i}e^{-t}$, we have

where the equalities are derived by the Itô’s Lemma. Then, we set the initial value ${\mathbf{r}}_{0}(\tilde{{\bm{x}}}_{t},0)=\tilde{{\bm{x}}}_{t,i}$, and integral on both sides of the equation. We have

For the specific construction of ${\bm{v}}_{k}$ to approximate the score, we defer it to Section 3.

As shown in SDE (2), we introduce an estimator ${\bm{v}}_{k}\approx\nabla_{{\bm{x}}}\ln p_{T-t}({\bm{x}})$ to implement the reverse diffusion. This section is dedicated to exploring viable estimators and the benefits of the reverse diffusion process. We found that we can reformulate the score as an expectation with the transition kernel of the forward OU process. We also derive the intuition that rdMC can reduce the isoperimetric dependence of the target distribution compared with conventional Langevin algorithm. Lastly, we introduce the detailed implementation of rdMC in real practice with different score estimators.

In this section, we analyze the overall complexity of the rdMC via ULA inner loop estimation. Since the complexity of the importance sampling estimate is discussed in Section 3.3, we only consider Algorithm 1 with direct ULA sampling of $q_{T-t}(\cdot|{\bm{x}})$ rather than the smart importance sampling initialization to make our analysis clear.

To provide the analysis of the convergence of rdMC, we make the following assumptions.

1. [A₁]

   For all $t\geq 0$, the score $\nabla\ln p_{t}$ is $L$-Lipschitz.

2. [A₂]

   The second moment of the target distribution $p_{*}$ is upper bounded, i.e., $\mathbb{E}_{p_{*}}\left[\left\|\cdot\right\|^{2}\right]=m_{2}^{2}$.

These assumptions are standard in diffusion analysis to guarantee the convergence to the target distribution (Chen et al., 2022a). Specifically, Assumption [A₁] governs the smoothness characteristics of the forward process, which ensure the feasibility of numerical discretization methods used for approximating the solution of continuous SDE. In addition, Assumption [A₂] introduces essential constraints on the moments of the target distribution. These constraints effectively prevent an excessive accumulation of probability mass in the tail region, thereby ensuring a more balanced and well-distributed target distribution.