Posterior Sampling by Combining Diffusion Models with Annealed Langevin Dynamics

Introduction

Diffusion models are currently the leading approach to generative modeling of images. Diffusion models are based on learning the “smoothed scores” $s_{\sigma^{2}}(x)$ of the modeled distribution $p(x)$. Such scores can be approximated from samples of $p(x)$ by optimizing the score matching objective ; and given good $L^{2}$-approximations to the scores, $p(x)$ can be efficiently sampled using an SDE or an ODE .

Much of the promise of generative modeling lies in the prospect of applying the modeled $p(x)$ as a prior: combining it with some other information $y$ to perform a search over the manifold of plausible images. Many applications, including MRI reconstruction, deblurring, and inpainting, can be formulated as linear measurements

Researchers have developed a number of heuristics to approximate posterior sampling using the smoothed scores, including DPS , particle filtering methods , DiffPIR , and second-order approximations . Unfortunately, unlike for unconditional sampling, these methods do not converge efficiently and robustly to the posterior distribution. In fact, a lower bound shows that no algorithm exists for efficient and robust posterior sampling in general . But the lower bound uses an adversarial, bizarre distribution $p(x)$ based on one-way functions; actual image manifolds are likely much better behaved. Can we find an algorithm for provably efficient, robust posterior sampling for relatively nice distributions $p$? That is the goal of this paper: we describe conditions on $p$ under which efficient, robust posterior sampling is possible.

A close relative to diffusion model sampling is Langevin dynamics, which is a different method for sampling that uses an SDE involving the unsmoothed score $s_{0}$. Unlike diffusion, Langevin dynamics is in general slow and not robust to errors in approximating the score. To be efficient, Langevin dynamics needs stronger conditions, like that $p(x)$ is log-concave and that the score estimation error satisfies an MGF bound (meaning that large errors are exponentially unlikely).

However, Langevin dynamics adapts very well to posterior sampling: it works for posterior sampling under exactly the same conditions as it does for unconditional sampling. The difference from diffusion models is that the unsmoothed conditional score $s_{0}(x\mid y)$ can be computed from the unconditional score $s_{0}(x)$ and the explicit measurement model $p(y\mid x)$, while the smoothed conditional score (which diffusion needs) cannot be easily computed.

So the current state is: diffusion models are efficient and robust for unconditional sampling, but essentially always inaccurate or inefficient for posterior sampling. No algorithm for posterior sampling is efficient and robust in general. Langevin dynamics is efficient for log-concave distributions, but still not robust. Can we make a robust algorithm for this case?

Can we do posterior sampling with log-concave $p(x)$ and $L^{p}$-accurate scores?

Our first result answers this in the affirmative. Algorithm˜1 uses a diffusion model for initialization, followed by an annealed version of Langevin dynamics, to do posterior sampling for log-concave $p(x)$ with just $L^{4}$-accurate scores. Annealing is necessary here; see Appendix˜F for why standard Langevin dynamics would not suffice in this setting.

The score estimates $\widehat{s}_{\sigma^{2}}(x)$ of the smoothed distributions $p_{\sigma^{2}}(x)=p(x)*\mathcal{N}(0,\sigma^{2}I_{d})$ have finite $L^{4}$ error, i.e.,

For precise bounds on the polynomials, see Theorem˜E.6. To understand the parameters, $\frac{\|A\|}{\eta\sqrt{\alpha}}$ should be viewed as the signal-to-noise ratio of the measurement.

Global log-concavity, as required by Theorem 1.1, is simple to state but a fairly strong condition. In fact, Algorithm 1 only needs a local log-concavity condition.

As motivation, consider MRI reconstruction. Given the MRI measurement $y$ of $x$, we would like to get as accurate an estimate $\widehat{x}$ of $x$ as possible. We expect the image distribution $p(x)$ to concentrate around a low-dimensional manifold. We also know that existing compressed sensing methods (e.g., the LASSO ) can give a fairly accurate reconstruction $x_{0}$; not as accurate as we are hoping to achieve with the full power of our diffusion model for $p(x)$, but still pretty good. Then conditioned on $x_{0}$, we know basically where $x$ lies on the manifold; if the manifold is well behaved, we only really need to do posterior sampling on a single branch of the manifold. The posterior distribution on this branch can be log-concave even when the overall $p(x)$ is not.

In the theorem below, we suppose we are given a Gaussian measurement $x_{0}=x+\mathcal{N}(0,\sigma^{2}I_{d})$ for some $\sigma$, and that the distribution $p$ is nearly log-concave in a ball polynomially larger than $\sigma$. We can then converge to $p(x\mid x_{0},y)$.

Then, there exist $K_{1},K_{2}=\operatorname*{poly}(d,m,\frac{\|A\|\sigma}{\eta},\frac{1}{\varepsilon})$ and $K_{3}=\operatorname*{poly}(d,m,\frac{\|A\|\sigma}{\eta},\frac{1}{\varepsilon},{L}\sigma^{2})$ such that: Given a Gaussian measurement $x_{0}=x+\mathcal{N}(0,\sigma^{2}I_{d})$ of $x\sim p$ with $\sigma\leq R/(K_{1}+2\tau)$. If $\varepsilon_{\text{score}}\leq\frac{1}{K_{2}\sigma}$, then there exists an algorithm that takes $K_{3}$ iterations to sample from a distribution $\widehat{p}(x\mid x_{0},y)$ such that

If $p$ is globally log-concave, we can set $\sigma=\infty$ so $x_{0}$ is independent of $x$ and recover Theorem 1.1; but if we have local information then this just needs local log-concavity. For precise bounds and a detailed discussion of the algorithm, see Section˜E.2.

The largest eigenvalue of $\nabla^{2}\log p(x)$ quantifies the extent to which the distribution departs from log-concavity at a given point. In Figure 1, we show an instance of a locally nearly log-concave distribution: $x$ is uniformly on the unit circle plus $\mathcal{N}(0,w^{2}I_{2})$. This distribution is very far from globally log-concave, but it is nearly log-concave within a $w$-width band of the unit circle. See Section˜E.4 for details.

Compressed Sensing.

In compressed sensing, one would like to estimate $x$ as accurately as possible from $y$. There are many algorithms under many different structural assumptions on $x$, most notably the LASSO if $x$ is known to be approximately sparse . The LASSO does not use much information about the structure of $p(x)$, and one can hope for significant improvements when $p(x)$ is known. Posterior sampling is known to be near-optimal for compressed sensing: if any algorithm achieves $r$ error with probability $1-\delta$, then posterior sampling achieves at most $2r$ error with probability $1-2\delta$. But, as we discuss above, posterior sampling cannot be efficiently computed in general.

We can use Theorem 1.2 to construct a competitive compressed sensing algorithm under a “local” log-concavity condition on $p$. Suppose we have a naive compressed sensing algorithm (e.g., the LASSO) that recovers the true $x$ to within $R$ error; and $p$ is usually log-concave within an $R\cdot\operatorname*{poly}$ ball; then if any exponential time algorithm can get $r$ error from $y$, our algorithm gets $2r$ error in polynomial time.

Consider attempting to accurately reconstruct $x$ from $y=Ax+\xi$. Suppose that:

Information theoretically (but possibly requiring exponential time or using exact knowledge of $p(x)$), it is possible to recover $\widehat{x}$ from $y$ satisfying $\left\lVert\widehat{x}-x\right\rVert\leq r$ with probability $1-\delta$ over $x\sim p$ and $y$.

We have access to a “naive” algorithm that recovers $x_{0}$ from $y$ satisfying $\left\lVert x_{0}-x\right\rVert\leq R$ with probability $1-\delta$ over $x\sim p$ and $y$.

For $R^{\prime}=R\cdot\operatorname*{poly}(d,m,\frac{\left\lVert A\right\rVert R}{\eta},\frac{1}{\delta})$,

Then we give an algorithm that recovers $\widehat{x}$ satisfying $\left\lVert\widehat{x}-x\right\rVert\leq 2r$ with probability $1-O(\delta)$, in $\operatorname*{poly}(d,m,\frac{\left\lVert A\right\rVert R}{\eta},\frac{1}{\delta})$ time, under Assumption 1 with $\varepsilon_{\text{score}}<\frac{1}{\operatorname*{poly}(d,m,\frac{\left\lVert A\right\rVert R}{\eta},\frac{1}{\delta},LR^{2})R}$.

That is, we can go from a decent warm start to a near-optimal reconstruction, so long as the distribution is locally log-concave, with radius of locality depending on how accurate our warm start is. To our knowledge this is the first known guarantee of this kind. Per the lower bound , such a guarantee would be impossible without any warm start or other assumption.

Figure 2 illustrates the sampling process of Corollary˜1.3. The initial estimate $x_{0}$ may lie well outside the bulk of $p(x)$; with just an $L^{4}$ error bound, the unsmoothed score at $x_{0}$ could be extremely bad. We add a bit of spherical Gaussian noise to $x_{0}$, then treat this as a spherical Gaussian measurement of $x$, i.e., $x+\mathcal{N}(0,RI)$; for spherical Gaussian measurements, the posterior $p(x\mid x_{0})$ can be sampled robustly and efficiently using the diffusion SDE. We take such a sample $x_{1}$, which now won’t be too far outside the distribution of $p(x)$, then use $x_{1}$ as initialization for annealed Langevin dynamics to sample from $p(x\mid y)$. The key part of our paper is that this process will never evaluate a score with respect to a distribution far from the distribution it was trained on, so the process is robust to error in the score estimates.

Notation and Background

There are several ways to sample from $p$ using the scores. Langevin dynamics is a classical MCMC method that considers the following overdamped Langevin Stochastic Differential Equation (SDE):

where $B_{t}$ is standard Brownian motion. The stationary distribution of this SDE is $p$, and discretized versions of it, such as the Unadjusted Langevin Algorithm (ULA), are known to converge rapidly to $p(x)$ when $p(x)$ is strongly log-concave . One can replace the true score $s(x)$ with an approximation $\widehat{s}$, as long as it satisfies a (fairly strong) MGF condition

In particular, showed that Langevin dynamics needs an MGF bound for convergence, and an $L^{p}$-accurate score estimator for any $1\leq p<\infty$ is insufficient.

An alternative approach, used by diffusion models, is to involve the smoothed scores. Starting from $x_{0}\sim\mathcal{N}(0,I_{d})$, one can follow a different SDE :

for a particular smoothing schedule $\sigma_{t}$; the result $x_{T}$ is exponentially close (in $T$) to being drawn from $p(x)$. This also has efficient discretizations , does not require log-concavity, and only requires an $L^{2}$ guarantee such as

to accurately sample from $p(x)$. One can also run a similar ODE with similar guarantees but faster .

Posterior sampling.

and want to sample from $p(x\mid y)$. The unsmoothed score $s_{y}(x):=\nabla_{x}\log p(x\mid y)$ is easily computed by Bayes’ rule:

Thus we can run the Langevin SDE (3) with the same properties: if $p(x\mid y)$ is strongly log-concave and the score estimate satisfies the MGF error bound (4), it will converge quickly and accurately.

Naturally, researchers have looked to diffusion processes for more general and robust posterior sampling methods. The main difficulty is that the smoothed score of the posterior involves $\nabla_{x}\log p(y\mid x_{\sigma_{t}^{2}})$ rather than the tractable unsmoothed term $\nabla_{x}\log p(y\mid x)$. Because the smoothed score is hard to evaluate exactly, a range of approximation techniques has been proposed . One prominent example is the DPS algorithm . Other methods include Monte Carlo/MCMC-inspired approximations , singular value decomposition and transport tilting , and schemes that combine corrector steps with standard diffusion updates . These approaches have shown strong empirical performance, and several provide guarantees under additional structure of the linear measurement; however, general guarantees for fast and robust posterior sampling remain limited beyond these restricted regimes.

Several recent studies use various annealed versions of the Langevin SDE as a key component in their diffusion-based posterior sampling method and achieve strong empirical results. Still, these methods provide no theoretical guidance on two key aspects: how to design the annealing schedule and why annealing improves robustness. None of these approaches come with correctness guarantees for the overall sampling procedure.

Comparison with Computational Lower Bounds.

One can efficiently obtain an $L^{p}$-accurate estimate of the smoothed score of $p$, so diffusion models can sample from $p$.

Any sub-exponential time algorithm that takes $y=Ax+\mathcal{N}(0,\eta^{2}I_{m})$ as input and outputs a sample from the posterior $p(x\mid y)$ fails on most $y$ with high probability.

To illustrate why the extra observation helps, consider the following simplified version of the hardness instance:

Here, $f:\{0,1\}^{d/2}\to\{0,1\}^{d/2}$ is a one‑way permutation — it takes exponential time to compute $f^{-1}(x)$ for most $x\in\{0,1\}^{d/2}$. $\delta(\cdot)$ is the Dirac delta function, and we choose $\sigma\ll d^{-1/2}$. Thus, $p(x)$ is a mixture of $2^{d/2}$ well‑separated Gaussians centered at the points $(s,f(s))$.

and let $\operatorname{rnd}(y)$ denote the vertex of $\{0,1\}^{d}$ closest to $y$. Then the posterior $p(x\mid y)$ is approximately a Gaussian centered at $(f^{-1}(\operatorname{rnd}(y)),\operatorname{rnd}(y))$ with covariance $\sigma^{2}I_{d}$. Generating a single sample would therefore reveal $f^{-1}(\operatorname{rnd}(y))$, which requires $\exp(\Omega(d))$ time.

However, suppose we have a coarse estimate $x_{0}$ satisfying $\|x_{0}-x\|<1/3$ (e.g., obtained by compressed sensing). Then, $x_{0}$ uniquely identifies the correct $(s,f(s))$ with $f(s)=\operatorname{rnd}(y)$, and the remaining task is just sampling from a Gaussian. Therefore, this hard instance becomes easy once we have localized the task and does not contradict our Theorem˜1.2.

We are able to handle the hard instance above well because it is exactly the type of distribution our approach is designed for: despite its complex global structure, it exhibits well-behaved local properties. This gives an important conceptual takeaway from our work: the hardness of posterior sampling may only lie in localizing $x$ within the exponentially large high-dimensional space.

Therefore, although posterior sampling is an intractable task in general, it is still possible to design a robust, provably correct posterior sampling algorithm — once we have localized the distribution. We view our work as a first step towards this goal.

Techniques

The algorithm we propose is clean and simple, but the proof is quite involved. Before we dive into the details, we provide a high-level overview of the intuitions behind the algorithm, concentrating on the illustrative case where the prior density $p(x)$ is $\alpha$-strongly log-concave. Under this assumption, every posterior density $p(x\mid y)$ is also $\alpha$-strongly log-concave. Therefore, posterior sampling could, in principle, be performed using classical Langevin dynamics.

The challenge arises because we lack access to the exact posterior score $s_{y}(x)$. We only possess an estimator derived from an estimate $\widehat{s}(x)$ of the prior score $s(x)$:

˜1 implies an $L^{4}$ accuracy of $\widehat{s}_{y}$ on average, but how do we use this to support Langevin dynamics, which demands exponentially decaying error tails?

Why can diffusion models succeed with merely $L^{2}$-accurate scores, whereas Langevin dynamics require MGF accuracy?

Both diffusion models and Langevin dynamics utilize SDEs. The $L^{2}$ error in the score-dependent drift term relates directly to the KL divergence between the true process (using $s(x)$) and the estimated process (using $\widehat{s}(x)$). Consequently, bounding the $L^{2}$ score error with respect to the current distribution $\widehat{p}_{t}$ controls the KL divergence.

Diffusion models leverage this property effectively. The forward process transforms data into a Gaussian, and the reverse generative process starts exactly from this Gaussian. At any time $t$, suppose $\widehat{p}_{t}$ is close to $p_{\sigma_{t}^{2}}$, then

by the $L^{2}$ accuracy assumption. This keeps the process close to the ideal process, ensuring overall small error.

Langevin dynamics, by contrast, often starts from an arbitrary, not predefined initial distribution $p_{\text{initial}}$. An $L^{p}$ score accuracy guarantee with respect to $p_{\text{target}}$ alone does not ensure accuracy for points $x_{t}$ that are not on the distributional manifold of $p_{\text{target}}$ (consider running Langevin starting from $x_{0}$ in Figure˜2). Therefore, a stronger MGF error bound is needed to prevent this from happening.

2 Adapting Langevin Dynamics for Posterior Sampling

While we can only use Langevin-type dynamics for posterior sampling, we possess a source of effective starting points: we can sample $x_{0}\sim p(x)$ efficiently using the unconditional diffusion model. Intuitively, $x_{0}$ already lies on the data manifold. The score estimator $\widehat{s}_{y}(x)$ initially satisfies:

As the dynamics evolves, the distribution $p(x_{t})$ transitions from $p(x)$ towards $p(x\mid y)$. If $x_{t}$ converges to $p(x\mid y)$, we again expect reasonable accuracy on average:

Hence the estimator is accurate at the start and at convergence. The open question concerns the intermediate segment of the trajectory: does $x_{t}$ wander into regions where the prior score $\widehat{s}(x)$ is unreliable? Ideally, the time-marginal of $x_{t}$, averaged over $y$, remains close to $p(x)$ throughout.

3 Annealing via Mixing Steps

In fact, even though $x_{0}$ and $x_{\infty}$ both have marginal $p(x)$, so the score estimate $\widehat{s}(x)$ is accurate on average at those times, this is not true at intermediate times. In Figure˜3, we illustrate this with a simple Gaussian example: $x_{0}$ and $x_{\infty}$ have distribution $\mathcal{N}(0,I)$ while $x_{t}$ has marginal $\mathcal{N}(0,cI)$ for a constant $c<1$. An $L^{p}$ error bound under $x\sim\mathcal{N}(0,I)$ does not give an $L^{2}$ error bound under $x\sim\mathcal{N}(0,cI)$, which means Langevin dynamics may not converge to the right distribution. A very strong accuracy guarantee like the MGF bound is needed here.

However, consider the case where the target posterior $p(x\mid y)$ is very close to the initial prior $p(x)$, such as when the measurement noise $\eta$ is very large (low signal-to-noise ratio). Langevin dynamics between close distributions typically converges rapidly. This suggests a key insight: if the required convergence time $T$ is short, the process $x_{t}$ might not deviate substantially from its initial distribution $p(x_{0})$. In such short-time regimes, an $L^{2}$ score error bound relative to $p(x_{0})$ could potentially suffice to control the dynamics. While $p(x)$ itself is already a good approximation for $p(x\mid y)$ when $\eta$ is very large, this motivates a general strategy.

Instead of a single, potentially long Langevin run from $p(x)$ to $p(x\mid y)$, we introduce an annealing scheme using multiple mixing steps. Given the measurement parameters $(A,\eta,y)$, we construct a decreasing noise schedule $\eta_{1}>\eta_{2}>\dots>\eta_{N}=\eta$. Correspondingly, we generate a sequence of auxiliary measurements $y_{1},y_{2},\dots,y_{N}=y$ such that each $y_{i}$ is distributed as $Ax+\mathcal{N}(0,\eta_{i}^{2}I_{m})$ and $y_{i}$ is appropriately coupled to $y_{i+1}$ (specifically, $y_{i}\sim\mathcal{N}(y_{i+1},(\eta_{i}^{2}-\eta_{i+1}^{2})I_{m})$ conditional on $y_{i+1}$). This creates a sequence of intermediate posterior distributions $p(x\mid y_{i})$.

An admissible schedule (formally defined in Definition˜D.1) ensures that:

$\eta_{1}$ is sufficiently large, making $p(x\mid y_{1})$ close to the prior $p(x)$.

Consecutive $\eta_{i}$ and $\eta_{i+1}$ are sufficiently close, making $p(x\mid y_{i})$ close to $p(x\mid y_{i+1})$.

Start with a sample $X_{0}\sim p(x)$. Since $\eta_{1}$ is large, $p(x)$ is close to $p(x\mid y_{1})$, so $X_{0}$ serves as an approximate sample $X_{1}\sim\widehat{p}(x\mid y_{1})$.

For $i=1$ to $N-1$: Run Langevin dynamics for a short time $T_{i}$, starting from the previous sample $X_{i}\sim\widehat{p}(x\mid y_{i})$, targeting the next posterior $p(x\mid y_{i+1})$ using the score $\widehat{s}_{y_{i+1}}(x)$. Let the result be $X_{i+1}\sim\widehat{p}(x\mid y_{i+1})$.

The final sample $X_{N}\sim\widehat{p}(x\mid y_{N})$ approximates a draw from the target posterior $p(x\mid y)$.

The core idea behind this annealing scheme is to actively control the process distribution $p(x_{t})$, ensuring it remains on the manifold of the prior $p(x)$. By design, each mixing step $i\to i+1$ connects two statistically close intermediate posteriors, $p(x\mid y_{i})$ and $p(x\mid y_{i+1})$. This closeness guarantees that a short Langevin run $T_{i}$ can mix them, and this short duration prevents $p(x_{t})$ from drifting significantly away from the step’s starting distribution $\widehat{p}(x\mid y_{i})$, and we can then argue that

This contrasts fundamentally with a single long Langevin run, where $x_{t}$ could venture far "off-manifold" into regions of poor score accuracy. By inserting frequent checkpoints that re-anchor the process, our annealing method substitutes such strong assumptions with structural control: the frequent “checkpoints” $p(x\mid y_{i})$ ensure the process is repeatedly localized to regions where the $L^{4}$ accuracy suffices. While error is incurred in each step, maintaining proximity to the manifold keeps this error small. The overall approach hinges on demonstrating that these small, per-step errors accumulate controllably across all $N$ steps.

This strategy, however, requires rigorous analysis of three key technical challenges:

How to bound the required convergence time $T_{i}$ for the transition from $p(x\mid y_{i})$ to $p(x\mid\penalty 10000\ y_{i+1})$? In particular, what happens when $p$ only has local strong log-concavity?

How to bound the error incurred during a single mixing step of duration $T_{i}$, given the $L^{4}$ score error assumption on the prior score estimate?

How to ensure the total error accumulated across all $N$ mixing steps remains small?

Addressing these questions forms the core of our proof.

In Appendix˜A, we show that for globally strongly log-concave distributions $p$, Langevin dynamics converges rapidly from $p(x\mid y_{i})$ to $p(x\mid y_{i+1})$. We extend this convergence analysis to locally strongly log-concave distributions in Appendix˜B. In Appendix˜C, we provide bounds on the errors incurred by score errors and discretization in Langevin dynamics. In Appendix˜D, we show how to design the noise schedule to control the accumulated error of the full process. In Appendix˜E, we conclude the analysis for Algorithm˜1, and apply it to establish the main theorems.

Experiments

To validate our theoretical analysis and assess real-world performance, we study three inverse problems on FFHQ–$256$ : inpainting, $4\times$ super-resolution, and Gaussian deblurring. Experiments use 1k validation images and the pre-trained diffusion model from . Forward operators are specified as in : inpainting masks $30\%\text{–}70\%$ of pixels uniformly at random; super-resolution downsamples by a factor of $4$; deblurring convolves the ground-truth with a Gaussian kernel of size $61\times 61$ (std. $3.0$). We first obtain initial reconstructions $x_{0}$ via Diffusion Posterior Sampling (DPS) , then refine them with our annealed Langevin sampler to draw samples close to $p(x\mid x_{0},y)$. To control runtime, we sweep the step size while keeping the annealing schedule fixed.

For each step size, we report the per-image $L^{2}$ distance to the ground truth and the FID of the resulting sample distribution (Figure 4). Across all three tasks, increasing the time devoted to annealed Langevin decreases $L^{2}$ but increases FID; in the inpainting setting, when the step size is sufficiently small, our method surpasses DPS on both metrics. Qualitatively, our reconstructions better preserve ground-truth attributes compared to DPS (Figures 5 and 6). All experiments were run on a cluster with four NVIDIA A100 GPUs and required roughly two hours per task.

Acknowledgments

This work is supported by the NSF AI Institute for Foundations of Machine Learning (IFML). ZX is supported by NSF Grant CCF-2312573 and a Simons Investigator Award (#409864, David Zuckerman).

References

Appendix A Langevin Convergence Between Strongly Log-concave Distributions

consider two random variables $y_{i}$ and $y_{i+1}$ defined as follows. First, draw $x\sim p$. Then, generate

This is the ideal (no discretization, no score estimation error) version of the process (2) that we actually run. Our goal is to establish the following lemma.

Suppose the prior distribution $p(x)$ is $\alpha$-strongly log-concave. Then, running the process (6) for time

In this section, our goal is to bound ${\chi^{2}\left(p(x\mid y_{i})\,\|\,p(x\mid y_{i+1})\right)}$. Since the posterior distributions can be expressed as

Let $Z_{1}=y_{i+1}-Ax$, and let $Z_{2}=y_{i}-Ax$. Then we have

where $f$ is the density function for $N(0,(\eta_{i}^{2}-\eta_{i+1}^{2})I_{m})$. Therefore,

Applying Markov’s inequality gives the result. ∎

Now we bound $\frac{p(y_{i+1})}{p(y_{i})}$. To make the lemma more self-contained, we abstract this a little bit.

where $p(Y_{1})$ and $p(Y_{2})$ are the densities of $Y_{1}$ and $Y_{2}$, respectively.

Write $Y_{1}-s=e_{1}$, and note that $Y_{2}-s=e_{1}+Z_{2}$. Then define

This gives that for any $Y_{1}$, $Y_{2}$, and $t$,

To bound $\sup_{\|e_{1}\|\leq t}G(e_{1})$, we expand $\phi$ as the $d$-dimensional Gaussian probability density function:

Using the quadratic expansion $\|e_{1}+Z_{2}\|^{2}=\|e_{1}\|^{2}+2\langle e_{1},Z_{2}\rangle+\|Z_{2}\|^{2}$, we rewrite:

Since $\|e_{1}\|\leq t$ and $\langle e_{1},Z_{2}\rangle\leq\|e_{1}\|\|Z_{2}\|$, we bound

Therefore, for any $Y_{1},Y_{2}$, and $t$, we have

Bounding expectation over Z2Z_{2}.

We can apply results on the Gaussian moment generating functions to bound this. Using Lemma˜A.10 by setting $\alpha=\frac{\eta_{2}^{2}}{2(\eta_{1}^{2}+\eta_{2}^{2})}$, $\beta=\frac{t\eta_{2}}{\eta_{1}^{2}+\eta_{2}^{2}}$, and $\gamma=\frac{\eta_{1}^{2}}{4(\eta_{1}^{2}+\eta_{2}^{2})}$, we have

where $p(Y_{1})$ and $p(Y_{2})$ are the densities of $Y_{1}$ and $Y_{2}$, respectively.

Let $t=(\sqrt{d}+\sqrt{2\ln(2\lambda)})\eta_{1}$. By applying Laurent-Massart bounds (Lemma˜A.11), we have

By applying Markov’s inequality, for a large enough constant $C>0$, we have

Cleaning up the bound a little bit, this implies that for a large enough constant $C>0$,

Combining this with the probability that $\left\lVert Z\right\rVert\leq t$, a union bound gives that

Now we can bound the $\chi^{2}$-diversity.

There exists a constant $C>0$ such that for any $\lambda>1$,

By Lemma˜A.5, there exists a constant $C>0$ such that

A union bound over these two implies that with probability of $1-1/\lambda$,

where $C^{\prime}$ is a positive constant. This concludes the lemma. ∎

A.2 Convergence time of Langevin dynamics

We present the following result on the convergence of Langevin dynamics:

Let $p$ and $q$ be probability distributions such that $q$ is an $\alpha$-strong log-concave distribution. Consider the Langevin dynamics initialized with $p$ as the starting distribution. Then, for any $t\geq 0$, we have

Let $p$ and $q$ be probability distributions such that $q$ is an $\alpha$-strong log-concave distribution. Consider the Langevin dynamics initialized with $p$ as the starting distribution. By running the diffusion for time

Now we show that the posterior distribution is even more strongly log-concave than prior distribution.

Suppose that $p(x)$ is $\alpha$-strongly log-concave. Then, the posterior density

By Bayes’ rule, the posterior density can be written (up to normalization) as

Since $p$ is $\alpha$-strongly log‑concave, its negative log–density satisfies

Moreover, the Gaussian likelihood term has

Hence $\varphi$ is $\alpha$-strongly convex, and the posterior density $p(x\mid Ax+N(\eta_{i}^{2}I_{m})=y_{i})\propto e^{-\varphi(x)}$ is $\alpha$-strongly log‑concave. ∎

By Lemma˜A.9, $p(x\mid y_{i+1})$ is $alpha$-strongly log-concave. This allows us to apply Lemma˜A.8. Therefore, to achieve $\varepsilon$ TV error in convergence, we only need to run the process for

Taking in the result in Lemma˜A.6, we have with $1-\frac{1}{\lambda}$ probability over $y_{i}$ and $y_{i+1}$, we only need

A.3 Utility Lemmas.

For all $r\geq 0$ and any $\gamma>0$, it is easy to check that by AM-GM inequality,

Taking $r=\|Z\|$ and exponentiating both sides, we obtain

Multiplying both sides by $\exp\Bigl(\alpha\|Z\|^{2}\Bigr)$ yields

For $Z\sim\mathcal{N}(0,I_{d})$ , when $\alpha+\gamma<\tfrac{1}{2}$ we have

Let $v\sim\mathcal{N}(0,I_{m})$. For any $t>0$,

Appendix B Convergence Between Locally Well-Conditioned Distributions

In the last section, we considered the convergence time between two posterior distributions of a globally strongly log-concave distribution. In this section, we will relax the assumption of global strong log-concavity and consider the convergence time between two distributions that are locally “well-behaved”. We give the following formal definition:

For $\delta\in[0,1)$ and $R,\widetilde{L},\alpha\in(0,+\infty]$, we say that a distribution $p$ is $(\delta,r,R,\widetilde{L},\alpha)$ mode-centered locally well-conditioned if there exists $\theta$ such that

$\Pr_{x\sim p}\left[x\in B(\theta,r)\right]\geq 1-\delta$.

For $x,y\in B(\theta,R)$, we have that $\|s(x)-s(y)\|\leq\widetilde{L}{\alpha}\left\lVert x-y\right\rVert$.

For $x,y\in B(\theta,R)$, we have that $\langle s(y)-s(x),x-y\rangle\geq\alpha\left\lVert x-y\right\rVert^{2}$.

Again, we consider the following process $P$, which is identical to process (6) we considered in the last section:

Our goal is to prove the following lemma:

Suppose $p$ is a $(\delta,r,R,\widetilde{L},\alpha)$ mode-centered locally well-conditioned distribution. Let $C>0$ be a large enough constant. We consider the process $P$ running for time

In this section, we will assume that $p$ is $(\delta,r,R,\widetilde{L},\alpha)$ mode-centered locally well-conditioned. Without loss of generality, we assume that the mode of $p$ is at 0, i.e., $\theta=0$.

We consider the process $P^{\prime}$ defined as the process $P$ conditioned on $x_{t}\in B(0,R)$ for $t\in[0,T]$.

Our goal is to prove the following lemma:

We start by decomposing the total variation distance between $P$ and $P^{\prime}$ as follows:

Recall that the process $P^{\prime}$ is defined as the law of $P$ conditioned on the event

where $\mathcal{F}^{c}=\{\exists\,t\in[0,T]:\,\|x_{t}\|\geq R\}$.

Let $\mathcal{E}:=\{x_{0}\in B(0,r)\}$ denote the event that the initial condition is “good.” Then, by the law of total probability,

Taking the expectation with respect to $y_{i}$ and $y_{i+1}$, we obtain

and by the law of total probability, we have

with $a\geq 0$. Then, for any time horizon $T>0$ and $\delta\in(0,1)$,

Define $r(t)=\|x_{t}\|$. Although the Euclidean norm is not smooth at the origin, an application of Itô’s formula yields that, for $x_{t}\neq 0$, one has

where $u(t)=x_{t}/\|x_{t}\|$. Using the bound $\|f(x_{t})\|\leq a$ and the hypothesis $\langle g(x_{t}),x_{t}\rangle\leq 0$, it follows by the Cauchy–Schwarz inequality that

Discarding the nonnegative Itô correction term $\frac{d-1}{\|x_{t}\|}\,dt$ (which can only increase the process), we deduce that

Since $\|u(s)\|=1$ for all $s$, the process $\beta(t)$ is a standard one-dimensional Brownian motion with quadratic variation $\langle\beta\rangle_{t}=t$. By a standard comparison theorem for one-dimensional stochastic differential equations, it follows that $r(t)\leq y(t)$ almost surely for all $t\geq 0$; hence,

A classical application of the reflection principle for one-dimensional Brownian motion shows that, for any $\rho>0$,

To incorporate the $d$-dimensional nature of the noise, one may use a union bound over the $d$ coordinate processes of $B_{t}$, which yields that

Combining the foregoing estimates, we deduce that

For any $\delta\in(0,1)$ and $T>0$, it holds that

Since $\|x\|\leq r$ with probability $1-\delta$. Thus, with probability at least $1-2\delta$, it follows that

Since the probability satisfying the condition is at least $1-2\delta$, we have

Putting Lemma˜B.4 and Lemma˜B.8 together, we directly obtain Lemma˜B.3.

B.2 Concentration of Strongly Log-Concave Distributions

Before moving futher, we first prove that a strongly log-concave distribution is highly concentrated.

is 1-Lipschitz (by the triangle inequality), it follows that

A standard calculation using the fact that the covariance matrix of $X$ satisfies $\operatorname{Cov}(X)\preceq\frac{1}{\alpha}I$ gives

Let $\mu$ and $\theta$ denote the mean and the mode of distribution $p$, respectively, where $p$ is $\alpha$-strongly log-concave and univariate. Then, $\left|\mu-\theta\right|\leq\frac{1}{\sqrt{\alpha}}$.

This immediately gives us the following corollary.

This also implies that every $\alpha$-strongly log-concave distribution is mode-centered locally well-conditioned.

Let $p$ be an $\alpha$-strongly log-concave distribution. Suppose the score function of $p$ is $L$-Lipschitz. Then, for any $0<\delta<1$, we have that $p$ is $(\delta,2\sqrt{\frac{d}{\alpha}}+\sqrt{\frac{2\log(1/\delta)}{\alpha}},\infty,L/\alpha,\alpha)$ mode-centered locally well-conditioned.

B.3 Convergence to Target Distribution

Since $p$ is not globally strongly log-concave, we need to extend the distribution $p$ to a globally strongly log-concave distribution. We will use the following lemma to extend the distribution.

For each fixed $z\in B(0,R)$ the mapping $\varphi_{z}$ has Hessian $-\alpha I_{d}$, hence is $\alpha$–strongly concave on the whole space. Because of (7) we have

Let $p$ be a $d$-dimensional $(\delta,r,R,\widetilde{L},\alpha)$ mode-centered locally well-conditioned probability distribution with $0<\delta\leq 1/2$ and $\alpha>0$. Assume

Since $\widetilde{p}=p/Z$ on $B$, we have

Therefore, $I_{B}\leq\frac{1}{2}\cdot 4\delta=2\delta$.

Now, we can consider process $\widetilde{P}$ defined as

Because $s(x)=\nabla\!\log\widetilde{p}(x)$ for every $x\in B(0,R)$, the drift coefficients of $P$ and $\widetilde{P}$ coincide on the event $\mathcal{E}$, and hence conditioning on $\mathcal{E}$ gives $P^{\prime}=\widetilde{P}^{\prime}$.

Taking expectation over $(y_{i},y_{i+1})$ gives

We start by considering another process $\widetilde{P}^{s}$ defined as

Combining this with Lemma˜B.15, we have that

Furthermore, by Lemma˜A.1 and our constraint on $T$, we have that

Appendix C Control of Score Approximation and Discretization Errors

In this section, we consider these processes running for time $T$:

Process $\widehat{P}$: Let $0=t_{1}<\dots<t_{M}=T$ be the $M$ discretization steps with step size $t_{j+1}-t_{j}=h$. For $t\in[t_{j},t_{j+1}]$,

Note that $\widehat{P}$ is exactly the process (2) we run in Algorithm˜1, except that we start from $x_{0}\sim p(x\mid y_{i})$.

We have shown that the process $P$ will converge to the target distribution $p(x\mid y_{i+1})$. We will show that the process $\widehat{P}$ will also converge to $p(x\mid y_{i+1})$ with a small error

Let $p$ be a $(\delta,r,R,\widetilde{L},\alpha)$ mode-centered locally well-conditioned. Suppose the followings hold for a large enough constant $C>0$:

$T>C\left(\frac{m\gamma_{i}+\log(\lambda/\varepsilon)}{\alpha}\right)$.

$\|A\|^{4}(T^{2}m+TR^{2})\leq\frac{\eta_{i}^{4}}{C\gamma_{i}^{2}}$.

$R\geq r+\frac{T\|A\|}{\eta_{i+1}^{2}}\Bigl(\|A\|r+\eta_{i+1}\bigl(\sqrt{m}+\sqrt{2\ln(1/\delta)}\bigr)\Bigr)+2\sqrt{dT\ln(2d/\delta)}$.

Then running $\widehat{P}$ for time $T$ guarantees that with probability at least $1-1/\lambda$ over $y_{i}$ and $y_{i+1}$, we have:

In this section, we assume $p$ is $(\delta,r,R,\widetilde{L},\alpha)$ mode-centered locally well-conditioned. Without loss of generality, we assume that the mode of $p$ is at 0, i.e., $\theta=0$. Let $L:=\widetilde{L}\alpha$, i.e., the Lipschitz constant inside the ball $B(0,R)$.

We will also consider the following stochastic processes:

Process $Q^{\prime}$ is the process $Q$ conditioned on $x_{t}\in B(0,R)$ for $t\in[0,T]$.

Process $P^{\prime}$ is the process $P$ conditioned on $x_{t}\in B(0,R)$ for $t\in[0,T]$.

Suppose $\|A\|^{4}(T^{2}m+TR^{2})\leq\frac{\eta_{i}^{4}\eta_{i+1}^{4}}{C(\eta_{i}^{2}-\eta_{i+1}^{2})^{2}}$ for a large enough constant $C$.

By Girsanov’s theorem, for any trajectory $x_{0,\dots,t}$,

where the Girsanov exponent $M_{t}$ is given by

Since $Q^{\prime}$ is supported in $B(0,R)$,

Now, for $\zeta_{y}:=\int_{0}^{t}\|\Delta b_{y}(x_{u})\|^{2}du$, we have that $M_{t}\sim\mathcal{N}\left(-\frac{1}{4}\zeta_{y},\frac{1}{2}\zeta_{y}\right)$

Note that $\|\eta_{i}^{2}y_{i+1}-\eta_{i+1}^{2}y_{i}\|^{2}$ has mean $\|(\eta_{i}^{2}-\eta_{i+1}^{2})Ax\|^{2}$ and is subgamma with variance $m\left(\eta_{i+1}^{2}\eta_{i}^{4}-\eta_{i+1}^{4}\eta_{i}^{2}\right)^{2}$ and scale $\eta_{i+1}^{2}\eta_{i}^{4}-\eta_{i+1}^{4}\eta_{i}^{2}$. Thus, for $t\|A\|^{2}\leq\frac{\eta_{i+1}^{2}\eta_{i}^{2}}{C\left(\eta_{i}^{2}-\eta_{i+1}^{2}\right)}$ we have

The first term can be bounded using Lemma˜C.4. Now we focus on the second term. Note that

Since $s$ is $L$-Lipschitz in $B(0,R)$, and using Lemma˜C.3, we have

where the last line follows from Markov’s inequality. The gives the result. ∎

We note that by our definition of $\gamma_{i}$,

Then, combining Corollary C.6 with Lemmas B.3 and C.2, we have

The conditions in Lemmas B.3 and C.2 are satisfied by our assumptions, noting that $\eta_{i+1}<\eta_{i}$ implies the bound on $R$ holds for both processes.

Applying Markov’s inequality and combining Lemma˜B.2 with the above, we conclude the proof. ∎

Appendix D Admissible Noise Schedule

Recall that we can define process $\widehat{P}_{i}$ that converges from $p(x\mid y_{i})$ to $p(x\mid y_{i+1})$: Let $0=t_{1}<\dots<t_{M}=T$ be the $M$ discretization steps with step size $t_{j+1}-t_{j}=h$. For $t\in[t_{j},t_{j+1}]$,

We have already proven that we can converge the process from $p(x\mid y_{i})$ to $p(x\mid y_{i+1})$ with good probability, as long as some conditions are satisfied. Those conditions actually depend on the choice of the schedule of $\eta_{i}$ and $T_{i}$. In this section, we will specify the schedule of $\eta_{i}$ and $T_{i}$.

Now we specify the schedule of $\eta_{i}$ and $T_{i}$.

We say a noise schedule $\eta_{1}>\dotsb>\eta_{N}$ together with running times $T_{1},\dotsb,T_{N-1}$ is admissible (for a set of parameters $C,\alpha,\lambda,A,d,\varepsilon,\eta,R$) if:

$\eta_{1}\geq\frac{\lambda\|A\|}{\varepsilon}\sqrt{\frac{d}{\alpha}}$;

For all $\gamma_{i}=(\eta_{i}/\eta_{i+1})^{2}-1$, we have $\gamma_{i}\leq 1$ and

The reason we need to satisfy the last inequality is to satisfy the conditions in Lemma˜C.1. We formalize this in the following lemma.

Let $C>0$ be a sufficiently large constant and $p$ be a $(\delta,r,R,\widetilde{L},\alpha)$ mode-centered locally well-conditioned distribution. For any $\delta,\varepsilon\in(0,1)$ and $\lambda>1$, suppose

For any admissible schedule $(\eta_{i})_{i\in[N]}$ and $(T_{i})_{i\in[N-1]}$, running the process $\widehat{P}_{i}$ for time $T_{i}$ guarantees that with probability at least $1-1/\lambda$ over $y_{i}$ and $y_{i+1}$:

It is straightforward to verify that an admissible schedule satisfies the first two conditions of Lemma˜C.1.

For the third condition regarding $R$, our assumption states:

Given that $T_{i}\lesssim\frac{m+\log(\lambda/\varepsilon)}{\alpha}$, this choice of $R$ is sufficient to satisfy the third condition in Lemma C.1.

Therefore, applying Lemma C.1 at each step $i$, we obtain that with probability at least $1-1/\lambda$ over $y_{i}$ and $y_{i+1}$:

We also want to prove the following two lemmas:

Let $p$ be a $d$-dimensional $(\delta,r,R,\widetilde{L},\alpha)$ mode-centered locally well-conditioned distribution. For any $\delta\in(0,1)$, suppose

Then, suppose $\eta_{1}\geq\frac{\lambda\|A\|}{\varepsilon}\sqrt{\frac{d}{\alpha}}$, with probability at least $1-\frac{1}{\lambda}$ over $y_{1}$,

There exists an admissible noise such that

where $\rho=\frac{\|A\|}{\eta\sqrt{\alpha}}$.

In this part, we prove Lemma˜D.3, showing that any admissible schedule has a large enough $\eta_{1}$, enabling us to use $p(x)$ to approximate $p(x\mid y_{1})$.

We have the following standard information-theoretic result.

Let $p$ be a $d$-dimensional $(\delta,r,R,\widetilde{L},\alpha)$ mode-centered locally well–conditioned probability distribution. Assume

Lemma B.14 provides an $\alpha$-strongly log–concave density $\widetilde{p}$ satisfying

Applying Lemma˜D.6 to $\widetilde{p}$ gives

Integrating in $y_{1}$ and using the elementary fact

together with the above calculaion, yields

Since all admissible noise schedules satisfy $\eta_{1}\geq\frac{\lambda\|A\|}{\varepsilon}\sqrt{\frac{d}{\alpha}}$. This implies

By Markov’s inequality, with probability at least $1-\frac{1}{\lambda}$ over $y_{1}$,

D.2 Bound for NN Mixing Steps

Let $a,x_{0}>0$, and let $c>0$. Consider the number sequence

For every $B>0$, let $k(B)$ be the minimum integer $i$ such that $x_{i}\geq B$. Then

We show in two steps that the time to go from $x_{0}$ to $1/a$, then to $B$. Define

We first show that $k_{1}\lesssim(ax_{0})^{-c}$. Consider the quantities

and let $j^{*}$ be the smallest $j$ such that $x_{N_{j}}\geq 1/a$. If instead $x_{0}\geq 1/a$ already, then $k_{1}=0$ and there is nothing to prove.

Assume $x_{0}<1/a$. For each $j<j^{*}$ define

By monotonicity of the sequence $(x_{i})$, it follows that $N_{j+1}\leq N_{j}+t_{j}$. Summing over $j$ up to $j^{*}-1$ gives

By definition, $N_{j^{*}}$ is the first index $i$ such that $x_{i}\geq 1/a$, so $k_{1}=N_{j^{*}}\lesssim(ax_{0})^{-c}$.

Bound to achieve BB.

If $B\leq 1/a$, the bound already holds. Now we analyze how many steps Note that for every $i\geq k_{1}$,

Given parameters $x_{0},a,b>0$, consider sequence inductively defined by $x_{i+1}=(1+\gamma_{i})x_{i}$, where

Given $B$, let $k(B)$ be the minimum integer $i$ such that $x_{i}\geq B$. Then,

Case 1: x0≥b2/ax_{0}\geq b^{2}/a.

We always choose $\gamma_{i}=\sqrt{x_{i}/a}$. We can verify that

and this satisfies the requirement for $\gamma_{i}$. By applying Lemma˜D.8, we have that

Case 2: x0≤B≤b2/ax_{0}\leq B\leq b^{2}/a.

We always choose $\gamma_{i}=\min(x_{i}/b,1)$. We can verify that

and this satisfies the requirement for $\gamma_{i}$. By applying Lemma˜D.8, we have that

Case 3: x0≤b2/a≤Bx_{0}\leq b^{2}/a\leq B.

We combine the bound for the first two cases, where we first go from $x_{0}$ to $b^{2}/a$, then go from $b^{2}/a$ to $B$. Then we have

Now we describe how we construct an admissible noise schedule. Consider we start from $\eta_{1}^{\prime}=\eta$, and for each $i$, we iteratively choose $\gamma_{i}^{\prime}$ to be the maximum $\gamma\leq 1$ such that

and then set $\eta_{i+1}^{\prime}=\sqrt{(1+\gamma_{i}^{\prime})(\eta^{\prime}_{i})^{2}}$. We continue this process until we reach $\eta_{N}^{\prime}\geq\frac{\lambda\|A\|}{\varepsilon}\sqrt{\frac{d}{\alpha}}$. It is easy to verify that $(\eta_{N}^{\prime},\eta_{N-1}^{\prime},\ldots,\eta_{1}^{\prime})$ is an admissible noise schedule. Now we bound the number of iterations $N$.

Since for all $\gamma$, we have $\|A\|^{4}(f_{T}^{2}(\gamma)m+f_{T}(\gamma)R^{2})\leq\|A\|^{4}(\sqrt{m}f_{T}(\gamma)+\frac{R^{2}}{2\sqrt{m}})^{2}$, a sufficient condition for $\|A\|^{4}(f_{T}^{2}(\gamma)m+f_{T}(\gamma)R^{2})\leq\frac{(\eta_{i}^{\prime})^{4}}{C\gamma^{2}}$ is that