[論文レビュー] Solving 3D Inverse Problems using Pre-trained 2D Diffusion Models

Hyungjin Chung, Dohoon Ryu+3 more|arXiv (Cornell University)|Nov 19, 2022

Advanced Neuroimaging Techniques and Applications被引用数 24

ひとこと要約

DiffusionMBIRを提案する。2D拡散事前分布をモデルベースの事前分布で補強し、単一GPU上で効率的に3D医用画像再構成タスクを解く方法。

ABSTRACT

Diffusion models have emerged as the new state-of-the-art generative model with high quality samples, with intriguing properties such as mode coverage and high flexibility. They have also been shown to be effective inverse problem solvers, acting as the prior of the distribution, while the information of the forward model can be granted at the sampling stage. Nonetheless, as the generative process remains in the same high dimensional (i.e. identical to data dimension) space, the models have not been extended to 3D inverse problems due to the extremely high memory and computational cost. In this paper, we combine the ideas from the conventional model-based iterative reconstruction with the modern diffusion models, which leads to a highly effective method for solving 3D medical image reconstruction tasks such as sparse-view tomography, limited angle tomography, compressed sensing MRI from pre-trained 2D diffusion models. In essence, we propose to augment the 2D diffusion prior with a model-based prior in the remaining direction at test time, such that one can achieve coherent reconstructions across all dimensions. Our method can be run in a single commodity GPU, and establishes the new state-of-the-art, showing that the proposed method can perform reconstructions of high fidelity and accuracy even in the most extreme cases (e.g. 2-view 3D tomography). We further reveal that the generalization capacity of the proposed method is surprisingly high, and can be used to reconstruct volumes that are entirely different from the training dataset.

研究の動機と目的

• 3D事前分布を訓練せずに、拡散モデルを用いて3D逆問題を解くことを動機づける。
• 2Dの事前学習済み拡散モデルを活用して、一貫性のある3D体積を再構成する。
• ADMMを介して拡散ノイズ除去とMBIR風正則化を統合し、スライス間のデータ整合性を強制する。

提案手法

• 2D拡散モデルをスライスごとに使用してノイズ除去を行い、3Dデータ整合性を介してz軸の一貫性を維持する。
• 3D体積上でのADMM更新を通じて、z方向のTV事前分布を用いて拡散事前分布を補強する。
• データ忠実性項とz方向正則化を持つ3D MBIR風最適化として再構成を定式化する。
• ADMMの部分問題を共役勾配法とソフトスレッショルディングで解き、z方向の疎性を得る。
• 計算量を削減するために、SDEステップ間でADMM状態を再利用する変数共有の高速版を採用する。

実験結果

リサーチクエスチョン

• RQ1 3D拡散モデルを訓練せずに、事前学習済みの2D拡散モデルを用いて3D逆問題を解くことは可能か？
• RQ2 2D拡散事前分布を3D MBIR風のz方向正則化で補強すると、全軸で一貫した3D再構成が得られるか？
• RQ3 提案手法のDiffusionMBIRは、SV-CT、LA-CT、CS-MRIのようなタスクを単一の一般的なGPUで動作させるのに十分なメモリ効率があるか？
• RQ4 既存の拡散ベースおよび完全教師あり手法と比較して、DiffusionMBIRは分布内および分布外の3D医用画像データでどのように性能を示すか。

主な発見

• DiffusionMBIRは、スパースビューCT、限角度CT、圧縮センサーMRIで最先端の再構成を達成する。
• 本手法は軸位、矢状、冠状の面全てで一貫した3D再構成をもたらす。
• DiffusionMBIRは高いPSNRとSSIMを維持し、分布外一般化能力が強いことを示す。
• 最小のADMM/CG反復（M=1, K=1）を用いた高速版は、計算負荷を抑えつつ高忠実度の結果を提供する。