[论文解读] Decomposed Diffusion Sampler for Accelerating Large-Scale Inverse Problems

Hyungjin Chung, Suhyeon Lee+1 more|arXiv (Cornell University)|Mar 10, 2023

Statistical Methods and Inference被引用 6

一句话总结

该论文提出了一种新颖的扩散采样策略——分解扩散采样器（Decomposed Diffusion Sampler, DDS），通过将Krylov子空间方法与扩散模型协同结合，加速医学影像中大规模逆问题的求解。通过证明由Tweedie去噪公式诱导的切空间构成Krylov子空间，DDS实现了共轭梯度（CG）初始化，无需流形约束梯度计算即可保持数据一致性，从而在多线圈MRI和3D CT任务中实现SOTA重建质量，并将推理速度提升80倍以上。

ABSTRACT

Krylov subspace, which is generated by multiplying a given vector by the matrix of a linear transformation and its successive powers, has been extensively studied in classical optimization literature to design algorithms that converge quickly for large linear inverse problems. For example, the conjugate gradient method (CG), one of the most popular Krylov subspace methods, is based on the idea of minimizing the residual error in the Krylov subspace. However, with the recent advancement of high-performance diffusion solvers for inverse problems, it is not clear how classical wisdom can be synergistically combined with modern diffusion models. In this study, we propose a novel and efficient diffusion sampling strategy that synergistically combines the diffusion sampling and Krylov subspace methods. Specifically, we prove that if the tangent space at a denoised sample by Tweedie's formula forms a Krylov subspace, then the CG initialized with the denoised data ensures the data consistency update to remain in the tangent space. This negates the need to compute the manifold-constrained gradient (MCG), leading to a more efficient diffusion sampling method. Our method is applicable regardless of the parametrization and setting (i.e., VE, VP). Notably, we achieve state-of-the-art reconstruction quality on challenging real-world medical inverse imaging problems, including multi-coil MRI reconstruction and 3D CT reconstruction. Moreover, our proposed method achieves more than 80 times faster inference time than the previous state-of-the-art method. Code is available at https://github.com/HJ-harry/DDS

研究动机与目标

• 解决尽管扩散模型在大规模医学逆问题中具有优越的重建质量与泛化能力，但其推理速度缓慢的问题。
• 将经典优化技术（特别是共轭梯度等Krylov子空间方法）与现代扩散采样相结合，以提升效率。
• 通过利用去噪样本中的几何结构，消除扩散基逆求解器中计算成本高昂的流形约束梯度（MCG）计算需求。
• 开发一种通用采样框架，兼容多种扩散参数化形式（VE、VP）与前向模型，增强在不同成像模态中的适应性。
• 在多线圈MRI与3D稀疏视角CT等具有挑战性的现实世界任务中，实现SOTA重建性能，且计算开销极低。

提出的方法

• 提出一种新型扩散采样策略——分解扩散采样器（DDS），通过Tweedie公式从去噪样本初始化共轭梯度（CG）。
• 理论上证明：若去噪样本处的切空间构成Krylov子空间，则CG更新将保持在此空间内，从而在无需MCG的情况下确保数据一致性。
• 利用由残差向量与前向模型Hessian生成的Krylov子空间结构，实现高效且一致的优化。
• 将方法集成至基于得分的扩散模型与去噪扩散模型（VE与VP设置），使其适用于多种扩散框架。
• 以单次CG迭代替代标准Langevin动力学或迭代MCMC步骤，显著减少函数评估次数（NFE）。
• 通过Tweedie公式执行单次去噪步骤，生成位于Krylov子空间内的初始估计，随后通过快速CG迭代强制实现数据一致性。

实验结果

研究问题

• RQ1 经典Krylov子空间方法能否与现代扩散采样有效结合，以加速医学影像中的逆问题求解？
• RQ2 在何种条件下，通过Tweedie公式从去噪样本诱导的切空间会形成Krylov子空间，从而实现基于CG的高效优化？
• RQ3 消除对流形约束梯度计算的需求是否能带来显著的速度提升，同时不牺牲重建质量？
• RQ4 所提出方法能否在不同扩散参数化形式（VE、VP）与前向模型（如并行成像、稀疏视角CT）间实现泛化？
• RQ5 该方法在保持或提升真实世界医学影像任务中重建保真度的前提下，能在多大程度上减少函数评估次数（NFE）？

主要发现

• DDS在多线圈MRI与3D稀疏视角CT任务中实现了SOTA重建质量，PSNR与SSIM指标均优于以往基于扩散的方法。
• 在1D×4加速的多线圈MRI中，DDS仅用19次NFE即达到32.12 dB的PSNR与0.887的SSIM，显著优于先前SOTA方法Score-MRI（4000次NFE，33.38 dB，0.887 SSIM）。
• 在3D稀疏视角CT重建中，DDS以49次NFE实现34.86 dB的PSNR与0.892的SSIM，优于DiffusionMBIR（4000次NFE，34.30 dB，0.877 SSIM）与E2E-VarNet。
• 与先前SOTA方法（Score-MRI）相比，推理时间减少80倍以上，仅用49次NFE即实现高保真重建。
• DDS在多种采样模式（均匀、高斯、变密度）与模态（2D与3D）下均保持高性能，展现出强大的泛化能力。
• 理论分析证实，在较弱条件下，Tweedie去噪公式诱导的切空间确实构成Krylov子空间，从而可通过CG实现精确数据一致性，无需MCG。