기술 요약: Patch-MoE Mamba

문제 제기

의료 영상 분할은 주로 합성곱 신경망 (CNN) 과 트랜스포머에 크게 의존합니다. 그러나 CNN 은 국소 수용 영역에 본질적으로 제한을 받아 장기 의존성을 모델링하는 데 어려움을 겪는 반면, 트랜스포머는 종종 2 차 계산 및 메모리 복잡도를 초래합니다. 최근, 선형 시퀀스 복잡성을 가진 효율적인 대안으로 상태 공간 모델 (SSM), 특히 Mamba 기반 아키텍처가 등장했습니다.

그럼에도 불구하고 기존 Mamba 기반 분할 모델은 두 가지 중요한 한계에 직면해 있습니다:

1. 국소 공간 구조의 붕괴: 대부분의 모델은 픽셀 단위 래스터 스캐닝을 사용하여 2D 특징 맵을 1D 시퀀스로 변환합니다. 이 과정은 시퀀스 내에서 공간적으로 인접한 픽셀들을 분리시켜 국소 2D 공간 일관성을 약화시킵니다. 이는 정확한 경계 위치 파악과 국소 구조 보존이 필요한 밀집 예측 작업에 해롭습니다.
2. 비유연한 방향성 융합: 현재 방법들은 일반적으로 단순 합계를 통해 여러 스캐닝 방향의 출력을 융합합니다. 이러한 고정된 집계는 모든 공간 위치에서 모든 방향이 동등하게 중요하다고 가정하여, 의료 대상의 다양한 크기, 모양, 질감 및 경계 복잡성 (예: 작고 대비가 낮은 병변 vs 크고 불규칙한 구조) 에 적응하지 못합니다.

방법론

저자들은 앞서 언급한 한계를 해결하기 위해 패치 순서형 전문가 혼합 (MoE) 상태 공간 아키텍처인 Patch-MoE Mamba를 제안합니다. 이 아키텍처는 U-Net 스타일 설계를 따르며, VM-UNetV2 의 표준 비주얼 상태 공간 (VSS) 블록을 새로운 Patch-MoE VSS 블록으로 대체합니다.

1. 계층적 패치 순서형 스캐닝

개별 픽셀을 래스터 스캐닝하는 대신, 이 방법은 특징 맵을 국소 패치로 분할하고 패치 순서형 방식으로 방향성 스캐닝을 수행합니다.

• 메커니즘: 특징 맵이 주어지면, 공간 그리드를 $p \times p$ 크기의 겹치지 않는 패치들로 나눕니다. 각 패치 내에서 픽셀들은 다음 패치로 이동하기 전에 행 우선 (row-major) 순서로 열거됩니다.
• 이점: 이는 패치 내의 모든 픽셀이 시퀀스에서 연속적으로 유지되도록 하여, 패치 간 장기 의존성 모델링을 가능하게 하면서도 국소 공간 이웃을 보존합니다.
• 계층적 전략: 다중 스케일 컨텍스트를 포착하기 위해 모델은 패치 크기 집합 $\{p^{(1)}, p^{(2)}, \dots\}$을 사용합니다. 작은 패치는 세밀한 경계 세부 사항을 포착하고, 큰 패치는 더 거친 구조를 모델링합니다. 서로 다른 스캐닝 방향 (전진, 역전, 너비 - 높이 전진, 너비 - 높이 역전) 은 이방성 구조를 처리하기 위해 서로 다른 패치 크기를 활용할 수 있습니다.

2. MoE 기반 방향성 융합

고정된 합계를 대체하기 위해 저자들은 방향성 출력을 "전문가"로 취급하는 적응형 융합 모듈을 도입합니다.

• 전문가 구성: 이 모듈은 다섯 가지 전문가를 활용합니다:
  • 네 가지 방향별 전문가 ($E_1$부터 $E_4$): 네 가지 방향 Mamba 스캐너의 정규화된 출력.
  • 하나의 연결 전문가 ($E_{concat}$): 네 개의 정규화된 맵을 채널 차원을 따라 연결한 후 $1\times1$ 합성곱을 통해 다시 투영하여 형성됩니다. 이 전문가는 방향 간 상관관계를 포착하고 방향에 구애받지 않는 전체적인 관점을 제공합니다.
• 공간 인식 게이트: 라우터 네트워크는 공간적 컨텍스트에 따라 가중치를 계산하여 이러한 전문가들을 적응적으로 융합합니다. 라우터는 국소 컨텍스트 (깊이별 합성곱을 통해) 와 전역 컨텍스트 (적응형 평균 풀링을 통해) 의 혼합에 의해 구동되며, 학습 가능한 스칼라 $\alpha$에 의해 제어됩니다.
• 잔차 안정화: 잔차 바이패스가 원시 방향성 출력을 융합된 결과에 다시 추가합니다 ($Z_l = \hat{Y}_l + \sum Y_l^{(i)}$). 이는 견고한 기본 응답을 보장하고 강력한 방향성 신호를 보존하여 훈련을 안정화합니다.

3. 의미 및 세부 정보 주입 (SDI)

인코더 출력은 SDI 모듈 (U-Net v2 에서 적응됨) 에 의해 처리되어, 더 높은 수준의 의미 정보를 주입하고 하위 수준의 더 세밀한 세부 사항을 헤다마르 곱을 통해 통합함으로써 특징 맵을 강화합니다.

주요 기여

1. 패치 순서형 스캐닝: 픽셀 단위 래스터 스캐닝을 패치 순서형 탐색으로 대체하는 새로운 스캐닝 메커니즘으로, 상태 공간 모델의 효율성을 유지하면서 국소 2D 공간 일관성을 효과적으로 보존합니다.
2. MoE 기반 방향성 융합: 단순 합계를 학습 가능한 게이트 네트워크로 대체하는 적응형 융합 모듈입니다. 이는 공간 컨텍스트에 기반하여 여러 방향 전문가와 연결 전문가에 동적으로 가중치를 부여하여, 모델이 다양한 병변 규모와 복잡성에 적응할 수 있게 합니다.
3. 계층적 다중 스케일 모델링: 계층적 패치 크기의 통합을 통해 네트워크가 동시에 세밀한 경계 세부 사항과 거친 의미 구조를 포착할 수 있게 합니다.
4. 잔차 안정화: 원시 방향성 출력을 위한 잔차 경로의 포함은 훈련 안정성을 보장하고 중요한 방향성 신호의 손실을 방지합니다.

실험 결과

저자들은 Patch-MoE Mamba 를 다섯 가지 공개 폴립 분할 데이터셋 (Kvasir-SEG, ClinicDB, ColonDB, ETIS, CVC-300) 과 두 가지 피부 병변 분할 데이터셋 (ISIC 2017, ISIC 2018) 에서 평가했습니다.

• 성능: Patch-MoE Mamba 는 모든 데이터셋에서 최첨단 (SOTA) 성능을 달성하여 CNN 기반 (U-Net, U-Net v2) 및 Mamba 기반 베이스라인 (VM-UNet, VM-UNetV2) 을 일관되게 능가했습니다.
  • 도전적인 ETIS 데이터셋 (낮은 대비와 작은 폴립으로 유명) 에서 모델은 VM-UNetV2 의 72.56% 에서 **74.04%**로 Dice 점수를 향상시켰습니다.
  • ISIC 2018에서는 이전 최고 기록인 88.36% 를 능가하는 **89.34%**의 Dice 점수를 달성했습니다.
• Ablation 연구:
  • 래스터 스캐닝을 패치 순서형 스캐닝으로만 대체한 것만으로도 평균 Dice 가 83.46% 에서 84.02% 로 향상되었습니다.
  • MoE 융합 모듈을 추가하면 평균이 **84.30%**로 추가로 상승했습니다.
  • 초기 단계에서는 거친 패치 크기를, 깊은 단계에서는 세밀한 패치 크기를 사용하는 계층적 패치 크기가 가장 좋은 결과를 낳았습니다.
• 계산 복잡성: 패치 순서형 스캐닝 자체는 계산 비용을 증가시키지 않습니다 (토큰 수가 변하지 않음). 그러나 연결 전문가의 포함으로 인해 파라미터는 약 70M 로, FLOPs 는 약 28G 로 증가합니다. 저자들은 MoE 메커니즘에 필요한 더 밀집된 채널 상호작용으로 인해 이러한 오버헤드가 예상된다고 지적합니다.

중요성 및 주장

이 논문은 Patch-MoE Mamba 가 의료 영상 분할을 위한 일반적이고 효과적인 솔루션을 제공한다고 주장합니다. 기존 Mamba 모델의 특정 한계인 국소 공간 구조의 손실과 방향성 융합의 경직성을 해결함으로써, 제안된 아키텍처는 다양한 크기, 모양 및 경계 복잡성을 가진 다양한 의료 대상에 대한 우수한 적응력을 보여줍니다. 결과는 패치 순서형 스캐닝을 통한 국소 공간 이웃의 보존과 적응형 전문가 융합의 결합이, 특히 낮은 대비나 작은 병변이 포함된 도전적인 시나리오에서 경계 위치 파악 및 분할 정확도를 크게 향상시킨다고 시사합니다.