DAMamba-UNet3D: A Parameter-Efficient Mamba State Space U-Net with Dynamic Adaptive Scan for 3D Medical Image Segmentation

기술 요약: DAMamba-UNet3D

문제 정의
3D 의료 영상 분할은 엄격한 계산 및 파라미터 예산 내에서 정밀한 국소 경계 묘사와 일관된 전역 공간 문맥 사이의 균형을 맞춰야 합니다. 기존 아키텍처들은 다음과 같은 뚜렷한 트레이드오프를 가집니다:

• 합성곱 U-Net (Convolutional U-Nets): 파라미터 효율성과 $O(n)$ 국소 혼합(local mixing)을 제공하지만, 명시적인 전역 문맥이 부족하여 장거리 의존성을 포착하기 위해 깊이를 확장해야 합니다.
• 비전 트랜스포머 (Vision Transformers): 셀프 어텐션을 통해 전역적 추론을 제공하지만, 대규모 3D 볼륨(예: $128^3$ 패치는 약 210만 개의 토큰을 포함)에 대해 과도한 $O(n^2)$ 메모리 및 연산 비용을 발생시킵니다.
• 상태 공간 모델 (SSMs/Mamba): 정보를 $O(n)$ 시간 및 메모리로 전역적으로 전파합니다. 그러나 기존의 의료용 SSM 분할 모델(예: SegMamba)은 고정된 스캔 패턴과 큰 파라미터 예산($\sim 70$M 파라미터)에 의존하며, 데이터 의존적 적응성이 부족합니다.

Dynamic Adaptive Scan (DAS)이 2D 자연 이미지 인식에서 선택적 스캔 전 데이터 의전적 재정렬을 학습하기 위해 도입되었으나, 의료 영상에 적용되거나 3D 볼륨으로 확장된 적은 없습니다.

방법론
저자들은 Tri-Plane 3D-DAS 블록을 U-Net 프레임워크에 통합한 하이브리드 인코더-디코더 아키텍처인 DAMamba-UNet3D를 제안합니다.

1. 아키텍처 설계:

   • 하이브리드 구조: 국소 세부 사항을 보존하기 위해 스템(E1), 병목(bottleneck), 그리고 전체 디코더에는 합성곱 블록을 유지합니다.
   • 인코더 통합: Tri-plane DAMamba 블록을 인코더 단계인 E2–E4에만 독점적으로 삽입합니다.
   • 경량 변체 (Compact Variant): 기본 모델은 베이스 너비를 $C=16$으로 사용하여, 약 530만 개의 파라미터($\sim 13\times$ 더 적은 SegMamba 대비)를 사용합니다.
   • 대규모 변체 (Large-Scale Variant): 직접적인 비교를 위해 SegMamba의 용량에 맞추도록 설계된, 인코더에 Stacked DAMamba 블록과 합성곱 병목을 갖춘 더 넓은 변체($\sim 70$M 파라미터)인 DAMamba-L입니다.

2. Tri-Plane Dynamic Adaptive Scan (DAS):

   • 메커니즘: 고정된 스캔(행 우선, 지그재그 또는 다중 방향 평탄화)과 달리, DAS는 입력 의존적인 공간 재정렬을 학습합니다. 오프셋 예측 네트워크(OPN)가 특징 맵으로부터 2D 오프셋을 예측하며, 이 오프셋은 선택적 스캔 전 쌍선형 보간(bilinear interpolation)을 통해 특징을 리샘플링하는 데 사용됩니다.
   • 3D 확장: 2D DAS 연산자를 세 가지 해부학적 평면인 축상면(Axial, $H \times W$), 관상면(Coronal, $D \times W$), 시상면(Sagittal, $D \times H$)에 적용하여 볼륨으로 확장합니다. 출력값들을 평균하여 3D DAS 블록을 형성합니다:
     $$DAS3D(X) = \frac{1}{3}(DAS_{ax} + DAS_{cor} + DAS_{sag})$$
   • 선택적 스캔 (Selective Scan): 재정렬된 특징 맵은 표준 Mamba 선택적 스캔 블록에 의해 처리되며, 이를 통해 이차 함수적(quadratic) 어텐션 맵 없이 $O(N)$ 비용으로 전역 문맥 전파를 가능하게 합니다.

3. 학습 및 평가:

   • 데이터셋: BraTS 2020 훈련 세트 (369명의 피험자, 4가지 모달리티).
   • 손실 함수: 어려운 픽셀과 경계에 집중하기 위한 Modulated Combo Loss ($L = (L_{CE} + \delta) \cdot (L_{mod}^{Dice} + \delta)$)를 사용합니다.
   • 프로토콜: 5-겹 피험자 수준 교차 검증과 전체 볼륨 추론(슬라이딩 윈도우, $128^3$ 패치)을 수행합니다. 평가지표는 향상된 종양(ET), 종양 핵(TC), 전체 종양(WT)에 대한 평균 Dice 및 HD95입니다.

주요 기여

1. DAS의 첫 의료 분야 적용: Dynamic Adaptive Scan을 2D 자연 이미지 분류에서 밀집된 3D 의료 영상 분할로 전이했습니다.
2. 2D-to-3D 확장: DAS 연산자를 세 가지 평면 적응형 스캐닝을 통해 볼륨 데이터로 격상시켜, 수작업으로 설계된 다중 방향 평탄화의 학습 가능한 대안을 제공합니다.
3. 파라미터 효율성: Tri-plane DAS 블록을 인코더에만 배치하는 하이브리드 U-Net을 통해, 훨씬 더 큰 모델들과 대등한 성능을 보이면서도 $\sim 13\times$ 더 적은 파라미터 비용을 달성함을 입증했습니다.
4. 두 가지 규모의 검증: 컴팩트 규모($\sim 5.3$M)와 대규모($\sim 70$M) 모두에서 접근 방식을 검증하여, 아키텍처 배치가 최적화되었을 때 학습된 DAS가 고정된 스캔 패턴보다 우수함을 보여주었습니다.

결과

• 경량 모델 (DAMamba-UNet3D, $\sim 5.3$M): BraTS 2020에서 0.815 $\pm$ 0.013의 평균 Dice를 달キ했습니다. 이는 재학습된 SegMamba(0.824 $\pm$ 0.014)와 거의 대등한 수준이며, 파라미터 수는 $\sim 13\times$ 더 적습니다.
• 대규모 모델 (DAMamba-L, $\sim 70$M): 재학습된 SegMamba를 능가하여 0.829 $\pm$ 0.012의 평균 Dice를 기록했으며(0.5 포인트 개선), 가장 큰 이득은 종양 핵(TC) 클래스(0.827 vs 0.842)에서 관찰되었습니다.
• 절제 연구 (Ablation Studies):
  • 배치 (Placement): DAS 블록을 인코더(E2–E4)에만 배치했을 때 최적의 결과(0.815)를 얻었습니다. DAS를 병목이나 디코더에 추가하면 성능이 낮아졌으며(각각 0.807 및 0.805), 이는 인코더 전용 배치가 중요하다는 것을 나타냅니다.
  • 스캔 패턴: 인코더 배치가 일치할 때, 학습된 tri-plane DAS가 고정된 스캔 순서(Raster, Snake, Cross, Tri-orientated)보다 우수한 성능을 보였습니다.

의의 및 주장
본 논문은 학습된 tri-plane DAS를 포함한 하이브리드 U-Net이 3D 뇌종양 분할을 위한 고정된 Tri-orientated Mamba (ToM) 스캐닝의 경쟁력 있는 대안이라고 주장합니다.

• 효율성: 경량 모델은 무차별적인 용량 증대 대신 적응형 스캐닝을 활용함으로써 높은 성능을 달성할 수 있음을 보여줍니다.
• 성능: 대규모 모델에서 학습된 적응형 메커니즘(DAS)은 동일한 훈련 및 평가 프로토콜 하에서 고정된 스캐닝 궤적(ToM)보다 성능을 향상시킬 수 있습니다.
• 신중함: 저자들은 BraTS 2020에서의 결과가 유망하지만, 본 연구가 단일 데이터셋에 국한되어 있음을 주의를 기울여 언급했습니다. 그들은 DAS와 ToM에 대한 일반적인 주장을 하기 전에 "다중 데이터셋 검증 및 통계적 테스트가 필요하다"고 명시했습니다. 주요 기여는 DAS를 3D 의료 볼륨으로 확장하고, 이를 파라미터 효율적인 하이브리드 아키텍처에서 실현 가능성을 입증한 데 있습니다.