Posterior Sampling by Combining Diffusion Models with Annealed Langevin Dynamics

AI summary

Overview

Research area: Machine learning theory — specifically provable algorithms for posterior (conditional) sampling with diffusion models and Langevin dynamics, with an empirical validation section on image inverse problems.

Technical level: Advanced. The paper is a theory contribution built on stochastic differential equations, score matching, strong log-concavity, total variation distance bounds, and polynomial-time complexity statements; the experiments are comparatively light.

Scope: The paper proves that combining a diffusion-model initialization with an annealed Langevin sampler yields polynomial-time posterior sampling for (globally or locally) log-concave priors using only an L⁴ bound on score-estimation error, and demonstrates the method on inpainting, super-resolution, and Gaussian deblurring with FFHQ–256 images.

What This Paper Is About

Given a noisy linear measurement y = Ax + ξ of a signal x drawn from a distribution p(x), and a good approximation to the prior p(x) in the form of a diffusion model, the goal is to sample from the posterior p(x|y). This matters because posterior sampling gives an accurate and fair framework for inpainting, deblurring, and MRI reconstruction, and is 2-competitive with the optimal reconstruction in any metric — but it is known to be computationally intractable in general, and existing diffusion-based heuristics such as DPS lack convergence guarantees. The paper asks: for which distributions p can we design a provably efficient and robust posterior sampling algorithm?

Key Contributions

1. A polynomial-time posterior sampling algorithm for globally log-concave priors. Theorem 1.1 shows that for an α-strongly log-concave p(x) over ℝ^d with L-Lipschitz score, if the smoothed-score error satisfies ε_score ≤ √α/K₁, then an algorithm (Algorithm 1) takes K₂ iterations to produce p̂(x|y) with E[TV(p̂(x|y), p(x|y))] ≤ ε, where K₁ = poly(d, m, ||A||/(η√α), 1/ε) and K₂ = poly(d, m, ||A||/(η√α), 1/ε, L/α). Only an L⁴ score-accuracy assumption is required.

2. Extension to local log-concavity via a Gaussian warm start. Theorem 1.2 assumes only that p is nearly log-concave inside a ball of radius R around a typical sample, and shows that given a Gaussian measurement x₀ = x + N(0, σ²I_d) with σ ≤ R/(K₁ + 2τ) and ε_score ≤ 1/(K₂σ), the algorithm takes K₃ iterations to reach total variation ≲ ε from p(x|x₀, y).

3. A competitive compressed-sensing corollary. Corollary 1.3 shows that if any (possibly exponential-time) algorithm can recover x to error r from y with probability 1−δ, and a naive algorithm (such as the LASSO) provides a warm start x₀ with error R, then under local log-concavity on a ball of radius R′ = R·poly(d, m, ||A||R/η, 1/δ), the proposed algorithm recovers error ≤ 2r in poly(d, m, ||A||R/η, 1/δ) time with probability 1−O(δ). The authors state this is the first known guarantee of this kind.

4. An annealing design and empirical validation. The paper introduces an "admissible" noise schedule η₁ > η₂ > … > η_N = η (Definition D.1) that keeps the Langevin process near the prior manifold, and validates the approach on three FFHQ–256 inverse problems where it can surpass DPS on both L² and FID in the inpainting setting given a small enough step size.

Main Findings

• L⁴ suffices where Langevin alone needs an MGF bound: Langevin dynamics with exact scores samples the posterior, but is brittle to score error and requires a sub-exponential (MGF) error bound; diffusion models succeed unconditionally with only an L² bound. The paper proves that diffusion initialization plus annealed Langevin achieves conditional sampling with merely an L⁴ score-error bound.

• Annealing is necessary: Appendix F explains why standard (non-annealed) Langevin dynamics would not suffice in this setting. Figure 3 illustrates this with p = N(0,1) and y = x + N(0,0.01): averaging over y, the marginal of X_t stays Gaussian, its variance first contracts and then returns toward the prior, so at an intermediate time t* the process has a constant factor lower variance — in high dimension an exponentially small region of p — meaning an Lᵖ bound under p does not control error under X_{t*}.

• Local log-concavity is enough with a warm start: Theorem 1.2 only needs the Hessian bound −LI_d ⪯ ∇²log p(x) ⪯ (τ²/R²)I_d to hold on balls of radius R with probability ≥ 1−ε. Figure 1 gives the motivating example: the uniform distribution on the unit circle convolved with N(0, w²I₂), whose λ_max(∇²log p(x)) reaches Ω(1/w⁴) near the center but is well behaved in a w-width band.

• The hardness of posterior sampling localizes: The paper's discussion of the cryptographic lower bound of [18] uses the instance p := q * N(0, σ²I_d) with q a mixture over 2^{d/2} points (s, f(s)) for a one-way permutation f, observed through y = Ax + N(0, η²I_d) with A = (0 I_{d/2}) and σ ≪ η ≪ d^{−1/2}. Sampling from p(x|y) would reveal f^{−1}(rnd(y)) in exp(Ω(d)) time, but a coarse estimate x₀ with ||x₀ − x|| < 1/3 uniquely identifies the correct component, reducing the task to sampling a Gaussian. The authors conclude the hardness may lie only in localizing x in the exponentially large space.

• Signal-to-noise interpretation: The quantity ||A||/(η√α) appearing in Theorem 1.1 is described as the signal-to-noise ratio of the measurement, and ||A||σ/η plays the analogous role in Theorem 1.2.

• Empirical trade-off on FFHQ–256: Across inpainting, 4× super-resolution, and Gaussian deblurring, increasing the time devoted to annealed Langevin decreases the per-image L² distance to ground truth but increases FID. In inpainting, with a sufficiently small step size, the method surpasses DPS on both metrics. Reconstructions qualitatively better preserve ground-truth attributes (Figures 5 and 6).

Methodology in Plain English

The researchers start from the observation that the unsmoothed conditional score is easy to write down by Bayes' rule: ∇ₓ log p(x|y) = s(x) + Aᵀ(y − Ax)/η². That means Langevin dynamics is naturally suited to posterior sampling — it adapts to conditioning at no extra cost — but it needs very reliable scores, because it can wander into regions where the prior score estimate is bad.

Their fix has two parts. First, they use the unconditional diffusion model to draw a starting sample from the prior, which by construction lies on the data manifold. Second, instead of running one long Langevin chain from the prior to the posterior, they build a sequence of intermediate, progressively less noisy measurement problems y₁, …, y_N = y, coupled so that y_i is a noisier version of y_{i+1}. Each consecutive pair of posteriors p(x|y_i) and p(x|y_{i+1}) is statistically close, so a short Langevin run suffices to move between them. Because each run is short, the process never drifts far from the distribution it started at, which keeps every score evaluation in a region where the L⁴ accuracy assumption actually holds. Small per-step errors then accumulate in a controlled way across the N steps.

For the locally log-concave case, they add a third ingredient: given a coarse estimate x₀, they treat x₀ as a spherical Gaussian measurement of x, sample from p(x|x₀) with the diffusion SDE (which is robust and efficient for spherical Gaussian measurements), and use that sample as the initialization for the annealed Langevin sampler targeting p(x|y).

The proof is organized across appendices: Appendix A handles convergence between consecutive posteriors for globally strongly log-concave p, Appendix B extends it to the local case, Appendix C bounds score-error and discretization error, Appendix D designs the noise schedule, and Appendix E assembles the main theorems.

Why This Matters

Impact on research. The paper closes part of a gap between two literatures. Unconditional diffusion sampling is provably efficient under an L² score bound; posterior sampling is provably hard in general under cryptographic assumptions; Langevin dynamics handles conditioning gracefully but needs unrealistically strong score guarantees. This work identifies a middle regime — log-concave or locally log-concave priors plus a warm start — where efficient, robust conditional sampling is possible with only an L⁴ bound, and it supplies the first compressed-sensing guarantee of this "go from a decent warm start to near-optimal reconstruction" form. It also gives theoretical guidance on two questions prior annealed-Langevin heuristics left open: how to design the annealing schedule and why annealing improves robustness.

Real-world applications (drawn from the settings the paper names):

• MRI reconstruction — the paper's motivating example for local log-concavity, where a compressed-sensing reconstruction serves as x₀.
• Image inpainting — evaluated on FFHQ–256 with masks covering 30%–70% of pixels uniformly at random.
• Image deblurring — evaluated with a 61×61 Gaussian kernel of standard deviation 3.0.
• Super-resolution — evaluated with a 4× downsampling factor.
• Fairness-sensitive reconstruction — the paper notes posterior sampling satisfies fairness guarantees with respect to protected classes [22], so a tractable posterior sampler inherits that property.

Industry relevance. The pipeline is directly compatible with existing pretrained diffusion models: the experiments reuse the pretrained model from [12] and simply refine DPS initial reconstructions, so the annealing component can be bolted onto deployed diffusion-based restoration systems rather than requiring retraining. The results also expose a practical tuning trade-off — more annealing time helps L² but hurts FID, and inpainting is the setting where the method dominates DPS on both — which is the kind of guidance practitioners need when allocating compute.

Future Directions

1. Weakening the log-concavity assumption further. The paper is explicit that it views this as "a first step"; the gap between the locally log-concave regime it handles and fully general distributions remains open, and the cryptographic lower bound of [18] shows that some assumption or side information is unavoidable.

2. Designing annealing schedules in practice. Definition D.1 characterizes admissible schedules for the analysis, but the experiments instead "sweep the step size while keeping the annealing schedule fixed," leaving the question of how to choose or optimize the schedule for real workloads.

3. Resolving the L²-versus-FID trade-off. Increasing annealed Langevin time decreases L² but increases FID in all three evaluated tasks; understanding and mitigating this tension — and characterizing when the method beats DPS on both metrics — is a concrete open empirical question.

4. Broadening the measurement model. The theory is developed for linear Gaussian measurements y = Ax + ξ with ξ ~ N(0, η²I_m); extending the guarantees to non-linear forward operators is not addressed.

Target Audience

Theoretical machine learning researchers working on sampling, diffusion models, and inverse problems; statisticians interested in Langevin dynamics and log-concave sampling; and applied researchers in medical imaging or computational photography who want to understand when diffusion-based

Authors’ abstract

Given a noisy linear measurement $y = Ax + ξ$ of a distribution $p(x)$, and a good approximation to the prior $p(x)$, when can we sample from the posterior $p(x \mid y)$? Posterior sampling provides an accurate and fair framework for tasks such as inpainting, deblurring, and MRI reconstruction, and several heuristics attempt to approximate it. Unfortunately, approximate posterior sampling is computationally intractable in general. To sidestep this hardness, we focus on (local or global) log-concave distributions $p(x)$. In this regime, Langevin dynamics yields posterior samples when the exact scores of $p(x)$ are available, but it is brittle to score--estimation error, requiring an MGF bound (sub-exponential error). By contrast, in the unconditional setting, diffusion models succeed with only an $L^2$ bound on the score error. We prove that combining diffusion models with an annealed variant of Langevin dynamics achieves conditional sampling in polynomial time using merely an $L^4$ bound on the score error.