PaperlayerView exact revision on arXiv

Source-linked AI summary

A Variational Perspective on Solving Inverse Problems with Diffusion Models

Morteza Mardani, Jiaming Song, Jan Kautz, Arash Vahdat

arXiv:2305.04391v2cs.LGcs.CVmath.NAstat.ML

TL;DR

Diffusion inverse problems require posterior sampling, but the conditional score is difficult to estimate because the denoising posterior can be complex and multimodal. The paper uses variational KL minimization to derive RED-diff, an SNR-weighted diffusion regularizer optimized with lightweight stochastic iterations. Experiments report stronger image-restoration quality than state-of-the-art samplers, alongside a mode-seeking limitation that reduces diversity.

  • Problem

    Conditional posterior sampling is challenging because diffusion models make the required likelihood score intractable when the denoising posterior is complex and multimodal.

  • Method

    RED-diff uses variational KL minimization with diffusion-prior score-matching regularization and weights denoisers by λ_t = λ/SNR_t.

  • Results

    RED-diff consistently outperforms alternative samplers across reported inpainting metrics and performs better across several nonlinear restoration tasks.

  • Takeaways & Limitations

    Viewing sampling as stochastic optimization enables off-the-shelf optimizers and lightweight, tunable inference iterations.

  • Takeaways & Limitations

    The variational sampler is mode-seeking, promotes MAP solutions, and therefore lacks diversity.

Abstract

from arXiv · show

Diffusion models have emerged as a key pillar of foundation models in visual domains. One of their critical applications is to universally solve different downstream inverse tasks via a single diffusion prior without re-training for each task. Most inverse tasks can be formulated as inferring a posterior distribution over data (e.g., a full image) given a measurement (e.g., a masked image). This is however challenging in diffusion models since the nonlinear and iterative nature of the diffusion process renders the posterior intractable. To cope with this challenge, we propose a variational approach that by design seeks to approximate the true posterior distribution. We show that our approach naturally leads to regularization by denoising diffusion process (RED-Diff) where denoisers at different timesteps concurrently impose different structural constraints over the image. To gauge the contribution of denoisers from different timesteps, we propose a weighting mechanism based on signal-to-noise-ratio (SNR). Our approach provides a new variational perspective for solving inverse problems with diffusion models, allowing us to formulate sampling as stochastic optimization, where one can simply apply off-the-shelf solvers with lightweight iterates. Our experiments for image restoration tasks such as inpainting and superresolution demonstrate the strengths of our method compared with state-of-the-art sampling-based diffusion models.

2 RELATED WORKS

Diffusion-based inverse-problem methods differ in how they incorporate measurements, extending from linear Langevin approaches to guidance methods for nonlinear tasks. Recent methods improve measurement guidance but retain scope or posterior-score approximation limitations.

  • Diffusion models for inverse problems: Diffusion models have been applied plug-and-play to inverse problems in natural and medical imaging and audio processing.The paper focuses primarily on image inverse problems and how methods handle measurements.
  • Diffusion models for inverse problems: Early approaches use Langevin dynamics for linear inverse problems and integrate observations into the sampling process.Kadkhodaie & Simoncelli and Jalal et al. are cited as examples.
  • Guidance methods: Reconstruction guidance and DPS extend measurement guidance to nonlinear inverse problems by differentiating through the diffusion model.DPS is described as an extension of reconstruction guidance to nonlinear tasks.
  • Guidance methods: ΠGDM improves guidance approximation through pseudoinverse measurement-model guidance, but its scope is limited to linear and certain semi-linear tasks.JPEG is given as an example of a semi-linear task.
  • Diffusion models for 3D: Diffusion-prior methods have also been connected to 3D generation through distillation-loss optimization, including DreamFusion, ProlificDreamer, and SparseFusion.These works treat generation from diffusion priors as an optimization problem.

3 BACKGROUND

The paper frames inverse problems as posterior sampling with a pretrained diffusion prior, but the required conditional score is difficult to estimate because the denoising posterior can be complex and multimodal. It introduces variational KL minimization, yielding a measurement-matching loss plus diffusion score-matching regularization and an optimization-based sampler.

  • 3.1 DENOISING DIFFUSION MODELS: Diffusion models add noise through a forward process and learn iterative denoising through a reverse process to generate images.The forward process is designed to end in a standard Gaussian distribution.
  • 3.2 SCORE APPROXIMATION FOR INVERSE PROBLEMS: Inverse problems seek x0 from nonlinear and noisy measurements y, commonly requiring a strong prior because tasks such as inpainting are severely ill-posed.The pretrained diffusion distribution is used as a plug-and-play prior, with the goal of sampling p(x0|y).
  • 3.2 SCORE APPROXIMATION FOR INVERSE PROBLEMS: The conditional score combines likelihood and prior scores, but the likelihood score is intractable because p(y|xt) integrates over a highly complex and multimodal p(x0|xt).Prior methods therefore use unimodal Gaussian or MMSE approximations.
  • 4 VARIATIONAL DIFFUSION SAMPLING: The variational approach minimizes KL divergence using q = N(µ, σ^2I), seeking a dominant data-distribution mode that matches the observations.The variational objective can be reduced to minimizing its variational bound.
  • 4 VARIATIONAL DIFFUSION SAMPLING: The resulting objective combines measurement matching with score-matching regularization from the diffusion prior, evaluated along a diffused trajectory.The variational Gaussian makes q(xt|y) Gaussian and its score analytically computable.
  • 4.1 SAMPLING AS STOCHASTIC OPTIMIZATION: Sampling is formulated as stochastic optimization over diffusion timesteps, with RED-style regularization linking denoisers to structural constraints across the trajectory.The method’s regularizer is similar to RED while retaining a generative diffusion prior and using the entire trajectory.
  • 4.3 WEIGHTING MECHANISM: The weighting mechanism uses λ_t = λ/SNR_t to convert the noise-prediction objective into a clean-data formulation and balance prior fitting against observation fitting.The hyperparameter λ controls the bias–variance trade-off between observations and the prior.

5 EXPERIMENTS

Experiments compare RED-diff with DPS, ΠGDM, and DDRM on ImageNet restoration tasks and ablate timestep weighting and sampling. RED-diff is reported to improve restoration quality across linear and nonlinear settings while using lightweight optimization iterations.

  • 5 EXPERIMENTS: RED-diff is evaluated against DPS, ΠGDM, and DDRM on linear and nonlinear image-restoration tasks using a pretrained unconditional guided diffusion model.The proof-of-concept experiments use a 1k ImageNet validation subset and report KID, LPIPS, SSIM, PSNR, and classifier accuracy.
  • 5.1 IMAGE INPAINTING: In ImageNet inpainting, RED-diff identifies context and adds fine details, whereas ΠGDM fails on a hard example and DPS and DDRM produce blurry content.The comparison uses representative hard and easy examples with tuned samplers.
  • 5.1 IMAGE INPAINTING: RED-diff consistently outperforms alternative samplers across reported inpainting metrics, including KID and PSNR, with a significant margin.The authors interpret this as more faithful restoration and better perceptual quality.
  • 5.1 IMAGE INPAINTING: RED-diff iterations require only forward passes through the score network, while DPS and ΠGDM differentiate through denoisers and DDRM performs costly SVD calculations.Time per step and GPU memory usage are compared empirically in Table 1.
  • 5.2 NONLINEAR INVERSE PROBLEMS: Across nonlinear HDR, phase retrieval, and deblurring tasks, RED-diff is reported to outperform or substantially improve on DPS, which struggles or is sensitive to tuning and initialization.ΠGDM and DDRM are excluded from these comparisons because they only address linear inverse problems.
  • 5.3.1 DENOISER WEIGHTING MECHANISM: For denoiser weighting, square-root inverse-SNR decay performs best for KID, while linear inverse-SNR weighting performs best for PSNR.The ablation compares monotonic functions of inverse SNR across denoiser timesteps.
  • 5.3.2 TIMESTEP SAMPLING: Descending timestep sampling performs significantly better than other tested strategies, progressively adding semantic structures before fine details.The comparison includes random, ascending, descending, minibatch-random, and minibatch-descending strategies.

6 CONCLUSIONS AND LIMITATIONS

RED-diff is a variational sampler for universal inverse problems that regularizes through denoising diffusion and uses SNR-based weighting. Experiments report strong performance against state-of-the-art alternatives, while the sampler remains mode-seeking and lacks diversity.

  • RED-diff applies a variational approach to universal inverse problems using diffusion priors and denoising-based structural constraints.Denoisers across diffusion steps impose constraints ranging from high-level semantics to low-level details.
  • SNR-based weighting tunes the contribution of denoisers across diffusion steps, while stochastic optimization enables efficient and adjustable sampling with standard optimizers.
  • Experiments on several image restoration tasks report strong RED-diff performance compared with state-of-the-art inverse-problem alternatives.
  • RED-diff lacks diversity because its mode-seeking variational formulation promotes MAP solutions.The authors suggest optimizer changes, more expressive variational distributions, or dispersion terms as possible directions.

A.1 PROOF OF PROPOSITION 1

The proof rewrites the variational KL objective using Bayes rule and separates it into a reconstruction term and a diffusion-prior term. Under the stated assumptions, the latter becomes a weighted score-matching objective along the denoising trajectory.

  • Bayes rule rewrites the KL objective into terms whose minimization separates measurement reconstruction from the diffusion-prior contribution.
  • The likelihood term becomes a reconstruction loss under the assumed i.i.d. Gaussian measurement noise model.
  • Under mild growth assumptions on log q(x_t|y) and p(x_t), weighted score matching over the continuous denoising trajectory equals the KL divergence.
  • Reweighting the trajectory objective can produce divergence measures other than the one corresponding directly to maximum likelihood.

A.2 PROOF OF PROPOSITION 2

The proof derives a generic timestep-weighted representation of the diffusion regularizer and obtains tractable gradients through reparameterization. Boundary assumptions make the endpoint contributions vanish.

  • The diffusion regularization term is a score-matching loss whose weighting can differ from β(t)/2, corresponding to alternative objectives beyond maximum likelihood.
  • Weighting the KL time-derivative by ω(t), with ω(0)=0, yields a generic reweighted objective.
  • The boundary contribution vanishes at t=0 because ω(t)=0 and at t=T because the terminal diffusion state is pure Gaussian noise.
  • For a deterministic variational distribution, reparameterization expresses the forward state as x_t=α_tμ+σ_tε and enables a tractable gradient with respect to μ.

A.3 ADDING DISPERSION TO VARIATIONAL APPROXIMATION

The variational approximation extends from a point mass to a Gaussian with optimized mean and dispersion, yielding closed-form gradients. The authors caution that Gaussian dispersion is poorly suited to generating diverse natural images.

  • Extending the variational distribution from a Dirac mass to a Gaussian preserves the mean gradient and provides a simple, tractable dispersion gradient.
  • For Gaussian dispersion, reparameterization represents the diffused signal using the variational mean, dispersion, timestep noise scale, and standard Gaussian noise.
  • The dispersion gradient has a closed form, but Gaussian perturbations do not generally produce another legitimate natural image.
  • The authors identify image-manifold dispersion models as a more sophisticated alternative for adding stochasticity and diversity.

B ADDITIONAL EXPERIMENTS

The experiments use a pretrained, class-unconditional guided diffusion model trained on 256 × 256 ImageNet images.

  • The score function comes from a pretrained guided diffusion model without class conditioning.The model was trained on the 256 × 256 ImageNet dataset.

B.2 IMAGE SUPERRESOLUTION

RED-diff is evaluated on image restoration tasks including 4x ImageNet superresolution, inpainting, noisy inpainting, and compressed sensing MRI. Results show a fidelity–perceptual-quality trade-off, diverse samples, and competitive reconstruction quality across settings.

  • B.2 IMAGE SUPERRESOLUTION: RED-diff balances image fidelity and perceptual quality in 4x ImageNet superresolution.Tuning the Adam learning rate shifts the balance: smaller rates favor fidelity, while larger rates favor perceptual quality.
  • B.2 IMAGE SUPERRESOLUTION: RED-diff significantly outperforms alternative samplers in PSNR and SSIM but trails them on perceptual quality.Adjusting λ changes the fidelity–perceptual-quality trade-off through the bias-variance balance.
  • B.2 IMAGE SUPERRESOLUTION: RED-diff produces sufficiently diverse ImageNet inpainting samples when the noise seed changes.Larger Adam learning rates or fewer optimization steps can further increase diversity.
  • B.2 IMAGE SUPERRESOLUTION: During restoration, earlier denoisers generate high-level structures before later denoisers refine lower-level features.The evolution is examined in both image and frequency domains through magnitude and phase visualizations.
  • B.2 IMAGE SUPERRESOLUTION: RED-diff is also evaluated on noisy inpainting and compressed sensing MRI using pretrained diffusion priors.MRI experiments compare against CSGM-Langevin on fastMRI brain and Mridata knee datasets under undersampling.

D.3 OPTIMIZATION STRATEGY

Optimization ablations show that RED-diff can use different optimizers, while learning-rate and regularization-weight tuning control the fidelity–perceptual-quality trade-off. Additional iterations provide negligible sample-quality improvement.

  • D.3 OPTIMIZATION STRATEGY: SGD with momentum can perform as well as Adam, although SGD is more sensitive to step size.This indicates that RED-diff is not sensitive to the optimizer choice when hyperparameters are tuned.
  • D.3 OPTIMIZATION STRATEGY: Smaller Adam learning rates improve reconstruction quality, whereas larger rates improve perceptual quality measured with KID.The reported examples compare learning rates such as 0.05 and 0.5.
  • D.3 OPTIMIZATION STRATEGY: Adding more optimization epochs produces negligible improvement in ImageNet inpainting sample quality.The comparison uses 1,000 Adam steps per epoch.
  • D.3 OPTIMIZATION STRATEGY: An optimal λ supports tuning the bias-variance trade-off for sampling quality.KID is evaluated across a range of λ values.

D.4 TIMESTEP SAMPLING STRATEGY

Timestep sampling and denoiser weighting are important to RED-diff optimization. Descending timestep sampling performs best, while denoiser residual behavior motivates nonuniform weighting and limited optimization steps.

  • D.4 TIMESTEP SAMPLING STRATEGY: Descending timestep sampling performs significantly better than random, ascending, and mini-batch alternatives.It starts from the denoiser at t = T and proceeds toward smaller timesteps.
  • D.4 TIMESTEP SAMPLING STRATEGY: Batching 25 denoisers per iteration is computationally parallelizable but smears the final result.The batch strategy sorts 1,000 timesteps in descending order and uses 40 iterations.
  • D.4 TIMESTEP SAMPLING STRATEGY: With 100 steps, ImageNet inpainting achieves its best perceptual quality at KID=1.93.More steps improve reconstruction PSNR by refining fidelity, indicating that relatively few steps suffice for the sampler.
  • D.4 TIMESTEP SAMPLING STRATEGY: Unlike classical RED, RED-diff adds noise to denoiser inputs and stochastically navigates toward the diffusion prior.The method treats sampling as stochastic optimization rather than a deterministic fixed-point problem.
Loading 2305.04391v2…