Mamba와 계층형 소프트 MoE의 결합으로 3D 의료 영상 분할 혁신

초록

본 논문은 Mamba 기반 선택적 상태공간 모델(SSM)에 계층형 소프트 Mixture‑of‑Experts(HoME) 레이어를 추가한 새로운 3D 의료 영상 분할 프레임워크를 제안한다. 토큰을 지역 그룹으로 나누어 로컬 전문가에게 라우팅하고, 이후 전역 전문가를 통해 그룹 간 정보를 교환함으로써 장거리 컨텍스트와 지역 특성을 동시에 포착한다. U‑shape 인코더‑디코더 구조에 통합된 Mamba‑HoME는 CT, MRI, 초음파 등 3가지 주요 모달리티와 다양한 데이터 품질에서 기존 최첨단 모델을 능가하는 정확도와 메모리 효율성을 보인다.

상세 분석

이 연구는 3차원 의료 영상에서 발생하는 두 가지 핵심 문제—(1) 고해상도 볼륨 데이터의 토큰 수가 급증해 전통적인 Transformer의 O(N²) 복잡도가 비현실적이며, (2) 다양한 해부학적 구조와 병변 패턴을 효과적으로 다루기 위해 지역 특화 전문가가 필요하다는 점—을 동시에 해결하고자 한다. 이를 위해 저자들은 먼저 Mamba SSM을 백본으로 채택한다. Mamba는 연속시간 재귀식 기반의 선형 복잡도 O(N·d) 연산을 통해 장거리 의존성을 효율적으로 학습한다는 장점이 있다. 그러나 순수 SSM은 모든 토큰을 동일한 파라미터 집합에 매핑하므로, 지역별 패턴 다양성을 반영하기 어렵다.

HoME 레이어는 이러한 한계를 보완한다. 입력 시퀀스를 Gᵢ개의 그룹으로 분할하고, 각 그룹 내에서 Kᵢ개의 토큰을 로컬 전문가(E₁ᵢ)에게 소프트 라우팅한다. 라우팅 가중치는 토큰과 학습 가능한 슬롯 임베딩 간의 내적을 softmax로 정규화한 값이며, 그룹 단위 라우팅을 통해 메모리 피크를 크게 낮춘다(전체 복잡도는 O(N·M·d) 유지). 첫 번째 레벨에서는 각 그룹이 독립적인 로컬 피처를 추출하고, 두 번째 레벨에서는 모든 그룹의 슬롯을 하나의 긴 시퀀스로 재구성한 뒤 전역 전문가(E₂ᵢ)에게 라우팅한다. 전역 전문가는 그룹 간 상호작용을 촉진하고, 최종 출력은 어텐션 기반 가중합을 통해 토큰 차원으로 복원된다.

Mamba‑HoME 블록은 Gated Spatial Convolution(GSC)으로 초기 공간적 선행 정보를 제공하고, 이를 1D 시퀀스로 펼친 뒤 Mamba와 HoME를 순차 적용한다. 정규화는 LayerNorm 대신 동적 tanh(DyT)를 사용해 연산 비용을 절감하고, 그래디언트 안정성을 확보한다. 인코더‑디코더 전체 구조는 SegMamba를 확장한 U‑shape 설계이며, 단계가 깊어질수록 첫 번째 레벨 전문가 수(Eᵢ)는 증가하고 그룹 크기(Kᵢ)는 감소하도록 설계돼 점진적인 세분화와 전문화를 구현한다.

실험에서는 PANORAMA, AMOS, FeTA 2022, MVSeg 등 네 개의 공개 데이터셋과 자체 CT 데이터셋을 사용해 CT, MRI, 초음파 3모달리티 전반에 걸쳐 평가하였다. Mamba‑HoME는 Dice 계수와 Hausdorff 거리 등 주요 지표에서 기존 Transformer‑ 기반 3D SegFormer, Swin‑UNet, SegMamba 등을 모두 앞섰으며, GPU 메모리 사용량과 추론 시간에서도 현저히 효율적이었다. 특히 데이터 품질이 낮은 경우(노이즈, 해상도 저하)에도 일반화 성능이 유지되는 점이 주목할 만하다.

한계점으로는 전문가 수와 그룹 크기 하이퍼파라미터 설정이 데이터 규모에 따라 민감할 수 있다는 점, 그리고 현재 구현이 3D 볼륨 전체를 한 번에 처리하므로 매우 큰 스캔(예: 전체 몸통 CT)에서는 추가적인 슬라이싱 전략이 필요할 수 있다는 점을 들 수 있다. 향후 연구에서는 자동화된 하이퍼파라미터 탐색, 멀티‑GPU 파이프라인, 그리고 멀티모달(이미지‑텍스트) 연계 학습으로 확장 가능성이 기대된다.