Posterior sampling via Langevin dynamics based on generative priors

TL;DR

This work proposes efficient posterior sampling by simulating Langevin dynamics in the noise space of a pre-trained generative model, exploiting the mapping between the noise and data spaces which can be provided by distilled flows or consistency models, offering superior efficiency and performance compared to existing diffusion-based posterior sampling techniques.

Abstract

Posterior sampling in high-dimensional spaces using generative models holds significant promise for various applications, including but not limited to inverse problems and guided generation tasks. Despite many recent developments, generating diverse posterior samples remains a challenge, as existing methods require restarting the entire generative process for each new sample, making the procedure computationally expensive. In this work, we propose efficient posterior sampling by simulating Langevin dynamics in the noise space of a pre-trained generative model. By exploiting the mapping between the noise and data spaces which can be provided by distilled flows or consistency models, our method enables seamless exploration of the posterior without the need to re-run the full sampling chain, drastically reducing computational overhead. Theoretically, we prove a guarantee for the proposed noise-space Langevin dynamics to approximate the posterior, assuming that the generative model sufficiently approximates the prior distribution. Our framework is experimentally validated on image restoration tasks involving noisy linear and nonlinear forward operators applied to LSUN-Bedroom (256 x 256) and ImageNet (64 x 64) datasets. The results demonstrate that our approach generates high-fidelity samples with enhanced semantic diversity even under a limited number of function evaluations, offering superior efficiency and performance compared to existing diffusion-based posterior sampling techniques.

Figures (20)

• Figure 1: (Left) : A schematic representation of posterior sampling via Langevin dynamics in our proposed framework. The sampling process begins with an initial sample $x_1^{(0)}$ from the noise space and maps to data space as $x_0^{(0)}$ using a deterministic mapper $\Phi$ and progressively updates the noise space input to obtain diverse posterior samples. (Right): Posterior samples generated by our method and DPS-DM. Our approach exhibits higher perceptual diversity, capturing variations in high-level features such as lighting, window style, and wall patterns. Uncertain semantic features are highlighted by red boxes, while persistent properties are shown by green boxes.
• Figure 2: Reconstruction time comparison between DPS-DM and our method for varying numbers of posterior samples. DPS-DM scales poorly with the number of samples, while our method maintains a nearly constant time, demonstrating significantly lower computational cost. The corresponding Number of Function Evaluations (NFEs) (including NFEs for the warmup stage, refer to Section \ref{['sec:algo']}) values per image are annotated.
• Figure 3: Image reconstructions for the linear and nonlinear tasks on LSUN-Bedroom (256 x 256).
• Figure 4: Image reconstructions for the linear tasks on ImageNet (64 x 64).
• Figure 5: Posterior samples for the inpainitng (10%) (top three rows) and nonlinear deblur (bottom three rows) tasks on LSUN-Bedroom (256 x 256). Green boxes highlight low-uncertainty features and red boxes highlight highly uncertain features.

Theorems & Definitions (12)

• Example 3.1: Inverse problem with Gaussian noise
• Theorem 4.1: TV guarantee
• Remark 4.1: $\kappa_y$ as a condition number
• Example 4.1: Well-conditioned problem
• Example 4.2: Ill-conditioned problem
• Lemma 4.2: Sampling error
• Corollary 4.3: TV of sampled posterior
• Lemma A.1
• proof : Proof of Lemma \ref{['lemma:x1-sde-equilibrium']}
• proof : Proof of Theorem \ref{['thm:post-tv']}