[论文解读] Improving Diffusion Models for Inverse Problems using Manifold Constraints

Hyungjin Chung, Byeongsu Sim+2 more|arXiv (Cornell University)|Jun 2, 2022

Advanced Neuroimaging Techniques and Applications被引用 101

一句话总结

本文提出一个流形约束梯度（MCG）项，作为无条件扩散模型求解器的补充，用于逆问题的迭代路径保持在数据流形上，并在修复、着色、稀疏视角 CT 等任务中提升重建效果。它提供理论依据和大量实验，显示相较于先前的扩散基方法与监督方法有显著提升。

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 suboptimal 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 https://github.com/HJ-harry/MCG_diffusion

研究动机与目标

• Motivate unsupervised diffusion-model solvers for general linear inverse problems without task-specific training.
• Introduce a manifold-constrained gradient term to keep diffusion paths on the data manifold.
• Integrate Tweedie-based denoising with a data-consistency gradient to improve stability and accuracy.
• Provide theoretical justification showing the gradient stays in the tangent space of the data manifold.
• Demonstrate empirical gains across inpainting, colorization, and CT reconstruction.

提出的方法

• Use an unconditional score function within a reverse diffusion step and a measurement-consistency projection as in prior work.
• Introduce manifold constrained gradient (MCG) that projects the gradient of the measurement term onto the data manifold via Q_i and J_{Q_i}.
• Derive the MCG correction as a tangent-space projection: ∂/∂x_i ||W(y−Hx̂0)||^2 on the tangent space T_{x̂0}M.
• Incorporate Tweedie’s denoising philosophy to justify the single-step denoising correction and connect it with the manifold constraint.
• Combine the MCG step with the usual data-consistency projection to preserve data fidelity while staying on the manifold.
• Provide theoretical results (Theorem 1) and geometric interpretation under a strong manifold assumption.

实验结果

研究问题

• RQ1 Can a manifold-aware correction term improve diffusion-based solvers for inverse problems without problem-specific training?
• RQ2 Does the manifold constrained gradient keep diffusion trajectories closer to the data manifold and reduce error accumulation?
• RQ3 How does the MCG interact with standard projection-based data consistency and Tweedie’s denoising in practice?
• RQ4 What performance gains arise across diverse inverse problems like inpainting, colorization, and sparse-view CT?
• RQ5 Is the method robust to different measurement masks and NFEs (number of function evaluations)?

主要发现

• MCG keeps the iterative path on the data manifold, reducing error accumulation.
• MCG combined with projection achieves data consistency and improved reconstruction over prior diffusion-based methods across multiple tasks.
• MCG outperforms Score-SDE, DDRM, RePAINT, and other baselines in inpainting (FFHQ/ImageNet), colorization (FFHQ/LSUN-bedroom), and CT reconstruction (AAPM dataset).
• Ablations show that using only MCG improves LPIPS but can harm measurement consistency, while combining MCG with projection yields best results.
• Qualitative results demonstrate higher fidelity reconstructions with MCG, including reduced boundary artifacts in inpainting.