[Paper Review] Divide-and-Conquer Posterior Sampling for Denoising Diffusion Priors

Yazid Janati, Moufad, Badr+3 more|arXiv (Cornell University)|Mar 18, 2024

Seismic Imaging and Inversion TechniquesEarth and Planetary Sciences3 citations

TL;DR

This paper proposes Divide-and-Conquer Posterior Sampling (DCPS), a novel method for sampling from posterior distributions in Bayesian linear inverse problems using denoising diffusion model (DDM) priors. By decomposing the posterior sampling task into a sequence of simpler intermediate problems via a divide-and-conquer strategy and leveraging variational inference with Langevin dynamics, DCPS achieves lower approximation error and superior reconstruction quality across image restoration tasks without task-specific tuning.

ABSTRACT

Recent advancements in solving Bayesian inverse problems have spotlighted denoising diffusion models (DDMs) as effective priors. Although these have great potential, DDM priors yield complex posterior distributions that are challenging to sample. Existing approaches to posterior sampling in this context address this problem either by retraining model-specific components, leading to stiff and cumbersome methods, or by introducing approximations with uncontrolled errors that affect the accuracy of the produced samples. We present an innovative framework, divide-and-conquer posterior sampling, which leverages the inherent structure of DDMs to construct a sequence of intermediate posteriors that guide the produced samples to the target posterior. Our method significantly reduces the approximation error associated with current techniques without the need for retraining. We demonstrate the versatility and effectiveness of our approach for a wide range of Bayesian inverse problems. The code is available at \url{https://github.com/Badr-MOUFAD/dcps}

Motivation & Objective

• Address the challenge of posterior sampling in Bayesian inverse problems when using complex deep generative priors, particularly denoising diffusion models (DDMs).
• Overcome the limitations of existing methods that rely on approximations to the drift term, which often lead to high approximation error and sensitivity to hyperparameters.
• Develop a principled, scalable sampling framework that exploits the structure of DDM priors to decompose the posterior into a sequence of simpler, intermediate problems.
• Achieve high-quality reconstructions in image restoration tasks—such as super-resolution, inpainting, and colorization—without task-specific hyperparameter tuning.
• Provide a theoretically grounded approach that reduces approximation error by shortening the length of Feynman–Kac models through increased decomposition levels (L).

Proposed method

• Introduce a sequence of pseudo-posteriors $ q_{k_ au}^y $ forming a smooth path from a standard Gaussian to the target posterior, parameterized by a discretized time grid $ \tau = 0, \dots, L $.
• For each step $ \ell \in \llbracket 0, L-1 \rrbracket $, define an intermediate posterior $ q_{k_{\ell}}^y $ as a Feynman–Kac model with a non-homogeneous Markov transition kernel.
• Approximate the backward kernel of each intermediate distribution using a Gaussian variational approximation $ q_{k_{\ell+1}}^{y,\ell} $, initialized from a sample of $ q_{k_{\ell+1}}^y $.
• Sample from $ q_{k_{\ell+1}}^{y,\ell} $ via an unadjusted Langevin algorithm with $ G $ steps, ensuring efficient exploration of the intermediate posterior.
• Use a backward sampling strategy that combines Langevin dynamics and non-homogeneous Markov chain transitions to propagate samples from $ q_{k_{\ell+1}}^y $ to $ q_{k_{\ell}}^y $, minimizing approximation error.
• Leverage theoretical results showing that reducing the length $ k_{\ell+1} - k_\ell $ (i.e., increasing $ L $) reduces Gaussian approximation error, enabling high-fidelity posterior sampling.

Figure 2: Illustration of DCPS with $L=3$ . Red dots refer to the distribution path ( 4.1 ) and $(\mathrm{L})$ denotes Langevin steps targeting $\smash{p^{y,\ell-1}_{{k_{\ell}}}}$ initialized at approximate samples from $\smash{p^{y}_{{k_{\ell}}}}$ . Arrows starting from $\smash{p^{y,\ell-1}_{{k_{\e

Experimental results

Research questions

• RQ1 Can a divide-and-conquer strategy reduce the approximation error in posterior sampling for DDM priors compared to existing drift-approximation methods?
• RQ2 How does the length of intermediate Feynman–Kac models affect the accuracy of the posterior approximation in the proposed framework?
• RQ3 Can the proposed method achieve high-quality reconstructions in image restoration tasks without task-specific hyperparameter tuning?
• RQ4 How does the method compare in terms of sample diversity and consistency with observed image regions to existing approaches like DPS and DDRM?
• RQ5 What is the impact of the variational approximation and Langevin dynamics on the mixing and convergence of the sampling chain?

Key findings

• DCPS achieves state-of-the-art reconstruction quality in image super-resolution (4× and 16×), outpainting, inpainting, and colorization tasks, with sharp, coherent, and diverse samples.
• On the Gaussian mixture experiment, DCPS with $ L=50 $ achieved a SW distance of $ 1.95/2.70 \pm 2.28 $, significantly outperforming baseline methods in terms of posterior approximation accuracy.
• DCPS with $ L=500 $ achieved a SW distance of $ 1.26/2.59 \pm 2.83 $, demonstrating that increasing the number of decomposition levels further reduces approximation error.
• Unlike DPS, which often fails to maintain consistency with observed image regions and generates inconsistent samples, DCPS produces coherent reconstructions across all tasks.
• DCPS does not require task-specific or data-specific hyperparameter tuning, unlike DPS, which required extensive tuning across different datasets and tasks.
• DDRM consistently produced blurry, low-diversity samples, indicating mode collapse, whereas DCPS maintained high diversity and sharpness without any tuning.