[论文解读] Solving Inverse Problems with Latent Diffusion Models via Hard Data Consistency

Bowen Song, Soo Min Kwon+4 more|arXiv (Cornell University)|Jul 16, 2023

Generative Adversarial Networks and Image Synthesis被引用 13

一句话总结

ReSample 使用 latent diffusion models 结合 hard data consistency 来解决一般线性和非线性反问题，在降低内存使用的同时取得了最先进的结果。

ABSTRACT

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 extit{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.

研究动机与目标

• 激励将 latent diffusion models (LDMs) 作为高维反问题的高效先验。
• 引入 hard data consistency，在 LDMs 的反向采样过程中强制确保测量保真度。
• 开发一种重采样方案，将测量一致的潜在样本映射回数据流形。
• 提供理论保证和实证证据，显示性能和内存效率的提升。
• 证明其在自然与医学成像中的广泛线性与非线性反问题上的适用性。

提出的方法

• 通过编码器 E 和解码器 D，在潜在变量 z 下在潜在空间中操作扩散模型。
• 通过对 z 求解一个受约束优化，以最小化 ||y - A(D(z))||^2，用硬数据一致性替代无条件的反向更新。
• 以基于 Tweedie 的估计 z0_hat 初始化潜在优化，使之尽可能接近真实潜在。
• 引入随机重采样，将优化后的潜在变量映射回带噪数据流形，形成一个后验样本以继续反向扩散。
• 使用跳跃步机制，仅在选定的反向时间步上应用数据一致性，以在精度和计算之间取得平衡。
• 给出理论结果，显示重采样潜在变量的方差下降且无偏，并通过超参数 sigma^2 讨论方差控制。

实验结果

研究问题

• RQ1 Can latent diffusion models (LDMs) effectively serve as priors for general (linear and nonlinear) inverse problems?
• RQ2 Does hard data consistency via latent-space optimization improve measurement fidelity and image quality compared to soft gradient-based updates?
• RQ3 Does the proposed stochastic resampling reduce variance and better preserve data-consistency during reverse diffusion in the latent space?
• RQ4 What are the memory and computational implications of ReSample versus pixel-space diffusion approaches across natural and medical imaging tasks?

主要发现

• ReSample 在自然图像和医学图像的若干线性与非线性反问题上实现了最先进的性能。
• Hard data consistency combined with latent-space optimization and stochastic resampling yields higher PSNR/SSIM and lower LPIPS than competitive methods.
• Stochastic resampling reduces variance and produces smoother reconstructions that remain measurement-consistent, especially as diffusion time t approaches 0.
• The method offers memory efficiency advantages by operating in latent space, with substantial reductions in memory usage compared to pixel-space baselines.
• Ablations show the benefit of resampling over stochastic encoding and the value of resampling frequency in CT reconstruction tasks.