Source-linked AI summary
Diffusion Posterior Sampling for General Noisy Inverse Problems
Hyungjin Chung, Jeongsol Kim, Michael T. Mccann, Marc L. Klasky, Jong Chul Ye
TL;DR
Existing diffusion inverse-problem solvers mainly address simple linear, noiseless settings, while intermediate-time likelihoods are difficult to compute. DPS approximates that likelihood to support noisy linear and nonlinear inverse problems, achieving strong reconstructions across a broad range of tasks. Its principal scope includes highly noisy and nonlinear imaging problems, though a common Gaussian likelihood approximation is exact only at t = 0.
Problem
Existing methods lack means to handle measurement noise and are difficult or unsuccessful for nonlinear inverse problems, while the intermediate-time likelihood is generally intractable.
Method
DPS uses a tractable approximation of the likelihood based on the efficiently computed posterior mean and incorporates its gradient into diffusion sampling without measurement-subspace projections.
Results
DPS handles highly noisy and nonlinear inverse problems and outperforms existing state-of-the-art methods by large margins across the widest range of problems.
Takeaways & Limitations
The proposed strategy provides a versatile diffusion posterior sampler for noisy signal-dependent or signal-independent inverse problems, including nonlinear imaging tasks.
Takeaways & Limitations
For Gaussian noise, the common likelihood equation is correct only at t = 0 and is wrong at other diffusion noise levels.
Abstract
from arXiv · showhide
Diffusion models have been recently studied as powerful generative inverse problem solvers, owing to their high quality reconstructions and the ease of combining existing iterative solvers. However, most works focus on solving simple linear inverse problems in noiseless settings, which significantly under-represents the complexity of real-world problems. In this work, we extend diffusion solvers to efficiently handle general noisy (non)linear inverse problems via approximation of the posterior sampling. Interestingly, the resulting posterior sampling scheme is a blended version of diffusion sampling with the manifold constrained gradient without a strict measurement consistency projection step, yielding a more desirable generative path in noisy settings compared to the previous studies. Our method demonstrates that diffusion models can incorporate various measurement noise statistics such as Gaussian and Poisson, and also efficiently handle noisy nonlinear inverse problems such as Fourier phase retrieval and non-uniform deblurring. Code available at https://github.com/DPS2022/diffusion-posterior-sampling
1 INTRODUCTION
Diffusion models provide data-distribution priors for inverse problems, but existing approaches are limited by intractable likelihood terms, costly spectral computations, and narrow linear or noiseless settings. This work proposes a general approximation-based solver for noisy linear and nonlinear inverse problems.
- Motivation: Diffusion models use the score of the learned data distribution as a prior for recovering signals from measurements.The inverse problem combines the forward operator and detector noise with posterior sampling.
- Motivation: The likelihood term required for diffusion posterior sampling is analytically intractable because measurements depend explicitly on x0 rather than the time-dependent xt.This makes direct posterior sampling difficult within the diffusion process.
- Limitations of prior work: Spectral-domain methods require costly SVD computations, restricting applications such as deblurring to forward models with efficiently computable decompositions.Prior work considered separable Gaussian kernels because of this restriction.
- Limitations of prior work: Prior work focused on linear inverse problems and had not explored nonlinear inverse problems.The paper targets this gap alongside noisy measurements.
- Contribution: The proposed approximation handles Gaussian and Poisson noise, extends to nonlinear operators through automatic differentiation, and runs in the image domain without SVD.Experiments cover inpainting, super-resolution, several deblurring settings, and Fourier phase retrieval.
2 BACKGROUND
Diffusion inverse-problem solvers combine a learned score prior with a measurement likelihood, but the likelihood at intermediate diffusion times is generally unavailable. Existing projection-based approximations are mainly suited to noiseless settings, motivating a tractable likelihood surrogate for noisy and nonlinear problems.
- 2.1 SCORE-BASED DIFFUSION MODELS: The forward SDE gradually transforms data into a tractable distribution, while reverse diffusion uses a learned time-dependent score to generate data samples.The data distribution occurs at t = 0 and an isotropic Gaussian can be reached at t = T.
- 2.2 INVERSE PROBLEM SOLVING WITH DIFFUSION MODELS: Posterior sampling adds the measurement likelihood score to the diffusion prior score when recovering x from partial measurements y.Bayesian inverse problems use p(x) as a prior and sample from p(x|y).
- 2.2 INVERSE PROBLEM SOLVING WITH DIFFUSION MODELS: The likelihood score ∇xt log pt(y|xt) is difficult to compute in closed form because the measurement explicitly depends on x0, not xt.The unknown dependency persists despite the known forward measurement model.
- 2.2 INVERSE PROBLEM SOLVING WITH DIFFUSION MODELS: Projection-based methods alternate unconditional diffusion updates with measurement-subspace projections under an approximately noiseless assumption.Other approaches use coarse-to-fine or linear likelihood-gradient approximations.
- 2.2 INVERSE PROBLEM SOLVING WITH DIFFUSION MODELS: For Gaussian noise, the common likelihood approximation is correct only at t = 0 and is wrong at the other noise levels used during generation.The paper notes that signal-dependent noise such as Poisson requires separate treatment from the additive-noise notation.
- 2.2 INVERSE PROBLEM SOLVING WITH DIFFUSION MODELS: Existing methods empirically work on noiseless problems but lack means to handle measurement noise and are difficult or unsuccessful for nonlinear inverse problems.The paper therefore seeks a solver that addresses noisy measurements and scales effectively to nonlinear operators.
3 DIFFUSION POSTERIOR SAMPLING (DPS)
DPS approximates the intractable likelihood needed for posterior diffusion sampling by evaluating the measurement model at the diffusion posterior mean. It incorporates Gaussian and Poisson noise without measurement-subspace projections and supports nonlinear forward operators.
- 3.1 APPROXIMATION OF THE LIKELIHOOD: DPS factorizes the unavailable p(y|x_t) through p(x_0|x_t), then approximates it as p(y|x̂_0), where x̂_0 is the posterior mean.This replaces the outer expectation over p(y|x_0) with evaluation at the posterior mean.
- 3.1 APPROXIMATION OF THE LIKELIHOOD: The approximation error is characterized by the Jensen gap, which can decrease as measurement noise increases under finite-gradient and finite-moment conditions.The paper states that the Jensen gap can approach zero as σ →∞ when the relevant quantities are finite.
- 3.2 MODEL DEPENDENT LIKELIHOOD OF THE MEASUREMENT: For Gaussian and Poisson measurements, differentiating the corresponding likelihoods yields model-specific gradients that are inserted into ancestral diffusion sampling.The Gaussian and Poisson cases use separate likelihood forms and gradient expressions; Poisson-direct can be unstable because of its log term.
- 3.2 MODEL DEPENDENT LIKELIHOOD OF THE MEASUREMENT: DPS combines the likelihood-gradient update with standard diffusion sampling and applies to nonlinear operators A(·), unlike prior methods restricted to linear operators.The framework can use gradients obtained through automatic differentiation for nonlinear inverse problems.
- 3.2 MODEL DEPENDENT LIKELIHOOD OF THE MEASUREMENT: Unlike projection-based MCG updates, DPS omits measurement-subspace projections, avoiding manifold departure and accumulated error under noisy measurements.The paper connects Gaussian DPS to MCG while distinguishing it from MCG’s additional projection step.
4 EXPERIMENTS
Experiments evaluate DPS on noisy linear and nonlinear inverse problems using diverse datasets, baselines, and perceptual metrics. Across these settings, DPS produces strong reconstructions and handles problem classes and noise models that constrain competing methods.
- Experimental setup: Experiments use FFHQ and ImageNet 256×256 validation images, with 1k images per dataset and diffusion models trained or adopted as described.The ImageNet model is used without task-specific finetuning, while the FFHQ model is trained from scratch on 49k images for 1M steps.
- Experimental setup: The evaluation compares DPS with diffusion, plug-and-play, and total-variation baselines, using FID and LPIPS as primary metrics alongside PSNR and SSIM.Diffusion-based methods use the same score function for a fair comparison.
- Noisy linear inverse problems: DPS outperforms comparison methods by large margins on noisy Gaussian linear inverse problems, while projection-based methods can overfit corrupted measurements when noise is non-negligible.Figure 4 compares DPS with DDRM and PnP-ADMM for Gaussian-noise settings.
- Noisy linear inverse problems: DPS produces crisp, realistic reconstructions across Gaussian-noise tasks and is not restricted to forward operators whose singular value decompositions are efficiently computable.This enables applications such as complex motion deblurring that are inaccessible to SVD-dependent methods described in the passage.
- Nonlinear inverse problems: For nonlinear inverse problems, DPS yields highly accurate phase-retrieval reconstructions and highly realistic nonlinear-deblurring samples, outperforming other methods by large margins.Phase-retrieval metrics are computed from the best of four generated samples because the problem can be non-unique and initialization-dependent.
5 CONCLUSION
The conclusion presents DPS as a strategy for general noisy inverse problems in imaging. It emphasizes broad empirical performance and coverage across noisy and nonlinear settings.
- 5 CONCLUSION: DPS addresses general noisy inverse problems in imaging, including both signal-dependent and signal-independent noise.The conclusion characterizes the strategy as versatile for highly noisy and nonlinear inverse problems.
- 5 CONCLUSION: Extensive experiments show that DPS outperforms existing state-of-the-art methods by large margins and covers the widest range of problems.
A PROOFS
The appendix develops proof components for the posterior-score approximation, including exponential-family identities, Gaussian bounds, and a Jensen-gap error bound. These results support the stated approximation under Gaussian measurement models and diffusion sampling distributions.
- Posterior identities: Tweedie’s formula is introduced for exponential-family likelihoods to relate the posterior mean of the canonical parameter to derivatives of the marginal density.The likelihood is written using the canonical vector, sufficient statistic, cumulant-generating function, and base density.
- Diffusion sampling: For VP-SDE or DDPM sampling, the conditional distribution p(x0|xt) has a unique posterior mean, with the forward transition identified as Gaussian.The proof uses the Gaussian form of the forward sampling distribution and its canonical decomposition.
- Gaussian bounds: The appendix establishes derivative bounds for univariate and isotropic multivariate Gaussian densities using continuity, bounded derivatives, and the mean value theorem.The univariate bound identifies a constant L, while the multivariate lemma extends the Gaussian-density control to R^d.
- Proof completion: The proof sequence concludes after combining the posterior-distribution measure, Gaussian bounds, intermediate-value reasoning, and Jensen-gap control.
- Approximation bound: The Gaussian measurement-model theorem quantifies the approximation error through a Jensen gap and bounds that gap using moment and regularity assumptions.The proof defines f from the forward operator and a multivariate normal likelihood, then invokes Gaussian-density bounds and the Jensen-gap proposition.
B INVERSE PROBLEM SETUP
The paper evaluates diffusion posterior sampling across linear and nonlinear inverse problems, including super-resolution, inpainting, deblurring, phase retrieval, and Poisson noise. It emphasizes robustness to noisy measurements and avoids projection-based noise amplification.
- Forward models: The setup covers super-resolution, inpainting, linear deblurring, phase retrieval, and nonlinear deblurring through task-specific forward measurement models.Nonlinear deblurring models blur as temporal integration rather than convolution, with a learned forward model available from blurry-sharp image pairs.
- Forward models: Phase retrieval removes Fourier phase information, making reconstruction highly ill-posed and motivating oversampling to encourage uniqueness.The oversampling matrix has ratio k/n.
- Measurement noise: Poisson measurements are simulated by sampling photon counts with means equal to clean measurements, then clipping to [0, 255] and normalizing to [-1, 1].The clean measurement is A(x0).
- Projection robustness: Projection-based methods can fail dramatically under excessive measurement noise because applying A^T during projection amplifies noise in tasks such as super-resolution and deblurring.The paper reports that MCG can also fail on noiseless super-resolution, whereas the proposed method works across a broader range of inverse problems.
- Step size: DPS step sizes that are too small produce measurement-inconsistent reconstructions, while values that are too large create saturation artifacts and amplify noise.The reported best-practice range for ζ′ is [0.1, 1.0].
C.3 OTHER STEP SIZE SCHEDULES
The paper compares alternative step-size schedules and likelihood choices, then examines how reconstruction quality changes with neural function evaluations. The proposed schedule is preferred qualitatively, while DPS has a high-NFE advantage over DDRM.
- Step-size schedules: The step-size ablation evaluates 100 Gaussian-deblurred FFHQ images using average LPIPS against ground truth.The study compares linear decay, exponential decay, and steps proportional to 1/σ2.
- Step-size schedules: The proposed schedule most closely matches ground truth perceptually, whereas decaying schedules lose fine details and 1/σ2 scaling performs poorly.The reported qualitative pattern is supported by corresponding quantitative metrics in Table 5.
- Poisson likelihoods: Poisson-direct is unstable, Poisson-Gaussian fails to converge properly because of its MSE weighting, and Poisson-shot preserves high-frequency details better than Poisson-LS.Poisson-LS remains stable but tends to blur relevant reconstruction details.
- Sampling speed: For noisy SR×4, DPS outperforms other methods at high NFE (≥250), whereas DDRM performs better at low NFE (≤100).The difference is attributed to DDRM's DDIM sampling strategy, which is known for better low-NFE performance.
C.6 LIMITATIONS
The proposed method inherits diffusion models’ relatively slow sampling, while achieving a perception-distortion trade-off that favors perceptual quality over standard distortion metrics.
- Runtime: The proposed method is relatively slow because it inherits the computational characteristics of diffusion-model sampling.The paper notes that advanced samplers could mitigate the slow sampling speed.
- Perception-distortion trade-off: The method preserves high-frequency image details and outperforms DDRM on FID and LPIPS, but underperforms it on PSNR.The paper explains this pattern through the perception-distortion trade-off.
- Task scope: Phase-retrieval reconstruction quality is less robust than performance on linear inverse problems and nonlinear deblurring.The paper attributes this limitation to stochasticity and reports frequent failure cases.
D.1 IMPLEMENTATION DETAILS
Implementation uses problem-specific DPS step sizes, pretrained score functions for FFHQ and ImageNet, and a single RTX 2080Ti GPU. Runtime is reported per image at 1000 neural function evaluations.
- Step sizes: The implementation lists step sizes separately for Gaussian-noise, Poisson-noise, and nonlinear inverse-problem settings.Most schedules use ζi proportional to 1/∥y − A(ˆx0(xi))∥, with task- and dataset-specific coefficients.
- Score functions: Pretrained score functions come from Choi et al. for FFHQ and Dhariwal and Nichol for ImageNet.These score functions are used for the corresponding dataset experiments.
- Compute time: Experiments run on a single RTX 2080Ti GPU, taking about 95 seconds per FFHQ image and 600 seconds per ImageNet image at 1000 NFE.The longer ImageNet runtime is attributed to its much larger network size.
- Code availability: The authors provide the implementation as publicly available code in the diffusion-posterior-sampling repository.The repository URL is listed in the paper.
D.2 COMPARISON METHODS
The comparison methods use varied reconstruction strategies, including diffusion sampling, projection-based data consistency, optimization, and constraint-based phase retrieval. Their implementations differ in sampling budgets, projection rules, denoisers, and tuned parameters.
- DDRM uses DDIM sampling for 20 NFEs with ηB = 1.0 and η = 0.85; Gaussian deblurring uses separable 1D convolutions for efficient SVD.
- MCG applies additional data-consistency steps as Euclidean projections onto the measurement set C, using the same α values as DPS.
- Score-SDE iteratively applies denoising followed by Euclidean data-consistency projections onto the measurement set C.
- PnP-ADMM uses ρ = 0.2, maxiter = 12, and a pretrained DnCNN denoiser for proximal mappings.
- ADMM-TV minimizes a total-variation objective with finite-difference operators, isotropic TV regularization, and ADMM parameter ρ.
- ER, HIO, and OSS initialize real signals from a normal distribution, impose non-negative and finite-support constraints, and run 10,000 iterations.Each algorithm is repeated four times per data, with the best reconstruction reported by measurement-amplitude mean squared error.
E FURTHER EXPERIMENTAL RESULTS
The paper reports quantitative PSNR and SSIM evaluations and additional visual experiments across super-resolution, inpainting, and deblurring. These experiments cover FFHQ and ImageNet at 256 × 256 resolution, with Gaussian and Poisson measurement noise.
- Quantitative evaluations use the standard PSNR and SSIM metrics on FFHQ and ImageNet 256×256 validation datasets.Tables 6 and 7 report linear inverse-problem results on 1k-image validation sets.
- Super-resolution experiments compare ×8 and ×16 settings on FFHQ and ImageNet 256 × 256 datasets.
- Inpainting experiments use a 128×128 box and 92% random missingness on FFHQ and ImageNet 256 × 256 datasets.
- Deblurring experiments evaluate Gaussian and motion blur on FFHQ and ImageNet 256 × 256 datasets.
- Poisson-noise experiments evaluate super-resolution at λ = 0.05 and deblurring at λ = 1.0 on FFHQ 256 × 256 data.