[論文レビュー] 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

ひとこと要約

本稿では、医療画像における大規模な逆問題を高速化するために、Krylov部分空間法と拡散モデルを包括的に統合する新規な拡散サンプリング戦略、分解型拡散サンプラー（DDS）を提案する。Tweedieのノイズ除去公式によって誘導される接空間がKrylov部分空間を形成することを証明することで、多様体制約付き勾配計算を必要とせず、データ整合性を維持した共役勾配初期化を可能にした。これにより、マルチコイルMRIおよび3D CTタスクにおいて、先行研究と比較して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

研究の動機と目的

• 拡散モデルの優れた再構成品質と一般化性能にもかかわらず、大規模な医療画像逆問題における推論速度の遅さに取り組む。
• 具体的には、共役勾配（CG）を含む古典的最適化手法であるKrylov部分空間法と、現代の拡散サンプリングを統合し、効率性を向上させる。
• ノイズ除去されたサンプルに内在する幾何的構造を活用することで、拡散ベースの逆問題ソルバーにおける計算コストの高い多様体制約付き勾配（MCG）計算を排除する。
• さまざまな拡散パrametrization（VE、VP）および前方モデルに対応可能な汎用的なサンプリングフレームワークを構築し、異なる画像モodalitiesへの適応性を高める。
• マルチコイルMRIや3DスパースビューCTといった現実世界の挑戦的タスクにおいて、最小限の計算オーバーヘッドで最先端の再構成性能を達成する。

提案手法

• Tweedieの公式を用いてノイズ除去されたサンプルから共役勾配（CG）を初期化する新規な拡散サンプリング戦略、分解型拡散サンプラー（DDS）を提案する。
• 理論的に、ノイズ除去されたサンプルにおける接空間がKrylov部分空間を形成する場合、CG更新がこの空間内に留まることが保証され、MCG計算を必要とせずデータ整合性を維持できることを証明する。
• 前方モデルのヘッセ行列と残差ベクトルによって生成されるKrylov部分空間の構造を活用し、効率的かつ整合性のある最適化を実現する。
• スコアベースおよびノイズ除去型拡散モデル（VEおよびVP設定）の両方へ統合し、異なる拡散フレームワークに適用可能である。
• 標準的なランジュバンダイナミクスや反復的MCMCステップの代わりに、ノイズ除去されたサンプルから初期化された1回のCGステップを導入し、関数評価回数（NFE）を著しく削減する。
• Tweedieの公式による1回のノイズ除去ステップでKrylov部分空間内に位置する初期推定値を生成し、その後に高速なCG反復を実行してデータ整合性を強制する。

実験結果

リサーチクエスチョン

• RQ1 古典的Krylov部分空間法を現代の拡散サンプリングと効果的に統合することで、医療画像における逆問題解決の高速化が可能か？
• RQ2 Tweedieの公式によるノイズ除去されたサンプルから誘導される接空間が、どのような条件下でKrylov部分空間を形成するか？これにより、CGに基づく精緻化が効率的に行えるか？
• RQ3 多様体制約付き勾配計算の必要性を排除することで、再構成品質を損なわず、顕著な高速化が達成できるか？
• RQ4 提案手法が、さまざまな拡散パラメータ設定（VE、VP）および前方モデル（例：並列イメージング、スパースビューCT）に一般化可能か？
• RQ5 関数評価回数（NFE）をどの程度削減できるか？また、実世界の医療画像タスクにおいて再構成忠実度を維持または向上させられるか？

主な発見

• DDSはマルチコイルMRIおよび3DスパースビューCTにおいて、先行する拡散ベース手法を上回る最先端の再構成品質を達成した。PSNRおよびSSIMの両面で優れた性能を示した。
• 1D×4の加速率を有するマルチコイルMRIでは、19 NFEで32.12 dBのPSNRおよび0.887のSSIMを達成し、従来のSOTA手法（Score-MRI：4000 NFEで33.38 dB、0.887 SSIM）を上回った。
• 3DスパースビューCT再構成では、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で高精度な結果を達成した。
• 均一、ガウス分布、可変密度の多様なサンプリングパターンおよび2D・3Dの多様なモダリティにおいても、高い性能を維持し、優れた一般化能力を示した。
• 理論的分析により、Tweedieのノイズ除去公式から誘導される接空間が、緩い条件下でもKrylov部分空間を形成することを確認した。これにより、MCGを必要とせず、正確なデータ整合性をCGによって実現できることが裏付けられた。