MoEMambaMIL: Structure-Aware Selective State Space Modeling for Whole-Slide Image Analysis

본 논문은 와이드 슬라이드 이미지 (WSI) 분석의 거대 해상도 및 계층적 구조를 효과적으로 처리하기 위해, 영역 중첩 선택적 스캐닝과 혼합 전문가 (MoE) 모델링을 통합한 구조 인식 선택적 상태 공간 모델 프레임워크인 MoEMambaMIL 을 제안하여 9 가지 하위 작업에서 최상의 성능을 달성함을 보여줍니다.

핵심

MoEMambaMIL: 거대한 병리 슬라이드를 분석하는 '지능형 팀워크'

이 논문은 **전체 슬라이드 이미지 **(Whole-Slide Image, WSI)라는 아주 거대한 디지털 병리 사진을 분석하는 새로운 인공지능 모델을 소개합니다. 이 모델의 이름은 MoEMambaMIL입니다.

이 복잡한 기술을 일상적인 언어와 비유로 쉽게 설명해 드리겠습니다.

---

1. 문제 상황: "수백만 개의 퍼즐 조각을 어떻게 처리할까?"

병리 의사는 현미경으로 조직을 보지만, 컴퓨터는 이를 **수십억 개의 픽셀 **(거의 10 억 개 이상)로 된 거대한 이미지로 봅니다. 이 이미지를 분석할 때, 우리는 이미지를 잘게 쪼개어 (패치) 하나씩 살펴봐야 합니다.

• 기존 방법의 한계:

  • 기존 AI 들은 이 수많은 조각들을 무질서하게 뒤섞인 상자처럼 다뤘습니다. "이 조각이 어디에 있었는지"나 "어떤 조직의 일부인지"를 무시하고, 그냥 무작위로 섞어서 분석했습니다.
  • 또 다른 방법 (Transformer) 은 모든 조각을 서로 비교하게 하는데, 조각이 너무 많으면 계산량이 기하급수적으로 늘어나서 컴퓨터가 지쳐버립니다.

• 핵심 질문: "어떻게 하면 이 거대한 퍼즐 조각들을 논리적인 순서로 정리하고, 효율적으로 분석할 수 있을까?"

2. 해결책: MoEMambaMIL 의 두 가지 핵심 전략

이 모델은 두 가지 똑똑한 전략을 합쳐서 문제를 해결합니다.

전략 1: "마트형 진열" (Region-Nested Selective Scan)

비유: 대형 마트의 진열 방식을 생각해보세요.

• 기존 방식: 모든 상품 (조각) 을 바닥에 무작위로 뿌려두고 찾는 것.
• MoEMambaMIL 방식:
  1. 먼저 **대분류 **(과일, 채소, 육류)를 먼저 봅니다.
  2. 그다음 **중분류 **(사과, 배, 바나나)를 봅니다.
  3. 마지막으로 **세부 품목 **(빨간 사과, 초록 사과)을 봅니다.

이 모델은 이미지를 **거친 조직 **(대분류)에서 시작해 점점 **세부 세포 **(세부 품목)로 내려가는 **중첩된 순서 **(Region-Nested)로 정리합니다. 이렇게 하면 AI 가 "이 세포는 이 조직의 일부야"라는 위치와 구조적인 관계를 자연스럽게 이해하게 됩니다.

전략 2: "전문가 팀워크" (Mixture-of-Experts, MoE)

비유: 거대한 병원을 운영하는 전문가 팀을 상상해보세요.

• **정적 전문가 **(Static Experts): "해부학자"와 "세포학자"처럼 역할이 고정된 팀입니다.

  • 저해상도 (거친 이미지) 를 보는 팀은 전체적인 조직 구조만 봅니다.
  • 고해상도 (선명한 이미지) 를 보는 팀은 세포의 미세한 변화만 봅니다.
  • 각 팀은 자신의 역할에 맞춰 훈련되어 있어, 서로의 일을 방해하지 않고 효율적으로 일합니다.

• **동적 전문가 **(Dynamic Experts): "진료 의뢰를 받는 의사"처럼 상황에 따라 변하는 팀입니다.

  • 어떤 조직이 "암"일지 "염증"일지 **내용 **(Content)을 보고, 가장 적합한 전문가를 동적으로 선택합니다.
  • 예를 들어, 이상한 세포가 보이면 "세포 전문가"를, 조직 구조가 비정상적이면 "구조 전문가"를 부르는 식입니다.

이 모델은 **고정된 역할 **(해부학적 구조)과 **유연한 판단 **(질병 패턴)을 동시에 수행할 수 있도록 설계되었습니다.

3. 왜 이 방법이 더 좋은가요? (Mamba 의 역할)

이 모델은 Mamba라는 최신 AI 기술을 사용합니다.

• **기존 AI **(Transformer)는 모든 조각을 서로 비교하느라 전체 팀이 함께 회의를 하느라 시간이 오래 걸립니다.
• Mamba는 한 줄로 이어진 컨베이어 벨트처럼 작동합니다. 정보를 순서대로 빠르게 흘려보내면서 선별적으로 중요한 정보만 기억합니다.
• 결과적으로 계산 속도는 빨라지고, 메모리 사용량은 줄었지만 정확도는 높아졌습니다.

4. 실제 성과: "모든 테스트에서 1 등"

이 모델은 신장암, 간암, 유방암 등 다양한 병리 데이터로 실험했습니다.

• 결과: 기존에 가장 잘하던 방법들보다 더 높은 정확도를 보였습니다.
• 특징: 어떤 종류의 이미지 분석 도구 (Feature Extractor) 를 쓰든, 어떤 종류의 암을 진단하든 일관되게 좋은 성능을 냈습니다.

5. 요약: 한 줄로 정리하면?

MoEMambaMIL은 거대한 병리 이미지를 **논리적인 순서 **(거친 것→세부 것)로 정리하고, 역할이 고정된 전문가와 상황에 따라 변하는 전문가가 팀을 이루어 협력하게 만든 초고속, 초정밀 AI입니다.

이 기술은 앞으로 병리학자들이 암을 더 빠르고 정확하게 진단하는 데 큰 도움을 줄 것으로 기대됩니다.

기술 요약

1. 연구 배경 및 문제 정의 (Problem)

• 전체 슬라이드 이미지 (WSI) 분석의 난제: WSI 는 기가픽셀 (gigapixel) 규모의 해상도를 가지며, 조직의 거시적 구조와 세포 수준의 미세한 패턴이 계층적으로 중첩된 다중 해상도 (multi-resolution) 구조를 가집니다.
• 기존 방법의 한계:
  • 기존의 다중 인스턴스 학습 (MIL) 기반 방법들은 WSI 를 순서가 없는 패치 (patch) 집합으로 간주하는 경우가 많아, 조직의 공간적 계층성과 지역적 의존성을 포착하는 데 한계가 있습니다.
  • 비전 트랜스포머 (Vision Transformer) 기반 모델은 패치 간 의존성을 모델링하지만, 이차 복잡도 (quadratic complexity) 로 인해 긴 시퀀스 처리 시 확장성이 떨어집니다.
  • 최근 등장한 상태 공간 모델 (SSM, 예: Mamba) 은 선형 시간 (linear-time) 에 긴 시퀀스를 처리할 수 있지만, WSI 의 2 차원적이고 다중 해상도인 계층적 구조를 1 차원 시퀀스로 단순화 (flatten) 할 때 공간적 국소성과 생물학적 계층 관계가 파괴되는 문제가 있습니다.

2. 제안 방법: MoEMambaMIL (Methodology)

저자들은 WSI 의 구조적 우선순위 (structural priors) 를 명시적으로 인코딩하는 MoEMambaMIL 프레임워크를 제안합니다. 이는 영역 중첩 선택적 스캔 (Region-Nested Selective Scan) 과 혼합 전문가 (Mixture-of-Experts, MoE) 모델링을 결합한 구조 인식 SSM 프레임워크입니다.

핵심 구성 요소:

1. 영역 중첩 선택적 스캔 (Region-Nested Selective Scan):

   • WSI 의 다중 해상도 피라미드 구조를 활용하여, 거시적 (coarse) 영역을 먼저 정의하고 해당 영역에 포함된 미세한 (fine-grained) 하위 영역들을 재귀적으로 확장하는 방식으로 1 차원 토큰 시퀀스를 구성합니다.
   • 이 방식은 동일한 해부학적 영역에 속하는 패치들이 시퀀상에서 연속적으로 배치되도록 하여, SSM 이 공간적 포함 관계와 계층적 의존성을 자연스럽게 학습할 수 있게 합니다.

2. 이중 전문가 아키텍처 (Dual Expert Design):

   • 정적 전문가 (Static Experts, Resolution-Aware Encoding):
     • 패치의 해상도 (magnification) 에 따라 결정적으로 할당되는 전문가들입니다.
     • 고해상도 패치는 세포 수준의 세부 사항을, 저해상도 패치는 조직의 전체적 구조를 학습하도록 하여, 각 해상도별 고유한 특징을 인코딩합니다.
     • 학습된 라우팅이 아닌 하드 할당 (hard assignment) 을 사용하여 초기 인코딩 단계의 복잡성을 줄이고 안정성을 확보합니다.
   • 동적 전문가 (Dynamic Experts, Region-Adaptive Modeling):
     • 스캔된 시퀀스 위에서 작동하는 희소 혼합 전문가 (SparseMoE) 모듈입니다.
     • 학습된 게이트 (gating) 네트워크를 통해 토큰의 내용 (content) 에 따라 상위 k 개의 전문가를 동적으로 선택합니다.
     • 이는 공간적 영역마다 다른 진단 패턴 (이질적인 조직 특성) 을 적응적으로 모델링할 수 있게 합니다.

3. 부하 균형 정규화 (Load Balancing Regularization):

   • 전문가 붕괴 (expert collapse, 일부 전문가만 과도하게 사용되는 현상) 를 방지하기 위해 중요도 (importance) 와 부하 (load) 간의 균형을 유지하는 보조 손실 함수를 도입합니다.

4. 최종 예측:

   • MoEMamba 백본을 통과한 토큰 특징들은 어텐션 기반 풀링 (attention-based pooling) 을 통해 슬라이드 수준의 표현으로 집계되어 최종 진단 분류를 수행합니다.

3. 주요 기여 (Key Contributions)

1. 구조 인식 시리얼라이제이션: WSI 의 다중 해상도 계층 구조를 상태 공간 모델링에 적합하도록 재구성한 '영역 중첩 선택적 스캔' 방식을 제안했습니다.
2. 새로운 MIL 프레임워크: 해상도 인식 정적 인코딩과 영역 적응적 동적 모델링을 분리하여 처리하는 MoEMambaMIL 을 개발했습니다. 이는 구조적 편향 (inductive bias) 과 유연한 전문화를 동시에 달성합니다.
3. 성능 및 효율성: 선형 복잡도를 유지하면서 WSI 분석에서 최신 (State-of-the-Art) 성능을 달성했습니다.

4. 실험 결과 (Results)

• 데이터셋: TCGA 신장암 (Kidney), 간암 (Liver Cancer), CAMELYON17 (유방암 전이) 등 3 개의 대규모 다중 클래스 WSI 데이터셋에서 평가되었습니다.
• 성능:
  • ResNet, UNI, GigaPath 등 다양한 특징 추출기 (Feature Extractor) 를 사용했을 때, 기존 MIL 방법 (TransMIL, CLAM 등) 과 최신 Mamba 기반 방법 (MambaMIL, BiMambaMIL 등) 보다 9 개의 하위 작업에서 일관되게 우수한 성능을 보였습니다.
  • 특히 TCGA Kidney 데이터셋에서 UNI 특징을 사용할 때 F1 점수 **95.78%**를 기록하여 최상위 성능을 달성했습니다.
  • CAMELYON17 과 같은 도전적인 데이터셋에서도 높은 일반화 성능을 입증했습니다.
• Ablation Study:
  • 해상도 인식 정적 전문가를 제거하거나 (WO-R), 동적 MoE 구조를 단일 전문가로 대체 (WO-MoE) 할 경우 성능이 크게 저하됨을 확인하여, 두 구성 요소의 상호 보완적 중요성을 입증했습니다.
  • MoE-Mamba 아키텍처의 최적 설정은 6 레이어, 토큰당 2 개의 전문가 활성화, 부하 균형 가중치 $\lambda=0.001$임을 확인했습니다.

5. 의의 및 결론 (Significance)

• 구조와 효율성의 통합: WSI 분석의 핵심인 '계층적 공간 구조'와 '긴 시퀀스 처리 효율성'을 동시에 해결하는 새로운 패러다임을 제시했습니다.
• 생물학적 해석 가능성: 단순한 패치 집합이 아닌, 조직의 계층적 포함 관계를 보존하는 시퀀스 모델링을 통해 병리학적 해석에 더 부합하는 표현 학습이 가능해졌습니다.
• 확장성: 선형 시간 복잡도를 가지면서도 MoE 를 통해 모델 용량을 확장하여, 기가픽셀 규모의 대규모 병리 이미지 분석에 매우 적합한 아키텍처임을 입증했습니다.

이 논문은 상태 공간 모델 (SSM) 과 혼합 전문가 (MoE) 를 병리 이미지 분석에 효과적으로 접목하여, 기존 트랜스포머 기반 방법들의 계산적 한계를 극복하고 구조적 정확도를 높인 획기적인 접근법으로 평가됩니다.