Improving Diffusion Models for Inverse Problems using Manifold Constraints
Hyungjin Chung∗,1 Byeongsu Sim∗,2 Dohoon Ryu1 Jong Chul Ye3,1,2 1 Dept. of Bio and Brain Engineering 2 Dept. of Mathematical Sciences 3Kim Jaechul Graduate School of AI ∗Equal contribution Korea Advanced Institute of Science and Technology (KAIST) {hj.chung, byeongsu.s, dh.ryu,
Abstract
Recently, diffusion models have been used to solve various inverse problems in an unsupervised manner with appropriate modifications to the sampling process. However, the current solvers, which recursively apply a reverse diffusion step followed by a projection-based measurement consistency step, often produce sub-optimal results. By studying the generative sampling path, here we show that current solvers throw the sample path off the data manifold, and hence the error accumulates. To address this, we propose an additional correction term inspired by the manifold constraint, which can be used synergistically with the previous solvers to make the iterations close to the manifold. The proposed manifold constraint is straightforward to implement within a few lines of code, yet boosts the performance by a surprisingly large margin. With extensive experiments, we show that our method is superior to the previous methods both theoretically and empirically, producing promising results in many applications such as image inpainting, colorization, and sparse-view computed tomography. Code available here
中文速览
扩散模型（diffusion model）在无监督求解图像重建等逆问题时，现有方法每一步都先做反向扩散、再做投影式测量一致性约束，但这两步叠加会把采样路径推离数据流形，导致误差逐步积累、结果变差。针对这一问题，作者提出了"流形约束梯度"（Manifold Constrained Gradient，MCG）：利用 Tweedie 公式一步估计出干净图像，再对测量误差关于该干净估计求梯度作为额外修正项，从理论上可证明这一修正能让迭代路径贴近数据流形的切空间，与原有投影步骤协同工作。该方法实现极为简洁（仅需数行代码改动），却在图像修复、上色、稀疏角度 CT 重建等多项任务上带来了显著的性能提升，同时避免了 RePaint 等方法反复正向/反向扩散所带来的巨大计算开销。这项工作揭示了现有扩散模型逆问题求解器的几何缺陷，并给出了一个理论严谨、实用高效的通用修复方案，对将扩散模型用于医学影像、图像修复等实际场景具有重要价值。
原文 arXiv:2206.00941;中英对照 + 大白话阅读 https://aha.fim.ai/paper/2206.00941v3