[논문 리뷰] Self-Consistent Recursive Diffusion Bridge for Medical Image Translation

Fuat Arslan, Bilal Kabaş외 3명|arXiv (Cornell University)|2024. 05. 10.

Brain Tumor Detection and Classification인용 수 13

한 줄 요약

SelfRDB는 노이즈가 추가된 소스를 엔드포인트로 사용하는 포워드 프로세스와 자기 일관적인 재귀 역샘플링으로 다중 모달 의료 영상 변환에서 우수한 성능을 달성하는 확산-브리지 프레임워크를 도입한다( MRI-CT, 다중 대조 MRI ).

ABSTRACT

Denoising diffusion models (DDM) have gained recent traction in medical image translation given improved training stability over adversarial models. DDMs learn a multi-step denoising transformation to progressively map random Gaussian-noise images onto target-modality images, while receiving stationary guidance from source-modality images. As this denoising transformation diverges significantly from the task-relevant source-to-target transformation, DDMs can suffer from weak source-modality guidance. Here, we propose a novel self-consistent recursive diffusion bridge (SelfRDB) for improved performance in medical image translation. Unlike DDMs, SelfRDB employs a novel forward process with start- and end-points defined based on target and source images, respectively. Intermediate image samples across the process are expressed via a normal distribution with mean taken as a convex combination of start-end points, and variance from additive noise. Unlike regular diffusion bridges that prescribe zero variance at start-end points and high variance at mid-point of the process, we propose a novel noise scheduling with monotonically increasing variance towards the end-point in order to boost generalization performance and facilitate information transfer between the two modalities. To further enhance sampling accuracy in each reverse step, we propose a novel sampling procedure where the network recursively generates a transient-estimate of the target image until convergence onto a self-consistent solution. Comprehensive analyses in multi-contrast MRI and MRI-CT translation indicate that SelfRDB offers superior performance against competing methods.

연구 동기 및 목표

• 소스 모달리티와 대상 모달리티 간의 표준 확산 모델을 넘어서는 향상된 의료 영상 변환의 필요성 동기화.
• 일반화 및 정보 전달을 강화하기 위한 소프트 프라이어 엔드포인트를 갖춘 확산-브리지 제안.
• 역 단계의 정확성을 높이기 위한 자기 일관 재귀 샘플링 메커니즘 도입.

제안 방법

• 타깃 소스 컨벡스 조합이 있는 포워드 프로세스: x_t ~ N(mu_x0,t x0 + mu_y,t y, sigma_t^2 I) 및 엔드포인트를 끝으로서 증가하는 sigma_t^2를 갖는 연속적 증가.
• 엔드포인트는 hard 소스 샘플(y가 아닌 노이즈가 추가된 소스 이미지 x_T = y_epsilon)이다.
• 역방향 프로세스는 회복 네트워크 G_theta를 사용하여 자기 일관 타깃 추정 x0~*를 반복적으로 생성하고 x_{t-1}를 q(x_{t-1}|x_t,y,x0~*)로 샘플링한다.
• 자기 일관 재귀: x0~{r+1} = G_theta(x_t,t,y,x0~{r})가 x0~*로 수렴할 때까지; 그런 다음 x0~*를 사용하여 q에서 x_{t-1}를 샘플링한다.
• 적대적 판별기 D_theta는 역샘플의 현실성을 안내하고, 픽셀 단위 L1 손실은 회복된 x0~*를 실제 x0에 정렬한다.

실험 결과

연구 질문

• RQ1 소스-대상 영상 간의 소프트 프라이어 엔드포인트를 갖춘 확산-브리지가 교차 모달 변환을 개선할 수 있는가?
• RQ2 자기 일관 재귀 샘플링 방식이 다중 모달 변환에서 중간 샘플의 정확도와 현실감을 향상시키는가?
• RQ3 SelfRDB가 다중 대조 MRI 및 MRI-CT 변환에서 GAN 및 확산 기반 베이스라인에 비해 어떤 성능을 보이는가?
• RQ4 엔드포인트 노이즈 수준 및 포워드 분산 스케줄링이 보이지 않는 소스 이미지에 대한 일반화에 어떤 영향을 미치는가?

주요 결과

• SelfRDB는 다중 대조 MRI 및 MRI-CT 작업에서 경쟁 방법(GAN 및 확산 모델)보다 우수한 변환 성능을 달성한다.
• 점진적으로 증가하는 엔드포인트 노이즈 스케줄은 소스 모달리티에 대한 하드 프라이어를 완만하게 하여 일반화를 향상시킨다.
• 자기 일관 재귀 샘플링 절차는 역 단계 추정의 정확성과 이미지 합성을 향상시킨다.
• IXI, BRATS, 골반 MRI-CT 데이터셋에서 SelfRDB는 SynDiff, DDPM, I2SB, pix2pix 베이스라인보다 높은 PSNR 및 SSIM을 보인다.