[논문 리뷰] Decomposed Diffusion Sampler for Accelerating Large-Scale Inverse Problems

Hyungjin Chung, Suhyeon Lee외 1명|arXiv (Cornell University)|2023. 03. 10.

Statistical Methods and Inference인용 수 6

한 줄 요약

이 논문은 의료 영상에서 대규모 역문제를 가속화하기 위해 키로프 부분공간 방법과 확산 모델을 유기적으로 융합한 새로운 확산 샘플링 전략인 분해된 확산 샘플러(DDS)를 제안한다. 트위디의 노이즈 제거 공식이 유도하는 탄젠트 공간가 키로프 부분공간을 이룬다는 것을 증명함으로써, 다각도적 일致성을 유지하면서도 다양체 제약 경사 하강 계산이 필요 없는 코그래디언(Conjugate Gradient, CG) 초기화를 가능하게 하여, 다중 코ils 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)과 같은 키로프 부분공간 방법을 현대적 확산 샘플링과 융합하여 효율성을 향상시킨다.
• 노이즈 제거된 샘플의 기하학적 구조를 활용해 확산 기반 역문제 해법에서 계산 비용이 큰 다각도 제약 경사 하강(MCG) 계산이 필요 없도록 한다.
• 다양한 확산 매개변수화(VE, VP) 및 전방 모델에 호환 가능한 일반 목적의 샘플링 프레임워크를 개발하여 다양한 영상 모odalities에 대한 적응성을 높인다.
• 다중 코일 MRI 및 3D 희소 시야 CT와 같은 도전적인 실세계 작업에서 최첨단 재구성 성능을 달성하면서도 계산 오버헤드를 최소화한다.

제안 방법

• 트위디의 공식을 사용해 노이즈 제거된 샘플에서 출발하는 코그래디언(CG) 초기화를 가능하게 하는 새로운 확산 샘플링 전략인 분해된 확산 샘플러(DDS)를 제안한다.
• 이론적으로, 노이즈 제거된 샘플에서 유도된 탄젠트 공간가 키로프 부분공간을 이룬다면, CG 업데이트가 이 공간 내에 유지되어 MCG 없이도 데이터 일치성을 유지함을 증명한다.
• 전방 모델의 잔차 벡터와 헤시안 행렬이 생성하는 키로프 부분공간의 구조를 활용하여 효율적이고 일관된 최적화를 가능하게 한다.
• 스코어 기반 및 노이즈 제거 기반 확산 모델(VE 및 VP 설정)에 모두 통합하여 다양한 확산 프레임워크에 적용 가능하도록 한다.
• 표준 랭글레인 역학 또는 반복적 MCMC 단계를 노이즈 제거된 샘플에서 초기화된 단일 CG 단계로 대체하여 기능 평가 횟수(NFE)를 극적으로 감소시킨다.
• 트위디의 공식을 통해 단일 노이즈 제거 단계를 수행해 키로프 부분공간 내에 위치한 초기 추정치를 생성한 후, 데이터 일치성을 확보하기 위해 빠른 CG 반복을 수행한다.

실험 결과

연구 질문

• RQ1 고전적인 키로프 부분공간 방법이 현대적 확산 샘플링과 효과적으로 융합되어 의료 영상에서의 역문제 해결을 가속화할 수 있는가?
• RQ2 트위디의 공식을 통해 노이즈 제거된 샘플에서 유도된 탄젠트 공간가 키로프 부분공간을 이루는 조건는 무엇이며, 이는 효율적인 CG 기반 정밀 조정을 가능하게 하는가?
• RQ3 다양체 제약 경사 하강 계산이 필요 없어지는 것이 재구성 품질을 손상시키지 않으면서도 상당한 속도 향상을 이끌 수 있는가?
• RQ4 제안된 방법이 다양한 확산 매개변수화(VE, VP) 및 전방 모델(예: 병렬 영상, 희소 시야 CT)에 일반화 가능한가?
• RQ5 실세계 의료 영상 작업에서 재구성 정밀도를 유지하거나 향상시키면서도, 기능 평가 횟수(NFE)를 얼마나 줄일 수 있는가?

주요 결과

• DDS는 다중 코일 MRI 및 3D 희소 시야 CT에서 최첨단 재구성 품질을 달성하며, 이전의 확산 기반 방법보다 PSNR 및 SSIM 모두에서 승리한다.
• 1D ×4 가속도를 적용한 다중 코일 MRI에서 DDS는 단지 19개의 NFE로 32.12 dB의 PSNR와 0.887의 SSIM 성능을 기록했으며, 이는 이전 최첨단 방법인 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을 모두 능가한다.
• 이전 최첨단 방법(Score-MRI) 대비 추론 시간을 80배 이상 단축시켜 단지 49개의 NFE로도 고해상도 재구성 결과를 도출했다.
• 균일, 가우시안, 가변 밀도 등 다양한 샘플링 패턴과 2D 및 3D 모달리티 전반에서 높은 성능 유지를 보이며 강력한 일반화 능력을 입증했다.
• 이론적 분석을 통해 트위디의 노이즈 제거 공식에서 유도된 탄젠트 공간가 미약한 조건 하에서도 키로프 부분공간을 이룬다는 것을 확인했으며, 이는 MCG 없이도 정확한 데이터 일치성을 보장하는 코그래디언 기반 최적화를 가능하게 한다.