[논문 리뷰] VM-DDPM: Vision Mamba Diffusion for Medical Image Synthesis

Zhihan Ju, Wanting Zhou|arXiv (Cornell University)|2024. 05. 09.

Image Retrieval and Classification Techniques인용 수 6

한 줄 요약

VM-DDPM은 상태공간 모델(SSM)과 CNN을 결합한 Vision Mamba 확산 모델을 도입하여 효율적이고 전역적으로 인식 가능한 의료 영상 합성을 구현하고, 여러 데이터셋에서 최첨단 FID를 달성한다.

ABSTRACT

In the realm of smart healthcare, researchers enhance the scale and diversity of medical datasets through medical image synthesis. However, existing methods are limited by CNN local perception and Transformer quadratic complexity, making it difficult to balance structural texture consistency. To this end, we propose the Vision Mamba DDPM (VM-DDPM) based on State Space Model (SSM), fully combining CNN local perception and SSM global modeling capabilities, while maintaining linear computational complexity. Specifically, we designed a multi-level feature extraction module called Multi-level State Space Block (MSSBlock), and a basic unit of encoder-decoder structure called State Space Layer (SSLayer) for medical pathological images. Besides, we designed a simple, Plug-and-Play, zero-parameter Sequence Regeneration strategy for the Cross-Scan Module (CSM), which enabled the S6 module to fully perceive the spatial features of the 2D image and stimulate the generalization potential of the model. To our best knowledge, this is the first medical image synthesis model based on the SSM-CNN hybrid architecture. Our experimental evaluation on three datasets of different scales, i.e., ACDC, BraTS2018, and ChestXRay, as well as qualitative evaluation by radiologists, demonstrate that VM-DDPM achieves state-of-the-art performance.

연구 동기 및 목표

• 의료 영상의 데이터 부족 문제를 해결하기 위해 고품질의 합성 이미지를 생성한다.
• CNN과 상태공간 모델(SSM)을 결합하여 선형 계산 복잡도로 전역 맥락 모델링을 달성한다.
• 의료 영상용 다중 수준 특징 융합(MSSBlock) 및 인코더-디코더 SSLayer 유닛 설계.
• 향상된 Cross-Scan Module(CSM)과 매개변수 없는 Sequence Regeneration 전략을 통해 공간적 연속성과 질감을 개선한다.

제안 방법

• CNN-SSM 하이브리드 백본 위에 구축된 Denoising Diffusion Probabilistic Model(DDPM)인 VM-DDPM을 제안한다.
• CSM과 CNN 경로를 결합하는 다중 수준 특징 추출 유닛으로 MSSBlock를 도입한다.
• 잔차 연결과 시간 임베딩 처리를 갖춘 인코더/디코더 기본 유닛으로 SSLayer를 구현한다.
• S6 작동 이전에 패치 순서를 섞는 Plug-and-Play Sequence Regeneration 전략으로 Cross-Scan Module(CSM)을 향상시킨다.
• 스케일 간 특징 융합을 위해 U-Net과 유사한 스킵 연결을 갖춘 인코더-병목-디코더 아키텍처를 사용한다.

실험 결과

연구 질문

• RQ1 CNN-SSM 하이브리드 확산 모델이 무조건적 의료 영상 합성에서 경쟁력 있거나 우수한 품질과 다양성을 달성할 수 있는가?
• RQ2 Sequence Regeneration 전략이 SSM 기반 확산 모델의 공간적 연속성과 일반화를 향상시키는가?
• RQ3 ACDC, BraTS2018, ChestXRay와 같이 크기와 모달리티가 서로 다른 데이터셋에서 VM-DDPM의 성능은 GAN 및 DDPM 기준선과 비교해 어떠한가?
• RQ4 합성 의료 영상의 질감과 구조에 대한 MSSBlock 기반 다중 수준 특징 융합의 영향은 무엇인가?

주요 결과

• VM-DDPM은 세 데이터셋에서 GAN 기반 및 여러 DDPM 근사값보다 우수한 FID 점수를 달성한다.
• ChestXRay, BraTS2018, 및 ACDC에서 VM-DDPM은 각각 FID 점수 11.783, 12.513, 34.525를 달성한다(낮을수록 더 좋다).
• 변수 제거(ablation) 분석은 Sequence Regeneration 전략이 원래 CSM 대비 데이터셋 전반에서 성능을 향상시킴을 보여준다.
• 정성적 방사선 의사 평가에 따르면 합성 영상은 실제 영상과 구분되기 어렵고 병리 및 질감이 유사하다.
• 이 방식은 의료 영상 합성에서 데이터셋 간 일반화 및 확장성 면에서 강한 성능을 보인다.