멀티디멘션 상태공간 모델 Mamba2D, 비전 작업 혁신

초록

Mamba2D는 2차원 스캔을 기반으로 한 네이티브 다차원 상태공간 모델(SSM)로, 이미지 분류에서 84.0%(M2D‑T)·85.3%(M2D‑S) top‑1 정확도를 달성하고, COCO 객체 검출 52.2 AP·ADE20K 세그멘테이션 51.7 mIoU를 기록했다. 1D 스캔의 한계를 극복하고 CUDA 파동프론트 스캔을 통해 효율적인 병렬 처리를 구현한다.

상세 분석

Mamba2D는 기존 비전 SSM이 1차원 시퀀스 모델을 이미지에 강제 적용하면서 발생하는 공간 구조 손실을 근본적으로 해결한다. 핵심 아이디어는 두 축(t, z)에 대해 독립적인 미분 방정식을 정의하고, ZOH(Zero‑Order‑Hold) 이산화를 각각 적용해 2D 상태 전이를 구현하는 것이다. 이렇게 하면 각 픽셀의 은닉 상태 h(t, z)는 바로 위와 왼쪽 이웃으로부터만 영향을 받으며, Manhattan 거리 기반의 전역 의존성을 유지한다.

선택적 스캔 메커니즘은 Mamba의 S6 구조를 2D에 그대로 확장한 것으로, 입력에 따라 A_t, B_t, A_z, B_z 를 동적으로 조정한다. 이는 기존 S4ND가 선형‑시간 불변 파라미터에 머물렀던 점과 달리, 입력‑의존적 가중치를 통해 장거리와 단거리 정보를 균형 있게 통합한다. 또한, 스캔 방향을 블록마다 교차(top‑left→bottom‑right ↔ bottom‑right→top‑left)시켜 인과적(causal) 편향을 완화하고 양방향 컨텍스트를 효율적으로 포착한다.

병렬 구현 측면에서는 2D 재귀가 대각선(파동프론트) 단위로 독립적이라는 사실을 이용해 CUDA 커널을 설계했다. 각 대각선은 동시에 계산되고, 다음 대각선은 이전 결과가 메모리에 기록된 뒤 실행된다. 연산량이 적은 메모리 바인드 특성을 고려해 A·Δ와 B·Δ·x 를 레지스터에서 즉시 계산하고, 은닉 상태 h는 역전파 시 재계산함으로 메모리 사용량을 크게 절감한다.

아키텍처는 MetaFormer 스타일의 계층적 설계와 결합된다. 초기 Conv‑Stem으로 패치 임베딩을 만든 뒤, 네 단계의 하이‑레졸루션 블록에 M2D‑SSM 토큰 믹서를 삽입한다. 믹서는 2D‑SSM 경로와 depthwise separable convolution 기반 로컬 경로를 병렬로 수행해 장거리와 지역 정보를 보완한다. 하위 레벨에서는 전통적인 self‑attention(QK‑Norm·RoPE)으로 전환해 연산 효율과 표현력을 균형 있게 유지한다.

실험 결과는 모델 규모별(7 M ~ 50 M 파라미터)로 기존 1D‑SSM(예: VMamba, MambaVision)과 2D‑SSM(예: V2M, 2D‑VMamba)보다 명확히 앞선다. 특히 M2D‑T(27 M)와 M2D‑S(50 M)는 각각 84.0%·85.3% top‑1 정확도를 기록해, 동일 파라미터 범위 내 최고 성능을 달성한다. 다운스트림 객체 검출과 세그멘테이션에서도 3× 스케줄 훈련 시 COCO AP 52.2, ADE20K mIoU 51.7을 얻어, 선형 복잡도와 높은 일반화 능력을 동시에 입증한다.

전반적으로 Mamba2D는 2D 상태공간 모델을 처음부터 수학적으로 재정의하고, 효율적인 파동프론트 병렬화를 구현함으로써 비전 분야에서 SSM의 활용 가능성을 크게 확장한다. 이는 차후 고해상도 의료 영상·위성 이미지·비디오 등 다차원 시계열 데이터에 대한 연구에 중요한 기반이 될 것으로 기대된다.