A Novel Approach for Integrating Mamba and Diffusion in 3D Medical Imaging

기술 요약: 3D 의료 영상에서의 Mamba와 Diffusion 통합을 위한 새로운 접근 방식

문제 정의
본 논문은 3D 의료 볼륨 재구성의 핵심 과제, 특히 움직임 아티팩트, 제한된 시야(field-of-view), 또는 획득 제약으로 인해 데이터가 누락된 부분적 볼륨(partial volumes)을 완성하는 문제를 다룹니다. 저자들은 기존의 최첨단 방법들이 다음과 같은 구체적인 한계를 가지고 있음을 식별했습니다:

• CNN (예: 3D U-Net): 수용 영역(receptive field)이 작아 장거리 해부학적 의존성을 포착하는 데 한계가 있음.
• Transformer (예: Swin-UNETR): 전역적 문맥 모델링은 가능하지만, $O(N^2)$의 이차 복잡도로 인해 고해상도 3D 볼륨 처리에 있어 계산 비용이 매우 높음.
• Diffusion Models: 고주파 세부 사항을 생성하는 데는 효과적이지만, 확률적(stochastic) 특성으로 인해 해부학적 일관성이 결여된 출력을 생성할 수 있으며, 정확한 불확실성 추정 능력이 부족함.

방법론: MedMamba3D
저자들은 불확실성 추정에 의해 가이드되는 단일 프레임워크 내에 네 가지 상호 보완적인 기술을 통합하도록 설계된 하이브리드 아키텍처인 MedMamba3D를 제안합니다. 이 아키텍처는 다음으로 구성됩니다:

1. 3D 선택적 상태 공간 인코더 (Mamba):

   • 양방향 스캐닝 메커니즘을 갖춘 계층적 Mamba 인코더를 활용합니다.
   • 시스템 파라미터($A, B, C$)가 입력에 의존하는 선택적 상태 공간 메커니즘을 채택하여, 모델이 관련 정보는 기억하고 무관한 데이터는 잊도록 동적으로 조절합니다.
   • 선형 계산 복잡도 $O(N)$을 달성하여, Transformer의 $O(N^2)$ 복잡도와 비교했을 때 고해상도 3D 볼륨(예: $256^3$ 복셀)을 효율적으로 처리할 수 있습니다.
   • 후속 브랜치들의 조건(condition)으로 사용될 잠재 특징(latent features)을 추출합니다.

2. 결정론적 U-Net 브랜치 (해부학적 사전 정보):

   • Mamba 인코더로부터 스킵 연결(skip connections)을 사용하는 3D U-Net 구조입니다.
   • 전역적인 해부학적 일관성과 구조적 무결성을 유지하는 데 집중합니다.
   • 관측된 데이터 및 해부학적 사전 정보에 재구성이 부합하도록 $L_1$ 및 $L_2$ 손실 함수를 사용하여 학습됩니다.

3. 조건부 Diffusion 브랜치 (고주파 생성):

   • 부분적 볼륨과 Mamba 인코더의 잠재 특징을 조건으로 사용하는 생성 브랜치입니다.
   • 역 확산 과정(reverse diffusion process)을 수행하여 결정론적 브랜치가 놓칠 수 있는 고주파 세부 사항과 사실적인 질감을 생성합니다.
   • 타겟 볼륨에 추가된 노이즈($\epsilon$)를 예측하도록 학습됩니다.

4. 불확실성 인지 베이지안 퓨전 모듈 (Uncertainty-Aware Bayesian Fusion Module):

   • 알레아토리 불확실성(데이터 노이즈: 결정론적 경로의 복셀당 분산 및 확률적 확산 궤적 간의 분산을 통해 추정)과 에피스테믹 불확실성(모델 불확실성: Monte Carlo Dropout 및 랜덤 시드를 통해 추정)을 구분합니다.
   • **역 분산 가중치(inverse variance weighting)**를 사용하여 결정론적 브랜치($\hat{X}_d$)와 생성적 브랜치($\hat{X}_g$)의 예측을 결합합니다. 이는 이론적으로 최적인 퓨전 방식으로, 추정된 신뢰도에 따라 예측에 가중치를 부여하여 최종 볼륨($\hat{X}_{fused}$)과 그에 따른 불확실성 맵을 생성합니다.

주요 기여

• 아키텍처 통합: 최근의 하이브리드 상태 공간 및 확산 모델의 발전을 바탕으로, 3D 의료 볼륨 완성을 위해 Mamba(상태 공간 모델), Diffusion, 그리고 베이지안 불확실성 퓨전을 통합한 새로운 아키텍처를 제안합니다.
• 효율성: Mamba 기반 인코더가 Transformer의 확장성 문제를 극복하고 선형 복잡도로 3D 볼륨 데이터를 처리할 수 있음을 입증합니다.
• 신뢰성: 결정론적 예측과 확률적 예측의 퓨전을 가이드하기 위해 알레아토리 및 에피스테믹 불확실성을 모두 추정하고 활용하는 메커니즘을 도입하여, 순수 확산 모델의 "해부학적 불일치" 문제를 해결합니다.
• 강건성: 상당한 데이터 손실(최대 90% 누락)이 발생하는 상황에서도 견고하게 작동하도록 설계되었습니다.

실험 결과
모델은 BraTS 2021 (뇌 MRI), fastMRI, MIMIC-CXR의 세 가지 데이터셋에서 평가되었습니다.

• 정량적 성능: MedMamba3D는 모든 지표에서 최첨단(SOTA) 결과를 달성했습니다.
  • BraTS에서 28.34 dB의 PSNR을 기록하여, Swin-UNETR(24.62 dB)과 LDM(24.94 dB)을 능가했습니다.
  • Transformer 대비 3.7 dB, Diffusion 모델 대비 3.4 dB의 PSNR 향상을 보였습니다.
  • 기대 교정 오차(ECE) 0.08로 우수한 보정(calibration) 성능을 나타냈습니다.
  • 세그멘테이션 성능(Dice score)이 베이스라인 대비 8.9% 향상되었습니다.
• 누락 데이터에 대한 강건성: 90%의 데이터가 누락된 상황에서도 높은 성능을 유지하며, 24.6 dB의 PSNR을 달성하여 3D U-Net(17.6 dB) 및 Swin-UNETR(20.1 dB)을 크게 앞질렀습니다.
• 효율성: fastMRI 데이터셋에서 MedMamba3D는 추론에 0.52초가 소요되어, LDM(3.1초) 및 Swin-UNETR(1.38초)보다 빠르면서도 더 높은 PSNR과 SSIM을 제공했습니다.
• 어블레이션 연구 (Ablation Studies): Mamba 인코더, Diffusion 브랜치, 그리고 베이지안 퓨전의 전체 통합이 최상의 결과를 낸다는 것을 확인했으며, 특히 퓨전 모듈이 보정 성능을 개선하고 지각적 오류(LPIPS)를 줄이는 데 기여함을 입증했습니다.

의의 및 주장
본 논문은 MedMamba3D가 결정론적 모델이 제공하는 해부학적 일관성과 확산 모델이 제공하는 고주파 세부 사항 생성 사이의 트레이드오프를 성공적으로 해결하는 동시에, 3D 볼륨에 대한 계산적 실행 가능성을 확보했다고 주장합니다. 불확실성 추정을 포함함으로써, 모델은 임상의에게 신뢰도 맵을 제공하여 재구성 과정에서 불확실성이 높은 영역을 식별할 수 있게 합니다.

저자들은 본 방법론이 메모리 집약적이며, 쌍을 이룬 데이터셋(paired datasets)이 필요하고, 긴 학습 과정 및 비용이 많이 드는 확산 샘플링 과정을 수반한다는 점을 한계점으로 인정합니다. 그러나 이 접근 방식이 임상 분석, 세그멘테이션 및 치료 계획을 위한 신뢰할 수 있는 고충실도(high-fidelity) 3D 의료 볼륨을 생성하는 데 있어 중요한 진전을 이루었다고 단언합니다.