MHMamba: Multi-Head Mamba for 3D Brain Tumor Segmentation

기술 요약: 3D 뇌 종양 분할을 위한 MHMamba

문제 제기
3D MRI 에 기반한 뇌 종양 분할은 종양 형태의 높은 이질성과 다중 모드 대비도로 인해 상당한 도전에 직면해 있습니다. 합성곱 신경망 (CNN) 은 국소적인 질감과 경계 모델링에 뛰어나지만, 제한된 수용野로 인해 장기 의존성을 처리하는 데 어려움을 겪습니다. 반면, Transformer 기반 모델은 전역적 맥락을 포착할 수 있지만, 2 차 계산 및 메모리 복잡도 ($O(N^2)$) 로 인해 공격적인 다운샘플링이나 윈도우링 없이 고해상도 3D 볼륨에 적용하기 어렵습니다. 이는 맥락의 불일치와 경계 불일치를 초래할 수 있습니다. 또한, 기존 하이브리드 접근법은 합성곱 및 어텐션 브랜치 간의 통계적 불일치 (예: 배치/인스턴스 정규화 vs 레이어 정규화) 로 인해 불안정한 학습과 특징 융합 문제를 겪는 경우가 많습니다.

방법론
본 논문은 고해상도 세부 정보를 유지하면서 선형 복잡도로 전역 특징을 모델링하도록 설계된 U 자형 인코더 - 디코더 아키텍처인 **MHMamba(멀티헤드 맘바)**를 제안합니다. 이 프레임워크는 세 가지 주요 구성 요소로 이루어져 있습니다:

1. 멀티헤드 맘바 (MHM) 모듈:

   • 인코더의 표준 블록을 멀티헤드 상태 공간 모델 (SSM) 로 대체합니다.
   • 채널 차원을 병렬 SSM 헤드 ($N_h$) 로 분할하여, 각 헤드가 고정된 모드 할당 없이 다중 모드 입력으로부터 상호 보완적인 장기 의존성 패턴을 학습할 수 있도록 합니다.
   • 각 헤드는 3D 공간 특징을 시퀀스로 직렬화하고, 선택적 상태 업데이트 (Mamba) 를 적용한 후, 연결 및 프로젝션을 통해 출력을 융합합니다.
   • 이 설계는 선형 복잡도를 유지하면서도 다중 모드 조건 하에서 전역 모델링 능력과 학습 안정성을 향상시킵니다.

2. 채널 - 공간 보정 어텐션 (CSCA) 모듈:

   • 멀티헤드 맘바 처리 후 특징 통계를 정렬하고 병변 반응을 강화하도록 설계되었습니다.
   • 이중 경로 메커니즘을 사용합니다:
     • 채널 경로: 전역 평균 및 최대 풀링을 MLP 에 적용하여 채널 가중치를 생성합니다.
     • 공간 경로: 이질성과 뚜렷한 피크를 포착하는 공간 가중치 맵을 생성하기 위해 네 가지 공간 통계 (평균, 표준 편차, 최대값, 최소값) 를 혁신적으로 통합합니다.
   • 게이트 메커니즘이 이러한 브랜치를 적응적으로 융합한 후, 원래 특징 반응을 보존하기 위해 잔차 연결이 뒤따릅니다.

3. 적응형 게이트 융합 (AGF) 모듈:

   • 인코더와 디코더 간의 스킵 연결에 위치합니다.
   • 고수준 의미론과 저수준 세부 사항 간의 충돌을 해결하기 위해 다중 스케일 특징을 동적으로 융합합니다.
   • 입력 특징은 채널 차원을 따라 하위 그룹으로 분할됩니다. 각 하위 그룹에 대해 모듈은 인코더와 디코더 특징 간의 균형을 맞추기 위해 독립적인 융합 가중치를 학습하여, 종종 일관성 없는 경계 예측으로 이어지는 특징 충돌을 방지합니다.

주요 기여

• 멀티헤드 맘바 아키텍처: 본 논문은 채널을 병렬 헤드로 분해하는 3D MRI 를 위한 멀티헤드 SSM 접근법을 소개합니다. 이는 Mamba 의 선형 복잡도 장점을 유지하면서 장기 모델링 및 다중 모드 적응성을 향상시킵니다.
• 특징 보정 및 융합: 특징 스케일과 통계를 통합하는 CSCA 모듈과 스킵 연결에서 다중 스케일 특징을 동적으로 융합하는 AGF 모듈의 통합은 경계 일관성과 소량 병변 검출을 크게 개선합니다.
• 체계적 검증: 이 방법은 BraTS2021 및 BraTS2023 데이터셋에서 엄격하게 평가되었으며, Dice 점수와 Hausdorff 거리 (HD95) 측면에서 nnU-Net, SegMamba, Transformer 기반 모델 등 강력한 베이스라인보다 안정적인 개선을 보였습니다.

실험 결과
BraTS2021 및 BraTS2023 데이터셋에서 수행된 실험은 MHMamba 가 최첨단 또는 최첨단에 근접한 성능을 달성함을 보여줍니다:

• BraTS2021: 평균 Dice 점수 91.02% 및 평균 HD95 3.38을 달성하여 SegMamba(90.32% Dice, 4.18 HD95) 및 nnU-Net(89.40% Dice, 5.58 HD95) 과 같은 베이스라인을 능가했습니다.
• BraTS2023: 평균 Dice 점수 90.23% 및 평균 HD95 3.67을 달성했습니다.
• 어블레이션 연구:
  • CSCA, AGF, 멀티헤드 맘바 모듈을 개별적으로 추가한 것은 전체 종양 하위 영역 (전체 종양, 종양 코어, 강화 종양) 에서 성능 향상에 기여했습니다.
  • 최적 구성은 4 개의 헤드로 밝혀졌으며, 이는 표현 능력과 계산 효율성 사이의 균형을 이루었습니다. 헤드를 8 개로 늘리면 작은 배치 크기로 3D 학습 시 과도한 분할과 최적화 노이즈로 인해 성능이 저하되었습니다.
  • 정성적 결과는 MHMamba 가 특히 이질적인 종양 영역과 작은 강화 영역에서 더 단단한 경계를 가진 덜 조각난 예측을 생성함을 나타냈습니다.

의의 및 주장
본 논문은 MHMamba 가 3D 뇌 종양 분할에서 전역 맥락 모델링과 국소 세부 사항 보존 간의 균형을 효과적으로 해결한다고 주장합니다. 멀티헤드 아키텍처를 통해 상태 공간 모델의 선형 복잡도를 활용함으로써, 이 방법은 Transformer 의 2 차 메모리 오버헤드를 피하면서도 CNN 의 국소 수용野 제한을 극복합니다. 제안된 모듈 (CSCA 및 AGF) 은 다중 모드 3D 분할에서 흔히 발견되는 통계적 불일치와 맥락 불일치를 특히 대상으로 하여, 방사선 치료 계획 및 수술 전 평가와 같은 임상 응용에 적합한 견고하고 효율적이며 정확한 프레임워크를 제공합니다. 저자들은 이 방법이 Mamba 기반 모델링의 효율성 장점을 유지하면서 종양 코어와 소량 병변에 대한 민감도를 크게 향상시킨다고 강조합니다.