PaperlayerView exact revision on arXiv

Source-linked AI summary

Improving Diffusion Inverse Problem Solving with Decoupled Noise Annealing

Bingliang Zhang, Wenda Chu, Julius Berner, Chenlin Meng, Anima Anandkumar, Yang Song

arXiv:2407.01521v3cs.LGcs.AIcs.CV

TL;DR

Existing diffusion posterior samplers struggle with nonlinear inverse problems because small denoising steps cannot correct large errors. DAPS decouples consecutive samples while annealing measurement-conditioned time-marginals toward the posterior, and it improves reconstruction quality and stability across challenging tasks.

  • Problem

    Small reverse-diffusion steps keep consecutive samples close, limiting correction of global errors in nonlinear inverse problems such as phase retrieval.

  • Method

    DAPS recursively samples measurement-conditioned time-marginals while decoupling consecutive samples and reducing noise toward zero.

  • Results

    DAPS improves reconstruction quality and stability across inverse problems, achieving 30.72 dB PSNR for FFHQ 256 phase retrieval, 1.98 dB above the previous state of the art.

  • Takeaways & Limitations

    Decoupled transitions let DAPS explore a larger solution space and produce better visual quality and stability, particularly for challenging nonlinear inverse problems.

  • Takeaways & Limitations

    The method assumes isotropic Gaussian measurement noise in its practical design choice.

Abstract

from arXiv · show

Diffusion models have recently achieved success in solving Bayesian inverse problems with learned data priors. Current methods build on top of the diffusion sampling process, where each denoising step makes small modifications to samples from the previous step. However, this process struggles to correct errors from earlier sampling steps, leading to worse performance in complicated nonlinear inverse problems, such as phase retrieval. To address this challenge, we propose a new method called Decoupled Annealing Posterior Sampling (DAPS) that relies on a novel noise annealing process. Specifically, we decouple consecutive steps in a diffusion sampling trajectory, allowing them to vary considerably from one another while ensuring their time-marginals anneal to the true posterior as we reduce noise levels. This approach enables the exploration of a larger solution space, improving the success rate for accurate reconstructions. We demonstrate that DAPS significantly improves sample quality and stability across multiple image restoration tasks, particularly in complicated nonlinear inverse problems.

1. Introduction

Bayesian inverse problems seek posterior-distributed reconstructions from noisy, non-unique measurements, with diffusion models providing learned data priors. DAPS addresses failures on complex nonlinear measurements by decoupling consecutive samples and annealing time-marginals toward the posterior, yielding stronger reconstruction results.

  • Inverse problems recover underlying signals from noisy, many-to-one measurements, so multiple solutions can be consistent with the observation.
  • Diffusion models estimate high-dimensional data priors by recursively converting noisy samples into less-noisy samples.
  • Small-step posterior sampling struggles with nonlinear measurements because consecutive samples remain close and cannot correct large global errors.
  • DAPS decouples consecutive samples, allowing larger changes while recursively sampling measurement-conditioned time-marginals that anneal toward the posterior.
  • 30.72 dB and 27.12 dB are achieved on FFHQ 256 phase retrieval and high dynamic range tasks, improving over prior state-of-the-art results by 1.98 dB and 4.39 dB.

2. Background

Diffusion models define noisy time-marginals and reverse their noising process to model data priors, while Bayesian inverse problems combine those priors with measurement likelihoods. Existing posterior-sampling approaches can fail on nonlinear problems because small reverse-SDE steps limit corrections to local errors.

  • 2.1. Diffusion Models: Diffusion models construct noisy data distributions across a noise schedule and generate samples by solving a reverse stochastic differential equation.
  • 2.1. Diffusion Models: Figure 2 depicts alternating posterior estimation, Langevin sampling, and forward noising until the noise level reaches zero.
  • 2.2. Bayesian Inverse Problems with Diffusion: Bayesian inverse problems combine a prior distribution with a noisy measurement likelihood to characterize the posterior over possible signals.
  • 2.2. Bayesian Inverse Problems with Diffusion: DPS and related methods estimate noisy likelihoods or interleave optimization and projection with diffusion sampling.
  • 2.2. Bayesian Inverse Problems with Diffusion: Small reverse-SDE step sizes keep consecutive samples close, so existing methods often correct minor errors but miss larger global errors in nonlinear inverse problems.

3. Method

DAPS reduces dependence between consecutive diffusion samples by sampling measurement-conditioned time-marginals through posterior estimation, Langevin dynamics, and controlled noising. The resulting framework extends to latent diffusion models and supports non-local transitions that better approximate posteriors in nonlinear problems.

  • 3.1. Decoupled Noise Annealing: DAPS samples a clean estimate conditioned on the current noisy sample and measurement, then adds Gaussian noise at the next level.
  • 3.1. Decoupled Noise Annealing: Decoupled noise annealing reduces noise progressively, starting from an approximately Gaussian high-noise marginal and ending with a sample from the posterior.
  • 3.1. Decoupled Noise Annealing: Langevin dynamics samples from the measurement-conditioned distribution using the prior conditional and measurement likelihood.
  • 3.2. Practical Design Choice of DAPS: The practical implementation commonly approximates p(x0 | xt) as a Gaussian centered at an unconditional probability-flow-ODE estimate, balancing accuracy and computational cost.
  • 3.3. LatentDAPS: DAPS can operate in latent space as LatentDAPS and achieves comparable or superior performance across several tasks.
  • 3.4. Discussion with Existing Methods: Unlike methods that solve a specific SDE or ODE, DAPS recursively samples time-marginals and permits non-local transitions that can correct accumulated errors.
  • 3.4. Discussion with Existing Methods: On a 2D nonlinear inverse problem, DAPS varies more between trajectory points and approximates the true posterior, whereas DPS converges to wrong solutions.

4. Experiments

Experiments evaluate DAPS across linear and nonlinear image inverse problems, showing strong reconstruction quality, stability, diversity, and practical efficiency. The method also extends to text-conditioned latent diffusion and compressed sensing MRI.

  • Main results: DAPS achieves comparable or better performance on five linear inverse problems and greater stability on nonlinear tasks than existing methods.The evaluation covers FFHQ and ImageNet using PSNR, SSIM, LPIPS, and FID, with nonlinear results reported using mean and standard deviation.
  • Main results: DAPS substantially improves phase-retrieval success rates, yielding higher PSNR and lower LPIPS than methods with high failure rates.The method also recovers more fine-grained measurement details than existing baselines.
  • Sampling-process analysis: PSNR and LPIPS improve stably as noise is annealed, while measurement error remains consistently low for samples from the measurement-conditioned time-marginal.The analysis compares the estimated Gaussian-approximation mean with samples obtained through Langevin dynamics.
  • Sample diversity: DAPS produces diverse samples for sparse measurements whose posterior distributions may contain multiple modes, while preserving measurement information.Demonstrated settings include 16× super-resolution and 192×192 box inpainting, with variation in expression, clothing, and hairstyle.
  • Further applications: LatentDAPS extends the approach to large-scale text-conditioned latent diffusion, while DAPS also outperforms diffusion-based baselines on compressed sensing multi-coil MRI.The text-conditioned demonstration uses Stable Diffusion v2.1 at 768×768 resolution; MRI reduces scan time through subsampling.
  • Ablation studies: Gaussian approximation makes DAPS 6–7 times faster than diffusion-score estimation with only a slight performance drop.The comparison is reported as an ablation of approaches for approximating the conditional distribution.

5. Conclusion

The paper concludes that DAPS addresses challenging inverse problems by decoupling consecutive diffusion samples. This enables broader exploration and improves visual quality and stability across difficult tasks.

  • DAPS targets inverse problems with complex nonlinear measurement processes, particularly phase retrieval.
  • Decoupling consecutive sample points allows larger trajectory changes and enables exploration of a larger solution space.
  • DAPS generates samples with better visual quality and stability than existing methods across challenging inverse problems.

Supplementary Material

The supplementary material extends DAPS with alternative samplers, latent and text-conditioned diffusion models, and implementation algorithms. It also reports sampling behavior and efficiency for these variants.

  • Pixel-space sampling: Langevin dynamics is the default DAPS sampler, while HMC and Metropolis-Hastings provide alternative posterior-sampling strategies.The algorithm can use Langevin dynamics, HMC, or Metropolis-Hastings; HMC is intended to reduce random-walk behavior, whereas Metropolis-Hastings supports settings without likelihood gradients.
  • Hamiltonian Monte Carlo: HMC explores p(x0 | xt, y) efficiently and can accelerate LatentDAPS with large-scale text-conditioned latent diffusion models.The paper reports that HMC requires fewer steps than Langevin dynamics for similar performance and speeds up LatentDAPS with large-scale text-conditioned LDMs.
  • Metropolis-Hastings: Metropolis-Hastings is more flexible for non-differentiable forward models but is generally less efficient and performs worse than Langevin dynamics and HMC.Its flexibility applies when the gradient of the data likelihood is unavailable.
  • Latent diffusion models: LatentDAPS transfers DAPS sampling to VAE latent space and supports pixel-space or latent-space Langevin dynamics.The paper notes a trade-off: pixel-space dynamics are faster, while latent-space dynamics can reduce approximation error and improve sample quality.
  • Text-conditioned models: Text-conditioned LatentDAPS uses Stable Diffusion v1.5 with classifier-free guidance and enables text-based control over posterior samples.Experiments evaluate eight FFHQ 256 tasks with guidance scale 7.5; text prompts also control modes for 8× super-resolution samples.

C. DAPS with Discrete Diffusion Models

DAPS extends posterior sampling to discrete diffusion models by replacing the continuous Gaussian approximation with an exponential model over Hamming distance. Experiments on discretized MNIST show strong classification accuracy on nonlinear categorical inverse problems.

  • Method: DAPS adapts to discrete diffusion models by approximating p(x0 | xt) with an exponential distribution over Hamming distance.The scale parameter rt is selected heuristically, and the resulting procedure is summarized in a discrete-DAPS algorithm.
  • Experiments: The discrete experiments use binary-string MNIST for 50% inpainting and an XOR inverse problem.The XOR task computes XOR over randomly paired MNIST binary strings and serves as a highly nonlinear measurement problem.
  • Results: DAPS achieves very high classification accuracy on both discrete inverse problems and outperforms both baselines by a large margin.The classification accuracy is evaluated with a ConvNet, while the baselines use 20 particles for sampling.
  • Marginal propagation: The theoretical construction assumes samples are drawn from measurement-conditioned time-marginals and propagates them to later marginals.The supplied propositions and proof passages express this marginal-preservation property for pixel and latent variables.
  • Efficiency: DAPS can achieve relatively better performance than baselines with small numbers of neural function evaluations.The paper presents NFE comparisons for phase retrieval on FFHQ 256 and identifies NFE as a key determinant of diffusion sampling cost.

E.2. Limitations and Future Extension

The paper identifies limitations in its latent implementation, annealing schedule, and ODE evaluation strategy, while noting possible extensions. It also cautions that biased training data may produce biased samples.

  • Latent implementation: LatentDAPS uses a naive latent-diffusion implementation, and pixel-space Langevin samples may leave the clean-image manifold.This can cause sub-optimal performance for diffusion autoencoders trained only on clean data.
  • Annealing schedules: The experiments use only decreasing annealing schedules, although non-monotonic schedules may improve solution-space exploration.The broader DAPS framework is stated to support arbitrary scheduler functions.
  • ODE computation: The method uses a fixed number of ODE-solver evaluations, while automatically reducing evaluations at later sampling steps is proposed as a future extension.The paper specifically suggests using fewer evaluations for smaller t.
  • Broader impacts: DAPS may generate biased samples when its diffusion model is trained on biased data.The paper recommends caution when applying the method in bias-sensitive scenarios.

F.2. DAPS Implementation Details

The implementation computes denoising estimates with probability-flow ODEs and anneals measurement-conditioned samples through a scheduled noise process. Experiments use task-dependent solver and annealing configurations alongside established baseline implementations.

  • ODE solver: DAPS computes x̂0(xt) by solving the probability-flow ODE initialized at xt and time t.The score model approximates the time-dependent score, and an Euler solver is used numerically.
  • Annealing scheduler: The annealing process samples xt from measurement-conditioned time-marginals while reducing the annealing noise to zero.The process starts at time T and uses an annealing scheduler across NA steps.
  • Hyperparameters: The implementation trades computation against quality through ODE-solver steps and annealing steps, using more steps for nonlinear tasks.The stated configurations are Node=5 and NA=200 for linear tasks, versus Node=10 and NA=400 for nonlinear tasks.
  • Hyperparameters: The Langevin step size decays linearly, while σmax and σmin are fixed separately for DAPS and LatentDAPS.The reported values are σmax=100 for DAPS, σmax=10 for LatentDAPS, and σmin=0.1.
  • Baselines: Baseline comparisons use the original or official implementations with their reported default settings, except for specified task-specific adjustments.Examples include DDIM step counts for DDRM and DDNM, official code for DCDP and ReSample, and a random-noise initialization for RED-diff.

G. Experiments on Synthetic Data Distributions

On a synthetic two-dimensional Gaussian-mixture inverse problem, DAPS more accurately recovers the single-mode posterior than DPS by sampling time-marginal distributions and correcting early-stage errors.

  • The synthetic prior is a two-dimensional Gaussian mixture, while the nonlinear measurement likelihood has two modes but the posterior is single-mode.The prior concentrates near one likelihood mode, making the posterior effectively single-mode.
  • DAPS recovers the target posterior more accurately than both SDE and ODE versions of DPS.DPS converges to two different modes because early likelihood-estimation errors propagate along its sampling trajectories.
  • DAPS samples a time-marginal distribution at each step, enabling recovery of the posterior despite errors that DPS cannot correct.This mechanism permits larger corrections than recursively conditioning each sample on its predecessor.
  • DAPS’s Wasserstein distance to the ground-truth posterior decreases quickly and remains small throughout sampling.This behavior matches the stated theory that x_t follows p(x_t | y) at every noise level.

H. Experimental Results on Compressed Sensing Multi-coil MRI

The CS-MRI experiment evaluates DAPS on reconstructing undersampled multi-coil MRI images using a diffusion prior, finding that it outperforms diffusion-based baselines.

  • CS-MRI reconstructs high-resolution source images from undersampled k-space measurements to reduce MRI acquisition time.Its forward model combines subsampling, Fourier transformation, and multi-coil sensitivity maps.
  • The experiment trains a diffusion model on preprocessed fastMRI knee images and compares DAPS with general restoration and MRI-specific baselines.Compared methods include DPS, DiffPIR, ScoreMRI, and CSGM.
  • DAPS outperforms existing diffusion-based baselines on compressed sensing multi-coil MRI.Additional visual reconstructions validate its capability on real-world medical imaging inverse problems.

I.1. More Ablation Study

The ablations examine sampler choice, computational budget, noise scheduling, measurement noise, and qualitative behavior across inverse problems, identifying accuracy–perceptual-quality trade-offs and strong nonlinear-task sensitivity to scheduling.

  • MCMC samplers: HMC provides the best balance between time efficiency and generation quality among the evaluated DAPS MCMC samplers.The comparison covers super-resolution and high-dynamic-range tasks on FFHQ 256.
  • ODE steps: Increasing ODE-sampler steps improves perceptual quality and fine-grained detail but slightly lowers PSNR.The resulting images become less blurry and gain high-frequency details, improving LPIPS while decreasing PSNR.
  • Noise scheduling: DAPS performance on nonlinear tasks increases dramatically with more annealing scheduling steps and reaches near-optimal quality beyond 200 steps.Linear-task performance improves only slightly as scheduling becomes denser.
  • Measurement noise: DAPS is robust to measurement noise below 0.05 and degrades almost linearly as the noise level increases.The result characterizes sensitivity to measurement noise rather than changing the reconstruction procedure.
  • Phase retrieval: DAPS shows better sample quality and stability than baselines on phase retrieval, an ill-posed task with multiple disjoint modes sharing identical measurements.Oversampling experiments further evaluate performance across phase-retrieval measurement ambiguities.
  • Qualitative results: Phase-retrieval trajectories converge from unconditional or noisy conditioned samples toward posterior samples as noise anneals down.Additional figures show independent runs, diverse multimodal reconstructions, and samples across restoration and MRI tasks.
Loading 2407.01521v3…