[논문 리뷰] Solving 3D Inverse Problems using Pre-trained 2D Diffusion Models

Hyungjin Chung, Dohoon Ryu외 3명|arXiv (Cornell University)|2022. 11. 19.

Advanced Neuroimaging Techniques and Applications인용 수 24

한 줄 요약

DiffusionMBIR를 제안하는 방법으로, 2D 확산 사전(prior)에 모델 기반 사전을 보강하여 단일 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 사전(prior)을 학습하지 않고 확산 모델로 3D 역문제 해결을 촉진한다.
• 2D 사전 학습된 확산 모델을 활용하여 일관된 3D 부피를 재구성한다.
• 슬라이스 간 데이터 일관성을 강제하기 위해 ADMM을 통한 MBIR 스타일의 규제와 확산 디노이징을 통합한다.

제안 방법

• 2D 확산 모델을 슬라이스별로 사용하여 디노이징하고, 3D 데이터 일관성을 통해 z축 일관성을 유지한다.
• 3D 부피에 대한 ADMM 업데이트를 통해 확산 사전에 z 방향 TV 사전을 보강한다.
• 데이터 충실도 항과 z 방향 규제를 가진 3D MBIR 유사 최적화로 재구성을 형식화한다.
• z 방향의 희소성을 위한 공액 기울기(CG) 스텝과 소프트 임계화를 사용하여 ADMM 하위 문제를 풀이한다.
• 계산량을 줄이기 위해 SDE 단계 간 ADMM 상태를 재사용하는 변수 공유 속도가 빠른 버전을 사용한다.

실험 결과

연구 질문

• RQ1 3D 확산 모델을 학습하지 않고도 사전 학습된 2D 확산 모델을 사용하여 3D 역문제를 해결할 수 있는가?
• RQ2 2D 확산 사전에 3D MBIR 유사한 z 방향 규제를 보강하면 축 전체에서 일관된 3D 재구성이 얻어지는가?
• RQ3 제안된 DiffusionMBIR이 SV-CT, LA-CT, CS-MRI와 같은 작업을 단일 일반 GPU에서 실행할 만큼 메모리 효율적인가?
• RQ4 기존의 확산 기반 및 완전 감독 방법과 비교하여, in-distribution 및 out-of-distribution 3D 의학 영상 데이터에서 DiffusionMBIR의 성능은 어떤가?

주요 결과

• DiffusionMBIR은 sparse-view CT, limited-angle CT, 및 compressed-sensing MRI에서 최첨단 재구성을 달성한다.
• 축, 시상, 관상면 전반에서 일관된 3D 재구성을 제공한다.
• DiffusionMBIR은 높은 PSNR과 SSIM을 유지하고, out-of-distribution 일반화가 강함을 보인다.
• ADMM/CG 반복이 최소화된 빠른 버전(M=1, K=1)은 낮은 계산으로도 높은 충실도 결과를 제공한다.