Partial Ring Scan: Revisiting Scan Order in Vision State Space Models

기술 요약: 부분 링 스캔(Partial Ring Scan): 비전 상태 공간 모델에서의 스캔 순서 재고

문제 정의
상태 공간 모델(SSM), 특히 비전 마마(Vision Mamba) 변형 모델들은 이미지의 2D 구조를 사전 정의된 스캔 순서(예: 래스터, 서펜타인)에 따라 1D 토큰 시퀀스로 직렬화하여 선형 시간 복잡도로 시퀀스를 처리하는 효율적인 대안으로 부상했습니다. 저자들은 이러한 스캔 순서가 공간적 인접성을 결정하는 매우 중요하면서도 흔히 간과되는 요소라고 주장합니다.

기존의 고정 경로 스캔(래스터, 양방향, 대각선)은 플립(flip)과 같은 특정 변환 하에서만 시퀀스-공간 정렬을 보존합니다. 평면 내 회전(in-plane rotation)이 발생할 경우, 이러한 경로는 전역적인 재인덱싱(reindexing)과 패딩 문제를 일으켜 객체의 연속성을 파괴하고 재귀적 커널의 이웃 관계를 어긋나게 만듭니다. 이는 SSM이 의미론적 모델링 대신 국소적 일관성을 복구하는 데 역량을 소모하게 만들며, 결과적으로 성능 저하(Top-1 정확도 1~2% 하락)와 기하학적 변환에 대한 강건성 감소를 초래합니다.

방법론: PRISMamba
본 논문은 회전 특화 학습 없이도 기하학적 구조와 시퀀스 인접성을 정렬하도록 설계된 회전 강건 아키텍처인 **PRISMamba(Partial Ring Scan Mamba)**를 제안합니다. 이 방법은 세 가지 핵심 구성 요소로 이루어집니다:

1. 링 단위 교차 스캔 (Ring-by-Ring Alternating Scan):

   • 분할(Partitioning): 이미지를 이미지 중심으로부터의 유클리드 거리에 따라 동심원 형태의 링(ring)들로 분할합니다.
   • 링 내부 집계(Intra-Ring Aggregation): 각 링 내에서 픽셀들은 폐쇄 루프(인접한 링끼리 시계 방향과 반시계 방향을 교차함)를 따라 이동합니다. 경량화된 선택적 SSM이 이 루프를 처리하며, 출력은 평균화를 통해 링 디스크립터(ring descriptor)로 집계됩니다. 이 집계 방식은 순서에 무관(order-agnostic)합니다. 즉, 회전이 발생하더라도 루프 내에서 순환 이동(cyclic shift)만 발생할 뿐 평균화된 디스크립터는 변하지 않습니다.
   • 반경 방향 전파(Inter-Radial Propagation): 링 디스크립터들은 두 번째의 짧은 선택적 SSM을 사용하여 안쪽 링에서 바깥쪽 링으로 순차적으로 처리되며, 이를 통해 반경 방향으로 컨텍스트를 전파합니다. 이 2단계 설계는 선형 시간 복잡도($O(HW)$)를 유지하면서도 전역 경로 기반 직렬화의 취약성을 피합니다.

2. 부분 채널 필터링 (Partial Channel Filtering, PCF):

   • 효율성을 높이기 위해, 저자들은 채널이 재귀적 링 경로로 진입하기 전에 필터링하는 하드 라우팅(hard routing) 메커즘을 도입했습니다.
   • 채널의 돌출도(salience)는 전역 평균 풀링(GAP)을 통해 결정됩니다. 평균 임계값 이상의 돌출도를 가진 채널은 링 기반 SSM을 통과하는 반면, 나머지 채널은 경량 잔차 브랜치(residual branch)를 통해 이 무거운 연산을 우회합니다.
   • 이 접근 방식은 모든 채널을 다운스트림에서 처리하는 소프트 어텐션 메커니즘(예: SE, CBAM)과 대조적으로, FLOPs와 활성화 메모리를 줄이면서도 정확도를 유지합니다.

3. 쓰기 백 및 융합 (Write-Back and Fusion):

   • 반경 방향 출력을 다시 픽셀 그리드로 브로드캐스트하고, 1×1 프로젝션 및 잔차 연결을 통해 백본 스트림과 융합하여 공간 차원을 보존합니다.

주요 기여

• 체계적인 스캔 순서 분석: 본 논문은 스캔 순서가 "제1순위(first-class)" 하이퍼파라미터임을 보여주는 최초의 체계적인 연구를 제공합니다. 일반적인 래스터 및 서펜타인 스캔이 객체의 연속성을 파괴하고 회전에 매우 민감한 반면, 링 기반 스캔은 안정성을 유지함을 입증했습니다.
• 회전 안정적 트래버스(Rotation-Stable Traversal): 픽셀을 동심원 링으로 그룹화하고 순서에 무관한 집계를 수행하는 링 기반 스캔 방식을 도입했습니다. 이는 극방 좌표 매핑(polar remapping)이나 회전 특화 데이터 증강 없이도 SSM의 선형 시간 핵심 기능을 유지하면서 기하학적 불일치를 해결합니다.
• 부분 채널 필터링 (PCF): 가장 정보량이 많은 채널만을 재귀 경로로 라우팅하는 새로운 효율성 모듈입니다. 이는 계산 비용을 줄이고 처리량을 개선하는 동시에, 고신호 특징에 재귀적 업데이트를 집중시켜 정확도를 약간 향상시킵니다.
• 통합 아키텍처 (PRISMamba): 이러한 요소들을 결합하여 최첨단 성능을 달ing하는 통합 모델을 구축했습니다.

실험 결과
저자들은 ImageNet-1K, MS COCO, 그리고 회전/가림(occlusion) 스트레스 테스트를 통해 PRISMamba를 평가했습니다:

• ImageNet-1K 분류: PRISMamba는 A100 GPU에서 84.5% Top-1 정확도, 3.9G FLOPs, 3,054 img/s의 처리량을 달성했습니다. 이는 VMamba(82.6% 정확도, 5.6G FLOPs, 1,686 img/s)보다 정확도는 1.9% 높으면서 FLOPs는 30% 적게 사용하고 처리량은 약 1.8배 더 빠릅니다.
• MS COCO 탐지/세그멘테이션: 1× 스케줄(1280×800)에서 PRISMamba는 48.9 APbox와 43.2 APmask를 달성하여, VMamba 및 GroupMamba보다 10~19% 적은 FLOPs를 사용하면서도 더 높은 성능을 보였습니다.
• 회전 강건성: 30° 및 60° 회전 조건에서 PR-ISMamba는 거의 완벽한 안정성(0.10.2% 하락)을 유지한 반면, 고정 경로 베이스라인(VMamba, GroupMamba 등)은 1.02.0%의 성능 저하를 보였습니다.
• 가림(Occlusion) 강건성: 무작위 마스크 가림 테스트에서 PRISMamba는 링 평균화와 반경 방향 전파를 통해 국소적 손상을 희석함으로써, 고정 경로 변형 모델(최대 1.3% 하락)보다 현저히 낮은 성능 저하(≤0.6%)를 보였습니다.
• Ablation (절제 실험): PCF 모듈 단독으로 FLOPs를 약 15% 줄이고 처리량을 약 40% 증가시키면서도 정확도를 0.4% 향상시켰으며, 이는 하드 채널 필터링이 재귀적 SSM에서 소프트 어텐션보다 더 효과적임을 입증합니다.

의의 및 주장
본 논문은 스캔 순서 설계가 비전 SSM의 정확도, 효율성 및 강건성을 결정짓는 핵심적이지만 미개척된 요소임을 주장합니다. PRISMamba는 모델 크기를 키우거나 특수한 학습을 요구하지 않고도 경로 기반 모델의 시퀀스-기하학적 불일치를 해결할 수 있음을 보여줍니다. 저자들은 본 연구가 기하학적 변환에 대해 비전 SSM을 더 강건하게 만들고, 선택적 채널 처리를 통해 효율성을 높이는 원칙적이고 저비용인 접근법이라고 강조합니다.

한계 및 향후 과제
저자들은 현재 구현이 고정된 이미지 중심과 이산적인 링 너비에 의존하고 있어, 중심에서 벗어난 피사체나 극단적인 종횡비에 대해서는 최적화되지 않을 수 있음을 인정합니다. 향후 과제로는 링의 기원을 학습하고 너비를 조절하는 것, 내용 적응형 분할(content-adaptive partitions) 설계, 그리고 링 트래버스 개념을 시공간 도메인(비디오) 및 3D 볼륨으로 확장하는 것이 포함됩니다.