USEMA: a Scalable Efficient Mamba Like Attention for Medical Image Segmentation

기술 요약: USEMA – 의료 영상 분할을 위한 확장 가능한 효율적 Mamba 유사 어텐션

문제 제기
정확한 의료 영상 분할은 진단 파이프라인에 필수적이지만, 제한된 학습 데이터와 복잡하고 중첩된 객체 기하학 등 고유한 도전에 직면해 있습니다. 비전 트랜스포머 (ViT) 는 자기 어텐션을 통해 전역 의존성을 포착할 수 있는 능력을 제공하지만, 의료 영상 분야에서의 적용은 두 가지 주요 문제로 인해 방해받고 있습니다:

1. 이차적 계산 복잡도: 순수 자기 어텐션은 시퀀스 길이에 대해 $O(n^2)$의 비용을 발생시켜, 고해상도 의료 영상의 경우 계산적으로 실행 불가능합니다.
2. 어텐션 분산: 시퀀스 길이가 증가함에 따라 어텐션 행렬의 점수 (소프트맥스 출력) 가 균일하게 0 을 향해 분산되는 경향이 있습니다. 이 현상은 약한 대수의 법칙에 이론적으로 근거하며, 모델이 중요한 키와 값을 선택하는 능력을 저해하여 어텐션 메커니즘을 실질적으로 균일한 평균화 과정으로 축소시킵니다.

방법론
저자들은 **SEMA(Scalable Efficient Mamba-like Attention)**라고 불리는 새로운 어텐션 메커니즘을 통합한 하이브리드 UNet 아키텍처인 USEMA를 제안합니다. 이 방법론은 다음과 같은 구성 요소를 통해 순수 트랜스포머와 기존 Mamba 기반 모델의 한계를 모두 해결합니다:

• SEMA 메커니즘: SEMA 는 선형 복잡도를 유지하면서 전체 자기 어텐션의 행동을 모방하기 위해 국소적 및 전역적 근사를 결합합니다.
  • 국소적 구성 요소: 긴 시퀀스에서 내재된 분산 문제를 피하기 위해 윈도우 어텐션을 사용하여 토큰의 국소화와 집중을 유지합니다.
  • 전역적 구성 요소: 이차적 비용 없이 전역 컨텍스트를 포착하기 위해 SEMA 는 이론적으로 일관된 산술 평균을 사용합니다. 이 항은 분산 부등식에서 관찰된 $O(1/n)$ 경계를 반영하며 어텐션의 전역적 측면을 근사화합니다. 전역 항은 모든 값 벡터의 평균으로 정의되어 시퀀스 전체에 브로드캐스트됩니다.
  • 통합: SEMA 블록은 Mamba 와 유사한 조건부 위치 임베딩, 레이어 정규화, 그리고 게이트 메커니즘 (예: SiLU 활성화 및 깊이별 합성곱) 을 활용하여 Mamba 매크로 구조에 삽입됩니다.
• USEMA 아키텍처: 이 모델은 대칭적인 U 자형 구조 (UNet) 를 채택합니다.
  • 인코더: 잔여 블록 다음에 SEMA 블록으로 구성됩니다. 특징은 어텐션 메커니즘에 들어가기 전에 $(B, C, H, W)$에서 $(B, C, HW)$로 재형상화됩니다. 회전 위치 임베딩과 국소 강화 위치 인코딩이 공간 인식을 증강하는 데 사용됩니다.
  • 병목: 공간 압축으로 인해 시퀀스 길이가 계산적으로 관리 가능한 정도로 충분히 작기 때문에, 병목 구간은 윈도우 어텐션 대신 전체 자기 어텐션을 활용합니다.
  • 디코더: 잔여 블록과 전치 합성곱으로 구성되며, 인코더의 계층적 특징을 융합하기 위해 스킵 연결을 활용합니다.

주요 기여

1. 하이브리드 어텐션 설계: 본 논문은 산술 평균을 자기 어텐션의 전역 근사로 이론적으로 정당화하는 SEMA 를 도입하여, 점수 분산을 방지하기 위해 국소 윈도우 어텐션을 보완합니다.
2. 아키텍처 통합: USEMA 는 UNet 프레임워크 내에서 CNN 의 국소 특징 추출 효율성과 Mamba 유사 어텐션의 전역 모델링 능력을 성공적으로 결합합니다.
3. 종합 벤치마킹: 이 모델은 다양한 해상도와 분할 목표 (장기, 도구, 세포 분할) 를 가진 세 가지 다양한 의료 영상 모드 (복부 MRI, 내시경, 현미경) 에서 평가되었습니다.

실험 결과
USEMA 는 세 가지 데이터셋에서 최신 트랜스포머 기반 모델 (UNETR, Swin-UNETR, nnFormer) 과 Mamba 기반/영향 받은 모델 (U-Mamba, Swin-UMamba, Mamba UNet, MLLA-UNet) 과 비교 벤치마킹되었습니다:

• 복부 MRI (MICCAI 2022 AMOS): USEMA 는 Dice 유사도 계수 (DSC) 0.7704와 정규화된 표면 거리 (NSD) 0.8345를 달성했습니다. 이는 최고의 트랜스포머 베이스라인 (nnFormer) 보다 DSC 에서 5.83% 더 높고, 최고의 Mamba 베이스라인 (U-Mamba Enc) 보다 DSC 에서 10.53% 더 높은 성능입니다. 주목할 만하게도 USEMA 는 52M 개의 파라미터로 이러한 결과를 달성했는데, 이는 최고 성능의 트랜스포머 모델보다 13% 적은 수치입니다.
• 내시경 (MICCAI 2017): USEMA 는 DSC 0.6463과 NSD 0.6621을 달성하여, DSC 기준 트랜스포머 베이스라인보다 약 5.6%, Mamba 베이스라인보다 약 13.5% 더 높은 성능을 보였습니다.
• 현미경 (NeurIPS 2022): 이 인스턴스 분할 작업에서 USEMA 는 F1 점수 0.5791을 달성하여 최고의 트랜스포머 베이스라인보다 8.61%, Mamba 영향 받은 MLLA-UNet 보다 19.23% 더 높은 성능을 보였습니다. 또한 U-Mamba Enc 보다 43% 적은 파라미터를 사용했습니다.

애블레이션 연구
애블레이션 연구는 전역 어텐션 근사의 필요성을 확인했습니다. 산술 평균 항을 제거하고 (윈도우 어텐션만 사용) 실험한 결과 모든 데이터셋에서 성능이 저하되었습니다 (예: 복부 MRI 데이터셋에서 DSC 0.013 감소). 이는 전역 항의 기여를 검증합니다.

의의 및 주장
저자들은 USEMA 가 계산 효율성과 분할 정확도 사이의 균형을 효과적으로 맞춘다고 주장합니다. 순수 자기 어텐션의 분산 특성과 Mamba 의 선형 복잡도를 활용함으로써, USEMA 는 다음과 같은 확장 가능한 솔루션을 제공합니다:

• 다양한 2D 의료 영상 벤치마크에서 순수 합성곱, 트랜스포머 기반, Mamba 기반 모델을 능가합니다.
• 전체 자기 어텐션 트랜스포머에 비해 계산 비용을 줄입니다.
• 이론적으로 근거한 효율적인 전역 근사 (산술 평균) 가 의료 영상 분석에서 국소 어텐션 메커니즘을 성공적으로 보완할 수 있음을 보여줍니다.

본 논문은 평균화 메커니즘을 학습되거나 가중된 평균으로 확장하고 USEMA 를 대규모 병리 영상에 적용할 계획을 conclude 하지만, 이러한 특정 미래 분야에서 현재 성공을 주장하지는 않습니다.