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

ひとこと要約

本稿では、ラティント・ディフュージョン・モデル（LDM）をデータ事前分布として統合するゼロショット手法、潜在変数拡散反復再構成（LDIR）を提案する。LDIRは、データ適合性項からの勾配ガイダンスを用いて、ペairedデータや逆投影行列の近似を必要とせず、反復的CT再構成にLDMを統合する。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はPSNR 39.01、SSIM 0.9552を達成し、次善の手法（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は1秒間に36.67イテレーションを実行し、メモリ使用量は4268 MBにとどまり、ピクセルベース手法（例：DIR：20.75 it/s、6338 MB）に比べて速度とメモリ効率が優れていた。
• 自然画像タスクにおいて、LDIRはSOTA手法のDPSと同等の結果を達成し、特に99％のランダム穴埋めと最小限のアーチファクトを伴う超解像において優れた性能を示した。
• アブレーションスタディの結果、履歴勾配ガイダンスを用いた潜在空間再構成は、ピクセル空間手法（DPS：28.55 PSNR）と比較して優れた性能（39.01 PSNR）を示し、速度とメモリ効率の面で顕著な向上が得られた。
• 本手法はタスク間で優れた一般化性能を示し、医療画像処理および自然画像修復の両方で強く、極めてスパースな測定にも対応した優れた性能を発揮した。