[논문 리뷰] Self-Supervised MRI Reconstruction with Unrolled Diffusion Models

Yılmaz Korkmaz, Tolga Çukur외 1명|arXiv (Cornell University)|2023. 06. 29.

Advanced Neuroimaging Techniques and ApplicationsMedicine인용 수 3

한 줄 요약

이 논문은 교차 어텐션 트랜스포머와 데이터 일致성 레이어를 결합하여 과도하게 샘플링된 k-공간 데이터를 노이즈 제거하는 자기지도 학습(unrolled) 확산 모델인 SSDiffRecon을 제안한다. 이는 오직 과도하게 샘플링된 k-공간 데이터로만 훈련함으로써 상태최저 수준의 이미지 품질과 속도를 달성하며, 풀스패일드 데이터가 필요한 보조 학습 및 자기지도 학습 기반 모델보다도 뛰어난 성능을 보이며, 역확산 단계를 오직 다섯 번만 수행한다.

ABSTRACT

Magnetic Resonance Imaging (MRI) produces excellent soft tissue contrast, albeit it is an inherently slow imaging modality. Promising deep learning methods have recently been proposed to reconstruct accelerated MRI scans. However, existing methods still suffer from various limitations regarding image fidelity, contextual sensitivity, and reliance on fully-sampled acquisitions for model training. To comprehensively address these limitations, we propose a novel self-supervised deep reconstruction model, named Self-Supervised Diffusion Reconstruction (SSDiffRecon). SSDiffRecon expresses a conditional diffusion process as an unrolled architecture that interleaves cross-attention transformers for reverse diffusion steps with data-consistency blocks for physics-driven processing. Unlike recent diffusion methods for MRI reconstruction, a self-supervision strategy is adopted to train SSDiffRecon using only undersampled k-space data. Comprehensive experiments on public brain MR datasets demonstrates the superiority of SSDiffRecon against state-of-the-art supervised, and self-supervised baselines in terms of reconstruction speed and quality. Implementation will be available at https://github.com/yilmazkorkmaz1/SSDiffRecon.

연구 동기 및 목표

• 기존 MRI 재구성 방법의 한계, 즉 낮은 문맥 감지 능력, 최적화되지 않은 데이터 분포 모델링, 풀스패일드 데이터를 훈련에 반드시 필요로 하는 문제를 해결한다.
• 감독 학습에서 풀스패일드 데이터에 의존하는 문제를 해결하기 위해, 마스크된 k-공간 재구성 기반의 자기지도 학습 전략을 도입한다.
• 확산 기반의 언롤드 아키텍처 내부에 물리 기반의 데이터 일치성 블록을 통합하여 이미지 품질과 재구성 속도를 향상시킨다.
• 역확산 과정에서 교차 어텐션 트랜스포머를 통합하여 문맥 이해력과 특징 표현 능력을 향상시킨다.
• 쌍체의 훈련 데이터가 필요 없이도 최소한의 역확산 단계로도 경쟁 가능한 성능을 달성할 수 있음을 입증한다.

제안 방법

• 역확산 과정에서 노이즈 제거를 위한 교차 어텐션 트랜스포머 블록과 물리적 제약 조건을 강제하는 데이터 일치성 투영을 번갈아가며 통합한 언롤드 아키텍처를 제안한다.
• 과도하게 샘플링된 데이터에서 마스크된 k-공간 샘플을 예측함으로써 자기지도 학습 방식으로 노이즈 제거 네트워크를 훈련시키는 조건부 확산 프로세스를 적용한다.
• 타임스텝과 입력에 특화된 조건부 정보를 트랜스포머 기반 노이즈 제거기 내에 주입하기 위해, 교차 어텐션을 활용한 매핑 네트워크를 사용한다.
• 각 노이즈 제거 단계 이후에 데이터 일치성 레이어를 통합하여 획득한 과도하게 샘플링된 k-공간 데이터에 대한 일치성을 강제한다.
• 확산 과정에서 예측된 노이즈와 진짜 노이즈 간의 차이를 최소화하는 노이즈 제거 손실을 사용하여 모델을 훈련시킨다.
• 오직 다섯 번의 역확산 단계만으로 추론를 수행함으로써, 표준 DDPM 대비 재구성 시간을 크게 단축시킨다.

실험 결과

연구 질문

• RQ1 쌍체의 풀스패일드 데이터가 필요 없이 자기지도 학습 기반의 확산 모델이 고해상도 MRI 재구성을 달성할 수 있는가?
• RQ2 트랜스포머 기반 언롤드 확산 아키텍처에 데이터 일치성 레이어를 통합할 경우 재구성 품질과 속도에 어떤 영향을 미치는가?
• RQ3 CNN 기반 모델 대비 교차 어텐션 기반 모델이 MRI 재구성에서 문맥 감지 능력과 세부 정보 유지 능력 향상에 얼마나 기여하는가?
• RQ4 최소한의 역확산 단계(예: 다섯 번)로도 표준 DDPM의 수천 번의 단계와 경쟁 가능한 성능을 달성할 수 있는가?
• RQ5 마스크된 k-공간 예측 기반의 제안된 자기지도 학습 전략이 감독 학습 및 기타 자기지도 학습 기반 기준 대비 강건성과 일반화 능력에서 어떤가?

주요 결과

• IXI 데이터셋에서 4× 가속도 조건에서 SSDiffRecon은 평균 PSNR 42.3 dB, SSIM 99.3을 달성하여, 두 번째로 우수한 자기지도 학습 방법보다 PSNR 3.55 dB, SSIM 1.96% 향상되었다.
• fastMRI 데이터셋에서 4× 가속도 조건에서 SSDiffRecon은 PSNR 40.1 dB, SSIM 96.5를 기록하여 감독 학습 기반 모델보다 PSNR 0.4 dB, SSIM 0.25% 향상되었다.
• 오직 다섯 번의 역확산 단계만으로도 상태최저 성능을 달성하여, 표준 DDPM 대비 추론 시간을 크게 단축시켰다.
• 시각적 비교 결과, 다른 자기지도 학습 방법들이 흔히 발생시키는 흐림이나 노이즈 문제를 보여주지 않으며, 고주파 세부 정보를 더 잘 유지하는 것으로 확인되었다.
• 절단 실험 결과, 교차 어텐션 트랜스포머와 데이터 일치성 레이어 모두 성능 향상에 필수적이며, 둘 중 하나를 제거할 경우 PSNR와 SSIM가 심각하게 저하됨을 확인했다.
• 자기지도 학습 전략 덕분에 풀스패일드 데이터가 필요한 감독 학습 방법과 동등한 성능을 달성할 수 있었으며, 이는 쌍체의 풀스패일드 데이터가 필요 없음을 의미한다.