PaperlayerView exact revision on arXiv

Source-linked AI summary

Solving Inverse Problems with Latent Diffusion Models via Hard Data Consistency

Bowen Song, Soo Min Kwon, Zecheng Zhang, Xinyu Hu, Qing Qu, Liyue Shen

arXiv:2307.08123v3cs.CV

TL;DR

Pixel-space diffusion priors require substantial computational resources and training data, while latent diffusion introduces challenges from encoder and decoder nonlinearities. ReSample addresses these challenges with hard data consistency and stochastic remapping, achieving state-of-the-art performance across diverse inverse problems and datasets.

  • Problem

    Pixel-space diffusion priors require substantial computational resources and large training datasets, while latent diffusion models introduce inverse-problem challenges because their encoder and decoder are nonlinear.

  • Method

    ReSample solves a hard-constrained optimization for measurement consistency, then stochastically remaps the consistent sample onto the noisy latent data manifold during reverse sampling.

  • Results

    ReSample achieves state-of-the-art performance across linear and nonlinear inverse problems on natural and medical images, outperforming baselines across three metrics on CelebA-HQ.

  • Takeaways & Limitations

    The method enables pre-trained latent diffusion models to solve general inverse problems while also supporting fine-tuning for medical imaging with reduced memory usage.

  • Takeaways & Limitations

    Latent-DPS often produces noisy or blurry reconstructions that are inconsistent with measurements, potentially because decoder-involved optimization is highly nonconvex.

Abstract

from arXiv · show

Diffusion models have recently emerged as powerful generative priors for solving inverse problems. However, training diffusion models in the pixel space are both data-intensive and computationally demanding, which restricts their applicability as priors for high-dimensional real-world data such as medical images. Latent diffusion models, which operate in a much lower-dimensional space, offer a solution to these challenges. However, incorporating latent diffusion models to solve inverse problems remains a challenging problem due to the nonlinearity of the encoder and decoder. To address these issues, we propose \textit{ReSample}, an algorithm that can solve general inverse problems with pre-trained latent diffusion models. Our algorithm incorporates data consistency by solving an optimization problem during the reverse sampling process, a concept that we term as hard data consistency. Upon solving this optimization problem, we propose a novel resampling scheme to map the measurement-consistent sample back onto the noisy data manifold and theoretically demonstrate its benefits. Lastly, we apply our algorithm to solve a wide range of linear and nonlinear inverse problems in both natural and medical images, demonstrating that our approach outperforms existing state-of-the-art approaches, including those based on pixel-space diffusion models.

1 Introduction

Inverse problems require priors because measurements often underspecify the unknown signal. ReSample uses latent diffusion models with hard data consistency to address pixel-space limitations and supports broad inverse-problem settings.

  • Inverse problems reconstruct an unknown signal from measurements produced by linear or nonlinear forward operators, often with fewer measurements than signal dimensions.
  • Pixel-space diffusion priors demand substantial computational resources and large training datasets, limiting their applicability to real-world imaging such as CT and MRI.
  • ReSample uses latent diffusion models as priors and enforces measurement consistency through a hard-constrained optimization problem.
  • A stochastic resampling scheme maps measurement-consistent samples back onto the noisy data manifold so reverse sampling can continue.
  • Experiments span linear and nonlinear inverse problems across natural and medical images, with state-of-the-art performance reported across multiple tasks and datasets.

2 Background

Diffusion-based inverse solvers use reverse sampling conditioned on measurements, while latent diffusion reduces dimensionality but introduces decoder nonlinearity. Existing latent-gradient updates can produce noisy, blurry, and measurement-inconsistent reconstructions.

  • Diffusion Models: DDPMs learn reverse-time dynamics from a score function, whereas DDIMs use a non-Markovian process to accelerate sampling.
  • Diffusion Models for Solving Inverse Problems: Posterior sampling conditions reverse diffusion on measurements, but the likelihood gradient generally lacks an analytical form, motivating projection and approximation-based methods.
  • Diffusion Models for Solving Inverse Problems: Pixel-space inverse-problem diffusion methods have limited applicability to real-world problems because they are built directly in pixel space.
  • Solving Inverse Problems with Latent Diffusion Models: Latent diffusion encodes images into a lower-dimensional latent space using an encoder and decodes them back through a nonlinear decoder.
  • Solving Inverse Problems with Latent Diffusion Models: Latent-DPS treats the decoder-mediated problem as nonlinear, but often yields noisy or blurry reconstructions that are inconsistent with measurements.

3 ReSample: Inverse Problems using Latent Diffusion Models

ReSample enforces measurement consistency during latent reverse sampling, then remaps consistent estimates onto the noisy latent manifold. Its stochastic resampling has theoretical variance and unbiasedness benefits, supporting inverse-problem reconstruction across multiple tasks.

  • Hard Data Consistency: Hard data consistency solves a measurement-constrained latent optimization problem, initialized from the previous reverse-sampling estimate to avoid unfavorable local minima.This strictly enforces measurements rather than applying a soft gradient update.
  • Sampling Procedure: ReSample alternates unconditional diffusion steps with selected resampling steps rather than solving the consistency optimization at every iteration.The skipped-step mechanism reduces cost while preserving semantic or structural information across nearby sampling steps.
  • Remapping Back to zt: After consistency optimization, ReSample maps the measurement-consistent estimate back to the noisy latent manifold using an auxiliary conditional distribution.This remapping enables continuation of the reverse sampling process.
  • Stochastic Resampling: Stochastic resampling conditions on the unconditional noisy sample to reduce the variance associated with stochastic encoding while balancing prior and data consistency.The variance parameter σ2_t controls this tradeoff, and σ2_t →∞ recovers stochastic encoding.
  • Theoretical Results: VAR(ẑ_t) < VAR(ṽ_t) when VAR(z′_t) > 0, while stochastic resampling remains unbiased under measurement consistency.Theorem 1 states E[ẑ_t|y] = E[z′_t|y] when ẑ0(y) satisfies y = A(D(ẑ0(y))).
  • Theoretical Results: As t →0, the variance between the ground-truth latent and estimated latent decreases to 0 under the stated Hessian boundedness assumption.The paper connects this decreasing variance to less noisy reconstructions than stochastic encoding.

4 Experiments

Experiments evaluate ReSample on linear and nonlinear inverse problems across natural and medical images, using comparisons with pixel- and latent-space baselines. ReSample generally achieves strong reconstruction quality, while resampling and repeated hard data consistency improve results.

  • Experimental setup: ReSample is evaluated against pixel-space and latent-space baselines on linear and nonlinear inverse problems using LPIPS, PSNR, and SSIM.Tasks include Gaussian deblurring, random-mask inpainting, super resolution, nonlinear deblurring, and CT reconstruction.
  • Natural images: ReSample significantly outperforms all baselines across all three metrics on CelebA-HQ and performs better than or comparably to baselines on FFHQ.The experiments use pretrained latent diffusion models and 100 validation images from the natural-image datasets.
  • Medical images: ReSample outperforms baselines on CT reconstruction in PSNR and SSIM.The medical-image experiments fine-tune latent diffusion models on 2,000 CT images and test on images from 10 held-out patients.
  • Resampling ablation: Stochastic resampling reconstructs smoother images with higher PSNRs than stochastic encoding on Gaussian deblurring and super-resolution.The ablation is conducted on LSUN-Bedroom and CelebA-HQ with measurement noise σy = 0.05.
  • Hard data consistency: More ReSample time steps produce more accurate CT reconstructions with higher PSNR.The ablation supports reinforcing hard data consistency during reverse sampling.
  • Efficiency: ReSample has lower overall memory usage, supporting its potential feasibility for medical imaging.The passage identifies memory complexity as important in medical-imaging applications.

5 Conclusion

The conclusion presents ReSample as an effective latent-diffusion algorithm for general inverse problems. It reports high-quality reconstructions and reduced memory usage after fine-tuning for medical imaging, while identifying hard data consistency as a computational bottleneck.

  • Conclusion: ReSample leverages latent diffusion models to solve general linear and nonlinear inverse problems.The method is presented as the paper’s central algorithmic contribution.
  • Conclusion: Experiments show high-quality reconstructions compared with many baselines and reduced memory usage for fine-tuned medical-imaging models.The conclusion summarizes results across natural and medical imaging applications.
  • Conclusion: The computational overhead of hard data consistency is identified as a fundamental limitation and challenge for future work.The authors leave reducing this overhead for future research.

A.1 Results on FFHQ

Additional FFHQ experiments show that ReSample remains competitive across inverse-problem tasks. It is strongest on nonlinear deblurring and super resolution, while matching DPS more closely on Gaussian deblurring and inpainting.

  • Results: ReSample outperforms all baselines on nonlinear deblurring and super resolution on FFHQ.The reported pattern is consistent with the authors’ broader multi-task evaluation.
  • Results: ReSample demonstrates performance comparable to DPS on Gaussian deblurring and inpainting.The smallest performance gap from DPS occurs in random inpainting.
  • Results: The largest performance gap from DPS occurs in nonlinear deblurring, while the smallest occurs in random inpainting.The authors note that this pattern mirrors results on CelebA-HQ.

A.2 Additional Results on Box Inpainting

Additional experiments extend ReSample to challenging box inpainting, high-resolution images, varied datasets, noise levels, and inverse-problem settings. Across these evaluations, the method is reported to outperform or remain competitive with the compared baselines.

  • A.2 Additional Results on Box Inpainting: ReSample outperforms PSLD and DPS in the challenging box-inpainting setting on CelebA-HQ.Box inpainting requires recovering images with large regions of missing pixels.
  • A.1 Results on FFHQ: ReSample outperforms all baselines in nonlinear deblurring and super resolution, while remaining comparable to DPS in Gaussian deblurring and inpainting.This task-dependent pattern is also observed in the FFHQ results.
  • High-resolution inverse problems: ReSample is evaluated on high-resolution 512×512×3 images for 70% random inpainting with measurement noise σy = 0.01.The high-resolution extension uses arbitrary-resolution random noise as the initial LDM input.
  • A.4 Discussion on PSLD Baseline: After PSLD hyperparameter tuning and relabeling the baseline as PSLD-LDM, ReSample still largely outperforms the baselines.The revised comparisons include additional FFHQ inpainting experiments.
  • Additional evaluations: Additional figures cover Gaussian deblurring, random inpainting, super resolution, and CT reconstruction under multiple Gaussian-noise levels.The supplementary evaluations include CelebA-HQ, LSUN-Bedroom, FFHQ, and medical CT settings.

A.6 Ablation Studies

The ablations examine stochastic resampling, hard-consistency frequency, computational efficiency, and LDM training efficiency. ReSample balances reconstruction quality against computational cost through its hyperparameters and latent-space design.

  • Stochastic Resampling: Increasing γ substantially improves CelebA-HQ reconstruction quality initially, then plateaus; larger values add fine detail but may introduce noise.γ balances prior consistency against measurement consistency.
  • Skip Step Size: Skip step sizes from 1 to 10 produce very similar CT reconstructions, while size 10 substantially reduces inference time.The authors select 10 to balance reconstruction quality and inference time.
  • Computational Efficiency: The study compares ReSample’s inference time and performance across Gaussian deblurring settings and varying hard-consistency skip sizes for chest CT reconstruction.These evaluations are summarized in Tables 10 and 11.
  • Computational Efficiency: ReSample requires minimal memory for hard data consistency, with its memory advantage widening as pretrained diffusion models grow larger.Applying consistency only on a subset of sampling steps slightly increases inference time relative to DPS, DMPS, and MCG, while ReSample is faster than PSLD.
  • Training Efficiency of LDMs: The authors evaluate LDM training efficiency against DDPMs for medical image synthesis, emphasizing reduced computational resource requirements for LDMs.The motivation is especially relevant to downstream medical-imaging applications.

B Discussion on Hard Data Consistency

Hard data consistency can be imposed through latent or pixel optimization, which differ in reconstruction appearance and computational cost. Combining the formulations can balance image quality and speed, while closed-form or conjugate-gradient solutions are available in specific linear settings.

  • Hard Data Consistency: Hard data consistency is formulated as an optimization problem, and the technique avoids iterative solvers in settings where a computationally efficient solution is available.For noisy measurements, the appendix discusses computing a closed-form solution under appropriate conditions.
  • Optimization Formulations: Latent optimization directly changes the latent variable, producing sharper details but noisier reconstructions; pixel optimization produces smoother but blurrier images.The pixel formulation preserves more high-level semantic information, whereas latent optimization operates directly on the latent representation.
  • Optimization Formulations: Latent optimization is slower because every gradient step backpropagates through the decoder, so combining latent and pixel optimization balances reconstruction speed and image quality.The combined strategy can also accelerate the optimization process.
  • Linear Measurement Operators: For linear noiseless measurements, pixel optimization admits a closed-form solution involving the pseudoinverse, while noisy cases can use a relaxed formulation.Conjugate gradients can implement the required linear-system computation.
  • Linear Measurement Operators: Conjugate-gradient least squares applies to linear inverse problems when the measurement matrix is available, including CT reconstruction using the Radon transform and its adjoint back projection.The same technique is not restricted to CT when the matrix representation of the forward operator is available.

C Implementation Details

The implementation uses staged hard data consistency within DDIM sampling, combining pixel and latent optimization with fixed experimental hyperparameters. Medical-image experiments require fine-tuned latent diffusion backbones, while the appendix also analyzes baseline behavior and Latent-DPS limitations.

  • Implementation Details: Natural-image experiments use 500 DDIM steps, perform pixel optimization in the middle stage and latent optimization in the final stage, and apply consistency every 10 iterations.The same hyperparameters are used across LSUN-Bedroom, FFHQ, and CelebA-HQ.
  • Implementation Details: The natural-image settings use τ = 10^-4, up to 2000 pixel-optimization iterations or 500 latent-optimization iterations, and generally γ = 40.The stochastic-resampling variance uses an adaptive schedule.
  • Medical Imaging: Medical-image models are fine-tuned from pretrained latent-diffusion checkpoints on 2,000 CT images for 100K iterations at learning rate 10^-5.The backbone uses the VQ-4 autoencoder and FFHQ-LDM, with CelebA-LDM also selected as a backbone.
  • Medical Imaging: Medical CT inference uses 1,000 DDIM steps, conjugate-gradient pixel optimization in the middle stage, latent optimization in the final stage, skip step size 10, and γ = 40.The pixel stage uses 50 iterations with κ = 0.9, while latent optimization uses τ = 10^-4.
  • Latent-DPS Analysis: ReSample achieves better measurement consistency than Latent-DPS, supporting hard data consistency as a motivation for reconstruction improvement.The comparison is based on measurement-consistency loss.
  • Failure Analysis: The decoder’s nonlinearity violates the linear-manifold assumption underlying DPS, and the resulting reverse path can leave the data manifold, producing blurry reconstructions.The appendix illustrates this failure by showing that averaging latent vectors can yield blurry, unrealistic images.
  • Failure Analysis: Posterior-mean estimates can be inaccurate at early timesteps, especially when t > 0.5T, causing inconsistent gradient directions and noisy reconstructions.Figure 21 compares predictions at t = 0.5T and t = 0.
  • Baselines: The related-work discussion positions PSLD as an unsupervised soft approach and notes broader limitations in latent diffusion inverse-problem methods.The appendix also documents implementation choices for Latent-DPS, DPS, MCG, DDRM, DMPS, PSLD-LDM, and supervised baselines.

F Deferred Proofs for ReSample

The deferred proofs formalize stochastic encoding and resampling in the latent diffusion process. They establish unbiasedness under measurement consistency, variance reduction, and a bound under a bounded-score assumption.

  • Notation and Pipeline: The appendix defines latent variables, posterior-mean estimates, measurement-consistent samples, and remapped noisy samples used to resume reverse diffusion.The remapped sample ẑ_t is the next sample in the reverse process.
  • Stochastic Encoding: Stochastic encoding yields a sample conditionally independent of the measurement y given the measurement-consistent estimate ẑ_0(y).The proof uses Tweedie’s formula and the diffusion forward process.
  • Proof Strategy: The posterior derivation uses Gaussian conditional distributions and the diffusion noise schedule to characterize the resampled latent variable.The appendix notes that the relevant Gaussian form follows from moment-generating functions.
  • Stochastic Resampling: Under measurement consistency, stochastic resampling is unbiased: E[ẑ_t|y] equals the corresponding unconditional-sample expectation.The theorem assumes y = A(D(ẑ_0(y))).
  • Stochastic Resampling: If the unconditional resampled variance is positive, the resampled sample has lower variance than the stochastically encoded sample.The result is stated as VAR(ẑ_t) < VAR(ẑ̃_t).
  • Theoretical Conditions: Theorem 2 derives an additional result for a diffusion process whose score function is bounded and whose α_t sequence decreases with t.The stated assumption constrains the score function and the noise schedule.
Loading 2307.08123v3…