Diffusion-Based Posterior Sampling Methods
- Diffusion-based posterior sampling methods are probabilistic inference algorithms that use deep diffusion models to approximate the posterior in inverse problems.
- They blend unconditional generative updates with soft likelihood gradients, employing surrogate calculations like Tweedie’s formula for both Gaussian and Poisson noise.
- This approach enables robust and scalable reconstruction in ill-posed and high-dimensional settings by operating entirely in the image domain.
Diffusion-based posterior sampling methods constitute a family of probabilistic inference algorithms leveraging score-based diffusion models for conditional generation and uncertainty quantification in inverse problems. These methods extend the application of diffusion models—originally developed for unconditional data synthesis—towards sampling from posteriors under noisy (linear or nonlinear) observation models, handling various measurement noise statistics, and enabling robust reconstruction in ill-posed or high-dimensional settings.
1. Theoretical Foundations and Problem Formulation
Diffusion models define a generative process in which a data variable is progressively corrupted by a stochastic differential equation (SDE) to produce a sequence of noisy states ; the generative model is learned to approximate the time-dependent score function using denoising score matching. The reverse-time SDE, parameterized with the estimated score function, enables sample generation from the learned data distribution.
In the context of inverse problems, one observes related to via a (potentially nonlinear and noisy) measurement process . The goal is to sample from the posterior , integrating the generative prior with the data likelihood. Direct computation of the conditional score is intractable for , since the likelihood acts only on .
To circumvent this, (Chung et al., 2022) proposes an approximation via Tweedie's formula, evaluating the likelihood at a posterior mean surrogate: given a DDPM/ADM-style discretization. The posterior gradient is then approximated as: enabling tractable posterior updates within the reverse diffusion steps.
2. Posterior Sampling Algorithm and Noise Models
The central scheme constitutes a blended update at each reverse diffusion step:
- Apply the unconditional score step from the pretrained model.
- Augment with a soft data-consistency gradient, i.e., .
For Gaussian noise, with forward model , ,
with step-size scaling set by . For Poisson noise, the likelihood for independent measurements at is: The update employs a coordinate-wise weighted norm, e.g.\
where models signal-dependent variance.
This blended update enables the method to robustly handle arbitrary noise statistics and nonlinear forward operators, by simply redefining the likelihood term and its gradient as appropriate for the measurement model.
3. Handling Nonlinear and Noisy Inverse Problems
A distinguishing feature of the DPS method (Chung et al., 2022) is its capacity to address general noisy inverse problems—both linear and nonlinear. Rather than imposing a hard projection to the measurement subspace (which can result in noise amplification and instability in ill-posed or high-noise regimes), the algorithm softly guides the generative trajectory via the measurement likelihood gradient.
The approach extends to settings where the measurement operator is a nonlinear mapping, provided that its Jacobian can be computed efficiently to evaluate the required gradient. This generality is demonstrated in the context of Fourier phase retrieval (nonlinear, non-Gaussian) and non-uniform deblurring.
4. Comparison with Prior Approaches
The DPS scheme diverges from earlier techniques in several important respects:
Approach | Measurement Enforcement | Domain | Applicability |
---|---|---|---|
MCG [Chung et al] | Hard projection | Image | Linear, noiseless |
DPS (this work) | Soft gradient blending | Image | Nonlinear, noisy (Gaussian, Poisson, etc) |
DDRM | Pseudo-inverse in spectral | Signal domain | Separable/linear, SVD required |
Unlike prior works that require SVDs of the measurement operator (limiting them to special cases or small-scale problems), DPS operates entirely in the image domain without such restrictions, and does not rely on strict hard constraints for measurement consistency, which can amplify error in practical, noisy settings.
DPS's robustness and flexibility are demonstrated both theoretically—via the developed surrogates for intractable likelihood terms—and empirically, through experiments on tasks such as denoising, inpainting, super-resolution, deblurring, and nonlinear inversion.
5. Implementation Details and Practical Considerations
DPS provides algorithmic variants for Gaussian (DPS-Gaussian) and Poisson (DPS-Poisson) noise and generalizes to arbitrary measurement processes as long as the likelihood gradient is feasible to compute. Core steps include:
- Evaluate using the pretrained network.
- Compute via Tweedie's formula.
- Compute the measurement gradient using the current estimate .
- Update by blending the diffusion step and the likelihood gradient (with a tunable step-size ).
Stability and performance are influenced by choice of schedule; ablation studies indicate both constant and decaying schedules can be effective, with the step size tuned to the noise level and task. The method is implemented entirely in the image domain, thereby evading the computational and memory cost of explicit eigenvalue decompositions.
Reproducibility is facilitated by the open-source codebase (https://github.com/DPS2022/diffusion-posterior-sampling). Empirical evaluation on FFHQ and ImageNet reveals state-of-the-art results across various tasks, with significant performance gains for complex and nonlinear inverse problems relative to prior plug-and-play diffusion methods.
6. Scalability, Limitations, and Extensions
DPS's avoidance of spectral transforms and SVDs yields scalability to high-dimensional regimes, subject primarily to the computational constraints imposed by the underlying diffusion model and neural network architecture. The main limitation is the reliance on surrogate likelihood gradients, which—though accurate in typical settings—can introduce bias if the approximation error is non-negligible, especially for highly nonlinear or adversarial measurement models.
Extensions to novel noise models or forward operators are possible with minimal algorithmic modification, requiring only an appropriate likelihood and differentiable measurement operator. In extremely noisy or strongly nonlinear regimes, additional regularization or adaptive step-size strategies may be beneficial. The method's generality enables its deployment as a universal solver for a broad array of inverse-imaging problems.
7. Broader Context and Impact
Diffusion-based posterior sampling methods, exemplified by the DPS approach (Chung et al., 2022), have catalyzed progress in both the theory and practice of inverse problem solving. By leveraging data-driven priors through deep diffusion models and integrating them flexibly with statistical noise models, these methods significantly advance the field beyond traditional Bayesian regularization, providing high reconstruction quality and robustness in diverse, realistic settings.
The ability to softly enforce data consistency without resorting to strict projections or domain-specific transformations has proven especially powerful in high-noise and nonlinear scenarios. This approach establishes a flexible and scalable framework for next-generation inverse problem solvers, broadly impacting imaging, signal processing, computational physics, and allied areas.