[论文解读] Fast and Stable Diffusion Inverse Solver with History Gradient Update

Linchao He, Hongyu Yan+9 more|arXiv (Cornell University)|Jul 22, 2023

MRI in cancer diagnosisMedicine被引用 3

一句话总结

该论文提出了一种零样本方法——潜在扩散迭代重建（LDIR），将预训练的潜在扩散模型（LDM）作为数据先验，集成于迭代CT重建中，利用数据保真项的梯度引导，无需成对数据或反投影矩阵近似。LDIR在稀疏视角和有限角度CT重建任务中达到最先进性能，优于所有无监督与有监督方法，同时实现了512×512高分辨率重建，具有更快推理速度和更低内存占用。

ABSTRACT

Diffusion models have recently been recognised as efficient inverse problem solvers due to their ability to produce high-quality reconstruction results without relying on pairwise data training. Existing diffusion-based solvers utilize Gradient Descent strategy to get a optimal sample solution. However, these solvers only calculate the current gradient and have not utilized any history information of sampling process, thus resulting in unstable optimization progresses and suboptimal solutions. To address this issue, we propose to utilize the history information of the diffusion-based inverse solvers. In this paper, we first prove that, in previous work, using the gradient descent method to optimize the data fidelity term is convergent. Building on this, we introduce the incorporation of historical gradients into this optimization process, termed History Gradient Update (HGU). We also provide theoretical evidence that HGU ensures the convergence of the entire algorithm. It's worth noting that HGU is applicable to both pixel-based and latent-based diffusion model solvers. Experimental results demonstrate that, compared to previous sampling algorithms, sampling algorithms with HGU achieves state-of-the-art results in medical image reconstruction, surpassing even supervised learning methods. Additionally, it achieves competitive results on natural images.

研究动机与目标

• 解决稀疏数据CT重建的病态问题，且无需成对训练数据或反投影矩阵近似。
• 通过使用预训练的潜在扩散模型作为无监督数据先验，实现零样本、高分辨率（512×512）图像重建。
• 通过数据保真项的梯度信息引导潜在扩散采样过程，提升重建质量和效率。
• 相比像素空间迭代方法，降低计算成本和内存消耗，同时保持或超越性能表现。
• 在多样化任务中实现良好泛化，包括医学影像与自然图像修复（如图像修复、超分辨率）。

提出的方法

• LDIR将迭代重建建模为优化问题，其中预训练的无条件潜在扩散模型作为数据先验，替代传统手工设计或学习得到的先验。
• 潜在扩散模型的反向采样过程由数据保真项的梯度引导，使模型能够在无需反投影矩阵的情况下生成与稀疏测量一致的重建结果。
• 引入历史梯度更新机制，在潜在空间的反向采样过程中融合历史梯度信息，以增强样本一致性。
• 该方法完全在潜在空间中运行，降低计算复杂度，实现高效的高分辨率（512×512）重建。
• 框架仅需一次训练，即可应用于多种CT重建任务（如稀疏视角和有限角度扫描），无需微调或结构修改。
• 通过在潜在空间中采用基于得分的采样，避免直接基于投影的优化，使其可适配不同探测器几何结构。

实验结果

研究问题

• RQ1 预训练的潜在扩散模型是否可在无需成对数据或反投影矩阵近似的情况下，有效用作迭代重建中的数据先验？
• RQ2 数据保真项的梯度引导在零样本设置下如何提升重建质量？
• RQ3 历史梯度更新机制是否能增强基于潜在扩散的重建中样本的一致性与收敛稳定性？
• RQ4 潜在空间迭代重建是否在速度、内存效率和重建质量方面优于像素空间方法？
• RQ5 单一统一模型是否能泛化于多样化稀疏测量任务，包括医学CT与自然图像修复？

主要发现

• 在稀疏视角CT重建中，LDIR实现39.01的PSNR和0.9552的SSIM，显著优于次佳方法DPS（28.55 PSNR，0.8140 SSIM）。
• 在有限角度CT重建中，LDIR实现39.77的PSNR和0.9612的SSIM，优于DPS（28.97 PSNR，0.8242 SSIM）和DIR（32.82 PSNR，0.8898 SSIM）。
• LDIR实现每秒36.67次迭代，内存消耗仅为4268 MB，相比像素基方法（如DIR：20.75次/秒，6338 MB）在速度和内存效率方面表现更优。
• 在自然图像任务中，LDIR与SOTA方法DPS相比表现相当，尤其在99%随机图像修复和超分辨率任务中，伪影极少。
• 消融实验表明，采用历史梯度引导的潜在空间重建性能更优（39.01 PSNR），显著优于像素空间方法（DPS为28.55 PSNR），且在速度与内存效率方面提升显著。
• 该方法在任务间泛化能力强，展现出在医学影像与自然图像修复任务中的优异性能，包括高度稀疏测量场景。