2DMamba 이미지 표현을 위한 효율적인 2차원 상태공간 모델

초록

본 논문은 2차원 이미지 구조를 그대로 보존하면서 선형 복잡도와 높은 병렬성을 제공하는 2DMamba를 제안한다. 2DMamba는 기존 1차원 Mamba의 선택적 스캔을 2차원 가로·세로 스캔으로 확장하고, GPU 친화적인 연산자를 설계해 giga‑pixel 수준의 병리 슬라이드 이미지와 일반 자연 이미지 모두에서 성능 향상을 입증한다.

상세 분석

2DMamba는 상태공간 모델(SSM)의 선택적 메커니즘을 2차원 이미지에 직접 적용한 최초의 시도이다. 기존 Mamba는 입력‑종속 Â, B̂, Ĉ 파라미터를 통해 불필요한 상태를 잊어버리는 방식으로 1차원 시퀀스에 최적화돼 있었다. 그러나 이미지와 같은 2D 데이터에 1D로 플래튼하면 인접 픽셀 간 거리가 멀어지는 ‘공간 불일치(spatial discrepancy)’가 발생한다. 2DMamba는 이를 해소하기 위해 가로 스캔을 먼저 수행하고, 그 결과를 세로 스캔에 입력함으로써 각 픽셀의 은닉 상태가 맨 왼쪽 위(또는 위‑좌) 영역들의 정보를 맨해튼 거리 기반으로 누적하도록 설계한다. 수식 (7)은 hᵢ,ⱼ = Σ_{i’≤i, j’≤j} Ā^{(i−i’+j−j’)} B̄ x_{i’,j’} 형태로, 이는 1D Mamba의 순차적 거리 i−i’에 비해 2D 거리 i−i’+j−j’를 사용해 공간 연속성을 보장한다.

연산 효율성 측면에서 2DMamba는 GPU 메모리 계층을 고려한 ‘2D 선택적 스캔 연산자’를 도입한다. 기존 1D Mamba는 긴 시퀀스를 타일링해 SRAM에 적재·스캔·재저장하는 방식으로 O(L) 메모리 접근을 달성했지만, 2D로 확장하면 타일 경계에서의 종속성이 복잡해진다. 저자들은 가로·세로 스캔을 각각 독립적인 1D 스캔으로 구현하고, 각 상태 차원 d에 대해 병렬화함으로써 메모리 전송을 최소화하고 연산량을 O(H·W)로 유지한다. 또한, 파라미터 C는 N개의 상태 차원을 단일 출력으로 집계해 역전파 시 중간 상태를 재계산하도록 설계해 메모리 사용량을 크게 줄였다.

임상 적용에서는 WSI를 패치 단위로 타일링하고, 조직 패치는 사전 학습된 병리 특성 추출기로, 비조직 영역은 학습 가능한 토큰 p로 대체해 2D 형태를 유지한다. 이렇게 만든 특성 맵을 U개의 2DMamba 블록에 통과시킨 뒤, 주의 기반 어그리게이터를 통해 슬라이드 수준의 표현을 얻는다. 실험 결과 10개의 공개 WSI 데이터셋에서 AUC, F1, Accuracy, C‑index 모두 25% 수준의 절대적 향상을 기록했으며, 자연 이미지에서는 VMamba와 결합해 ADE20K에서 mIoU를 0.50.7%, ImageNet‑1K에서 Top‑1 정확도를 0.2% 개선했다.

전반적으로 2DMamba는 2D 구조 보존, 선택적 상태 업데이트, GPU 친화적 병렬 구현이라는 세 축을 동시에 만족시켜, 대규모 의료 영상 및 일반 비전 작업에서 기존 SSM이나 Transformer 기반 모델이 갖는 계산·성능 한계를 효과적으로 극복한다는 점이 가장 큰 공헌이다.