⚡ electrical engineering

Mamba Goes HoME: Hierarchical Soft Mixture-of-Experts for 3D Medical Image Segmentation

본 논문은 Mamba 백본을 기반으로 구축되어 긴 문맥 모델링을 강화하고 다양한 모달리티와 데이터 품질에 걸쳐 3D 의료 영상 분할에서 최첨단 성능을 달성하기 위해 2단계 토큰 라우팅 메커니즘을 활용하는 계층적 소프트 전문가 혼합(HoME) 프레임워크를 소개한다.

원저자: Szymon Płotka, Gizem Mert, Maciej Chrabaszcz, Ewa Szczurek, Arkadiusz Sitek

게시일 2026-02-09

📖 4 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Szymon Płotka, Gizem Mert, Maciej Chrabaszcz, Ewa Szczurek, Arkadiusz Sitek

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

당신은 수백만 권의 책이 담긴 거대한 도서관을 정리하려고 노력 중이라고 상상해 보세요 (이 책들은 인간의 몸을 찍은 3D 의료 스캔 데이터를 나타냅니다). 당신의 목표는 작은 종양이나 특정 장기의 경계와 같은 아주 미세한 세부 사항을 찾아내는 동시에, 전체 도서관이 어떻게 배치되어 있는지에 대한 큰 그림을 이해하는 것입니다.

이 논문은 이 문제를 해결하기 위해 Mamba-HoME라는 새로운 AI 시스템을 소개합니다. 다음은 이 시스템이 어떻게 작동하는지 쉬운 비유를 통해 설명한 것입니다.

문제점: 하나의 두뇌로는 너무 큼

기존의 AI 모델들은 도서관 전체를 한꺼번에 보려고 시도했습니다.

기존 방식 A (CNNs): 한 번에 하나의 책장만 바라보는 사서와 같습니다. 이들은 빠르고 국소적인 세부 사항을 보는 데 능숙하지만, 전체 방 안에서 책장들이 어떻게 연결되어 있는지는 볼 수 없습니다. 즉, "큰 그림"을 놓칩니다.
기존 방식 B (Transformers): 모든 연결 관계를 이해하기 위해 도서관에 있는 모든 책을 동시에 읽으려는 사서와 같습니다. 이는 큰 그림을 보는 데는 매우 훌륭하지만, 시간이 너무 오래 걸리고 이를 처리하기 위해 슈퍼컴퓨터 수준의 메모리가 필요합니다. 고해상도 3D 스캔을 처리하기에는 너무 느려집니다.

해결책: "Mamba-HoME" 팀

저자들은 방식 A의 속도와 방식 B의 거시적 시야를 결합하되, HoME(Hierarchical Soft Mixture-of-Experts)라는 영리한 트위스트를 더한 새로운 AI 팀을 만들었습니다.

HoME를 거대한 건설 현장의 2단계 관리 시스템이라고 생각해 보세요.

레벨 1: 지역 현장 소장 ( "로컬" 전문가)

팀이 거대한 3D 스캔을 받으면, 모든 픽셀을 본사 책임자에게 보내는 대신, 스캔을 작고 관리 가능한 이웃 단위(그룹)로 나눕니다.

각 구역에는 해당 구역의 특정 세부 사항을 처리할 지역 현장 소장(전문가)이 배정됩니다.
만약 어떤 구역에 뼈의 질감이 많다면 "뼈 전문가"가 담당하고, 액체가 있다면 "액체 전문가"가 업무를 맡습니다.
이는 효율적입니다. 왜냐하면 전문가들은 먼저 자기 구역 안에 있는 사람들하고만 소통하기 때문입니다. 그들은 도시 전체를 향해 소리를 지르며 시간을 낭비하지 않습니다.

레벨 2: 도시 계획가 ( "글로벌" 전문가)

지역 현장 소장들이 업무를 마치면, 그들은 요약본을 도시 계획가(글로벌 전문가)에게 보냅니다.

이 계획가들은 모든 구역에서 올라온 요약본들을 살펴봅니다.
이들은 각 구역이 어떻게 서로 맞물리는지 파악합니다. 예를 들어, "간(구역 A)이 위(구역 B) 바로 옆에 있다"는 식입니다.
이들은 최종 지도를 정교하게 다듬어, 국소적인 세부 사항이 전체적인 맥락 속에서 제대로 맞아떨어지는지 확인합니다.

왜 "Mamba"인가?

이 논문은 Mamba라는 백본 기술을 사용합니다. Mamba를 매우 효율적인 컨베이어 벨트라고 생각해 보세요.

기존의 시스템들은 전체 라인을 이해하기 위해 벨트 위의 모든 물건을 하나하나 멈춰 서서 확인해야 했지만(이는 느립니다), Mamba는 벨트를 따라 선형적으로 이동합니다. Mamba는 이전에 보았던 것을 기억하고 그 기억을 사용하여 현재 보고 있는 것을 이해하며, 이 과정에서 과부하가 걸리지 않습니다.
이를 통해 시스템은 이전 방식들보다 훨씬 빠르게, 그리고 훨씬 적은 컴퓨터 메모리를 사용하여 거대한 3D 스캔(예: 인간의 몸 전체)을 처리할 수 있습니다.

결과: 더 정확한 지도, 더 적은 메모리

저자들은 이 시스템을 세 가지 다른 유형의 의료용 "카메라"로 테스트했습니다:

CT 스캔 (뼈와 장기를 3D로 보여주는 X선과 유사).
MRI 스캔 (뇌나 간과 같은 연조직을 매우 상세하게 보여줌).
초음파 (보통 노이즈가 많고 입자가 거침).

연구 결과:

정확도: Mamba-HoME는 현재 존재하는 다른 어떤 방법보다 장기나 종양의 경계를 더 정확하게 그려냈습니다. 특히 사물의 "가장자리"를 찾는 데 뛰어났습니다.
효율성: 매우 똑똑함에도 불구하고, 3D 의료 영상을 처리할 때 큰 문제가 되는 컴퓨터 메모리(RAM)를 많이 필요로 하지 않았습니다.
다재다능함: 이 시스템은 CT와 MRI 스캔으로부터 학습하였으며, 장기의 "형태"를 이해하는 능력이 매우 뛰어나서, 초음파 데이터로 명시적인 학습을 하지 않았음에도 초음파 스캔에서도 잘 작동할 수 있었습니다.

요약

이 논문은 Mamba-HoME가 컴퓨터에게 인간의 몸 내부를 "보는" 법을 가르치는 새로운 방법이라고 주장합니다. 이는 두 단계의 전문가 팀을 사용합니다. 즉, 작은 세부 사항을 빠르게 처리하는 지역 전문가와, 그 세부 사항들이 완벽하게 조화를 이루도록 만드는 글로벌 전문가를 활용하는 것입니다. 이 접근 방식은 더 빠르고, 더 적은 컴퓨터 자원을 사용하며, 현재의 최상급 도구들보다 더 정확한 의료 지도를 만들어냅니다.

참고: 이 논문은 AI 모델의 세그멘테이션 작업(장기 주변의 선을 그리는 작업)에 대한 기술적 성능에만 전적으로 집중합니다. 이 모델이 현재 병원에서 환자 진단이나 치료 계획 수립에 사용되고 있다고 주장하거나, 제시된 기술적 벤치마크를 넘어선 미래의 임상적 적용에 대해 논하지 않습니다.

기술 요약: Mamba Goes HoME

문제 정의

3차원(3D) 의료 영상 분할은 CT, MRI, 초음파(US)와 같은 다양한 모달리티에서 진단 및 치료 계획 수립에 있어 매우 중요합니다. 그러나 기존 아키텍처들은 다음과 같은 상당한 트레이드오프(trade-offs)에 직면해 있습니다:

**합성곱 신경망(CNNs)**은 선형적인 계산 복잡성을 제공하지만, 수용 영역(receptive field)이 제한되어 국소적 병변과 전역적 해부학적 구조 사이의 계층적 관계를 이해하는 데 필수적인 전역 공간 패턴을 포착하는 데 어려움이 있습니다.
**비전 트랜스포머(ViTs)**는 전역 어텐션(global attention)을 활용하여 장거리 의존성을 모델링하지만, 토큰 수에 대해 이차(quadratic)적인 계산 복량을 가집니다. 이는 고해상도 3D 볼륨 데이터에 대해 계산적으로 매우 부담스럽습니다.
**선택적 상태 공간 모델(Selective State-Space Models, SSMs)**인 Mamba와 같은 모델은 긴 시퀀스 모델링에 대해 선형 복잡도를 제공하여 ViT의 효율적인 대안이 됩니다. 그러나 표준 SSM은 복잡하고 다중 스케일인 의료 데이터에서 발견되는 다양한 국소 패턴을 적응적으로 처리하도록 설계되지 않았습니다.
혼합 전문가(Mixture-of-Experts, MoE) 프레임워크는 국소 패턴을 처리하기 위해 특징을 특화된 서브 네트워크로 동적으로 라우팅할 수 있지만, SSM의 전역적 효율성과 MoE의 국소적 적응성을 3D 의료 영상에서 결합하는 연구는 여전히 미개척 분야로 남아 있습니다.

방법론: Mamba-HoME

저자들은 Mamba 선택적 상태 공간 모델(SSM) 백본과 계층적 소프트 혼합 전문가(Hierarchical Soft Mixture-of-Experts, HoME) 레이어를 통합한 새로운 아키텍처인 Mamba-HoME를 제안합니다. 이 설계의 목적은 선형 계산 복잡도를 유지하면서 국소-전역 패턴 계층을 효율적으로 모델링하는 것입니다.

1. 계층적 소프트 혼합 전문가 (HoME)

핵심 혁신은 국소 전문화와 전역 컨텍스트 통합의 균형을 맞추기 위해 설계된 2단계 토큰 라우팅 레이어입니다:

레벨 1: 국소 그룹화 및 라우팅. 입력 시퀀스는 국소 그룹으로 분할됩니다. 소프트 혼합 전문가(SMoE) 레이어는 이러한 그룹 내의 토큰을 특화된 국소 전문가에게 라우팅합니다. 이를 통해 전역 라우팅에 비해 메모리 오버헤드를 줄이면서 공간적 국소성을 보존하고 국소적 특징 추출을 가능하게 합니다.
레벨 2: 전역 집계. 첫 번째 레벨의 출력은 집계되어 두 번째 레벨의 전역 SMoE 레이어를 통과합니다. 이는 그룹 간 정보 융합과 전역 컨텍스트 정교화를 촉진합니다.
확장성: 그룹 내에서 집계 전 라우팅을 수행함으로써, 토큰 수가 수백만 개에 달할 수 있는 고해상도 3D 데이터에서의 과도한 메모리 비용을 피하고 $O(Nd)$의 계산 복잡도를 유지합니다.

2. Mamba-HoME 블록

이 아키텍처는 HoME 레이어를 U자형 인코더-디코더 네트워크(SegMamba에서 영감을 받음)에 내장합니다:

백본: 선형 시간 복잡도로 효율적인 장거리 의존성 모델링을 위해 Mamba 레이어를 활용합니다.
게이트형 공간 합성곱(Gated Spatial Convolution, GSC): 데이터를 Mamba 처리를 위한 1D 시퀀스로 평탄화(flattening)하기 전에 국소 공간 사전 지식(spatial priors)을 포착하기 위해 통합되었습니다.
다이내믹 탄젠트(Dynamic Tanh, DyT): 표준 레이어 정규화(Layer Normalization)를 포인트 와이즈 비선형 함수( $f(x) = w \cdot \tanh(\alpha \cdot x) + b$ )로 대체합니다. 이는 평균과 분산 계산을 피함으로써 그래디언트를 안정화하고 계산 오버헤드를 줄여 훈련 및 추론을 가속화합니다.
계층적 진행: 전문가의 수( $E_1 < E_2 < \dots$ )는 깊이에 따라 단조 증가하는 반면, 그룹 크기( $K_1 > K_2 > \dots$ )는 감소하여, 네트워크가 깊어질수록 점진적으로 더 세밀한 처리가 가능하도록 합니다.

주요 기여

HoME 레이어: 토큰을 그룹화하여 국소 전문가 라우팅을 수행한 후 전역 전문가 정교화를 거치는 2단계 토크 라우팅 MoE 레이어를 도입하여, 국소-전역 패턴 계층을 효율적으로 포착합니다.
통합 아키텍처: Mamba의 SSM을 HoME와 결합하여 메모리 효율적인 장기 시퀀스 처리와 계층적 전문가 라우팅을 동시에 구현한 Mamba-HoME 블록을 설계했습니다.
다단계 분할 네트워크: 이러한 구성 요소들을 3차원 의료 영상 분할에 특화된 U자형 아키텍처에 내장하여, 공간 및 깊이 차원에 걸친 컨텍스트 표현력을 향상시켰습니다.

실험 결과

저자들은 네 개의 공개 데이터셋(PANORAMA, AMOS, FeTA 2022, MVSeg)과 하나의 자체 제작 CT 데이터셋을 포함하여 CT, MRI, 초음파(US) 모달리티를 평가했습니다.

성능: Mamba-HoME는 테스트된 모든 모달리티에서 분할 정확도(Dice Similarity Coefficient - DSC) 및 경계 정밀도(Hausdorff Distance 95% - HD95) 측면에서 최첨단(SOTA) 결과를 달성했습니다.
- PANORAMA 데이터셋에서 사전 학습 모델 기준 56.7%, 처음부터 학습(from scratch) 기준 54.8%의 DSC를 달성하여 SegMamba 및 Swin UNETR와 같은 베이스라인을 능가했습니다.
- FeTA 2022 태아 MRI 데이터셋에서 87.7%의 DSC에 도달했습니다.
- MVSeg 3D 초음파 데이터셋에서 85.0%의 DSC를 달성했습니다(사전 학습 모델).
효율성: 일부 경쟁 모델에 비해 더 많은 파라미터 수(170.1M)를 가짐에도 불구하고, 고해상도 3D 데이터에 중요한 이점인 낮은 GPU 메모리 사용량을 보여주었습니다. 추론 속도는 일부 베이스라인보다 약 30% 느렸으나, 저자들은 이것이 높은 정확도 향상을 위한 매력적인 트레이드오프라고 언급했습니다.
일반화 능력: 모델은 강력한 교차 모달 일반화 성능을 보여주었습니다. CT 및 MRI 데이터로 사전 학습된 모델은 최소한의 미세 조정(fine-tuning)만으로도 3D 초음파 작업에 성공적으로 적응했으며, DSC 및 HD95 지표 모두에서 SOTA 방법들을 능가했습니다.
절제 연구(Ablation Studies):
- 특정 전문가 수( $E_1 \in \{4, 8, 12, 16\}$ , $E_2 \in \{8, 16, 24, 32\}$ )와 그룹 크기( $K \in \{2048, 1024, 512, 256\}$ )에서 최적의 성능이 발견되었습니다.
- Dynamic Tanh(DyT) 정규화는 정확도를 희생하지 않으면서 레이어 정규화 대비 훈련 및 추론 속도를 약 6% 개선했습니다.
- 대규모 CT 및 MRI 데이터셋에 대한 지도 학습 기반 사전 학습은 다양한 해부학적 영역 및 모달리티 전반에 걸쳐 성능과 견고함을 크게 향 old 향상시켰습니다.

의의 및 주장

본 논문은 Mamba-HoME가 SSM의 효율성과 MoE 프레임워크의 적응성 사이의 간극을 효과적으로 메움으로써 3D 의료 영상 분할 분야에서 중요한 진전을 이루었다고 주장합니다. 저자들은 계층적 설계가 이전의 CNN 또는 트랜스포머 기반 접근 방식보다 의료 데이터의 고유한 국소-전역 공간 계층을 더 효과적으로 포착할 수 있게 한다고 강조합니다.

이 연구는 모달리티 불가지론적(modality-agnostic) 특징 표현의 잠재력을 부각하며, CT 및 MRI로 사전 학습된 모델이 초음파 작업에 효과적으로 일반화될 수 있음을 입증하여 데이터 변동성 문제를 해결했습니다. 저자들은 M가 Mamba-HoME를 높은 정확도를 제공하면서도 경쟁력 있는 메모리 및 계산 효율성을 유지하는 고해상도 3D 데이터를 위한 확장 가능한 솔루션으로 제시합니다.

인정된 한계점: 저자들은 매우 큰 규모의 데이터셋(예: >10,000 스캔)에 대한 확장성과 대규모 자기 지도 학습(>200,000 스캔) 하에서의 동작이 아직 탐구되지 않았음을 언급하며, 이를 향후 연구 과제로 남겨두었습니다.