[论文解读] A Variational Perspective on Solving Inverse Problems with Diffusion Models

Morteza Mardani, Jiaming Song+2 more|arXiv (Cornell University)|May 7, 2023

Generative Adversarial Networks and Image Synthesis被引用 36

一句话总结

本文提出 RED-diff，一种用于在带扩散先验的逆问题中求解的变分采样器，采用基于 KL 的目标函数以及利用所有扩散时间步的去噪正则化。

ABSTRACT

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.

研究动机与目标

• 推动使用扩散先验对多种逆问题进行普遍、无需训练的求解。
• 通过采用变分推断来近似 p(x0|y)，以克服难以处理的后验难题。
• 发展一种去噪正则化（通过去噪）视角（RED-diff），利用跨越扩散时间步的去噪器。
• 通过带有轻量迭代的随机优化，提供一种实用、高效的采样机制。

提出的方法

• 将后验推断表述为 q(x0|y) 与 p(x0|y) 之间的 KL 最小化。
• 推导出一个变分界，通过扩散先验得到一个分数匹配（score-matching）正则化项。
• 引入一个带权重的分数匹配目标，利用多个扩散步骤的去噪（RED-diff）。
• 证明在使用合适的权重函数时，正则化项的梯度可以在不通过分数网络反向传播的情况下计算。
• 将采样近似为对 mu 的随机优化，使用简单、可处理的损失，结合重建项和正则化项。
• 将该方法与正则化-去噪（RED）框架联系起来，并讨论优化上的优势。

Figure 1: The schematic diagram of our proposed variational sampler (RED-diff). The forward denoising diffusion process gradually adds noise to the estimate $\mu$ . The denoisers of the backward diffusion process apply score-matching regularization to the measurement matching loss. The refined estim

实验结果

研究问题

• RQ1 在基于扩散模型的逆问题中，变分 KL 框架是否能够为给定 y 的 x0 提供一个原理性后验推断？
• RQ2 在所有扩散时间步中加入带权重的去噪分数是否会比现有方法提升恢复质量？
• RQ3 采样是否可以表述为高效的带有轻量迭代的随机优化，而不需要分数雅可比矩阵？
• RQ4 与最先进的采样器相比，RED-diff 在线性和非线性图像恢复任务中的表现如何？
• RQ5 时间步权重和采样策略对保真度与感知质量的实际影响是什么？

主要发现

• RED-diff 在图像修复中的表现优于最先进的采样器（DPS、Pi-GDM 和 DDRM），在 PSNR、SSIM、KID、LPIPS 和 top-1 精度方面。
• 该方法产生轻量迭代，无分数雅可比矩阵，提高了内存效率和对 GPU 的友好性。
• 基于去噪信噪比的加权机制提高了早期扩散步的权重、降低了后续步的权重，以获得更好的正则化。
• 降序时间步采样（从 T 到 1）相对于其他采样策略，提升了保真度和感知质量。
• 消融实验表明，去噪器加权和时间采样策略是平衡内容与细节和稳定性的关键参数。
• 在非线性逆问题（相位检索、HDR、去模糊等）上，RED-diff 相对于 DPS 显示出更优的指标，在适用场景下也优于 Pi-GDM 和 DDRM。

Figure 2: Comparison of the proposed variational sampler with alternatives for inpainting representative ImageNet examples. Each sampler is tuned for the best performance.