논문 리뷰 : Vision Mamba: Efficient Visual Representation Learning with Bidirectional State Space Model

용어 설명

• State Space Model(SSM): 연속적인 시스템을 이산화하여 sequence data를 처리하는 수학적 모델입니다. 기존 RNN이나 CNN의 대안으로 긴 sequence 처리에 효율적입니다.
• Mamba: 하드웨어 인식(hardware-aware) 알고리즘을 적용하여 학습과 추론 속도를 극대화한 최신 SSM 기반 architecture입니다.
• Bidirectional SSM: 시각적 데이터의 공간적 맥락(spatial context)을 파악하기 위해 sequence를 정방향(forward)과 역방향(backward) 양쪽으로 스캔하여 정보를 압축하는 Vim의 핵심 메커니즘입니다.
• Position embeddings: Sequence 모델링에서 patch들의 공간적 위치 정보를 모델에 알려주기 위해 더해지는 신호(signal)입니다.

Purpose of the Paper

• 기존 Vision Transformer(ViT)의 self-attention 메커니즘은 이미지 해상도가 커질수록 연산량과 메모리 사용량이 이차(quadratic) 수준으로 폭증하는 치명적인 한계가 있었습니다.
• 최근 language model에서 각광받는 SSM(특히 Mamba)은 선형적인 연산 복잡도를 가지지만, 단방향(unidirectional) 처리 방식과 위치 인식(positional awareness)의 부재로 인해 visual data 처리에 부적합했습니다.
• 본 연구는 self-attention 없이도 전역 문맥(global context)과 위치 정보를 효과적으로 학습할 수 있는 순수 SSM 기반의 generic vision backbone인 Vim(Vision Mamba)을 제안하여, 고해상도 visual data 처리의 병목을 해결하고자 합니다.

Key Contributions

• Bidirectional State Space Model 도입: 이미지를 patch sequence로 변환한 뒤, 정방향과 역방향 모두로 SSM을 적용하여 데이터 의존적인 전역 시각 문맥(data-dependent global visual context)을 모델링합니다.
  • Novelty: 기존 language 모델용 단방향 Mamba를 vision task에 맞게 양방향 구조로 혁신적으로 확장했습니다.
• 위치 인식(Position-Aware) 시각 이해 달성: Position embeddings를 결합하여 위치 민감도가 높은 visual data의 특성을 반영합니다.
  • Novelty: Attention 없이도 Transformer와 동일한 수준의 공간적 표현 능력을 확보했습니다.
• Middle class token 전략: Classification을 위해 sequence의 맨 앞이나 뒤가 아닌 중간(middle)에 class token을 배치합니다.
  • Novelty: 기존 ViT의 헤드 토큰 방식을 탈피하여, SSM의 순환(recurrent) 특성과 ImageNet의 중심 객체 prior를 극대화하는 맞춤형 구조를 설계했습니다.

Experimental Highlights

• 성능 달성: ImageNet classification, COCO object detection, ADE20K semantic segmentation 등 주요 vision task에서 기존의 강력한 baseline인 DeiT보다 우수한 성능을 기록했습니다.
• 효율성 극대화: 고해상도(1248x1248) 이미지 배치 추론 시, DeiT 대비 2.8배 더 빠르고 GPU 메모리를 86.8% 절약했습니다.
• High-resolution 환경 적용: 기존 ViT 기반 모델들이 고해상도 처리를 위해 window attention 같은 2D prior를 억지로 주입해야 했던 것과 달리, Vim은 2D prior 없이 sequence representation을 그대로 유지하며 고해상도 object detection과 instance segmentation을 성공적으로 수행했습니다.

Limitations and Future Work

• 명확한 한계점(Limitations): 본 논문은 supervised learning 형태의 image classification 사전 학습에 일차적인 초점을 맞추었으며, self-supervised learning이나 대규모 multi-modality 환경에서의 검증은 아직 초기 단계입니다. 이는 모델의 범용성을 입증하기 위해 반드시 넘어야 할 산입니다.
• 향후 연구 방향(Future Work): Vim의 양방향 구조와 위치 임베딩은 Mask Image Modeling(MIM)이나 CLIP 스타일의 멀티모달 pre-training에 매우 적합하므로 이 방향으로의 확장이 필요합니다.
• 발전 가능성: 선형적인 메모리 효율성을 바탕으로 의료 이미지, 위성(remote sensing) 이미지, 긴 비디오(long videos) 등 기존 Transformer가 감당하기 어려웠던 초고해상도 시각 데이터 분석에 혁신적인 도구가 될 잠재력이 큽니다.

Overall Summary

이 논문은 효율적인 sequence 모델링에 탁월한 Mamba 모델을 컴퓨터 비전 영역으로 확장하여, attention 메커니즘 없이도 강력한 시각 표현 학습이 가능한 Vision Mamba(Vim)를 제안합니다. 양방향 SSM과 위치 임베딩을 통해 visual data의 공간적 맥락을 정확히 포착하면서도 선형적인 계산 복잡도를 달성하여 고해상도 이미지 처리의 근본적인 한계를 극복했습니다. 이는 향후 거대 vision-language model이나 초고해상도 분석이 필수적인 다양한 응용 분야의 차세대 backbone으로 자리매김할 수 있는 중요한 연구입니다.

쉬운 설명

농구 경기나 단체 무용 퍼포먼스를 상상해 보세요. 코트 전체의 흐름(global context)을 파악해야 할 때, 기존 Transformer(ViT)는 코트 위의 모든 선수들이 서로 1:1로 일일이 눈을 마주치며(self-attention) 상황을 파악하는 방식이었습니다. 선수가 많아질수록(고해상도) 확인해야 할 경우의 수가 기하급수적으로 늘어나 체력(연산량)이 고갈되죠.

반면 Vim은 마치 코트 끝에서 끝으로 한 번 달려가면서(정방향 스캔), 그리고 다시 반대로 돌아오면서(역방향 스캔) 코트 전체의 상황을 머릿속에 효율적으로 압축해서 기억하는 뛰어난 포인트 가드와 같습니다. 체력(메모리) 소모는 아주 적으면서도 경기장 전체의 흐름을 완벽하게 읽어내는 혁신적인 전술이라고 볼 수 있습니다.

 

 

HWC 크기의 이미지 준비 패치 사이즈는 16으로 쪼갬
각 패치를 리니어 프로젝션 시켜서 패치 토큰들 생성..

토큰 중간이나 맨앞에 미리 준비한 클래스 토큰 넣어둠

포지셔널 인코딩 주입. (위치 정보 보존)




그리고 vim 블록을 통과 하는데 이는 어케 생겼냐면
입력된 패치 시퀀스가 가장 먼저 정규화(Norm) 층을 통과 후 선형 변환 거침
x랑 z 투배럭으로 쪼개짐 z는 게이팅 용도

x를 복사해서 다시 양방향으로 쪼갬. (무슨 말이냐면 데이터 하나의 순서를 뒤집음)

각각 1d 컨브 때려서 silu 통과 후 지역 정보 교류 

이 후 B C 델타 만들고 각각 ssm 통과 시킴.

아까 빼둔 z빼서 양쪽 모두 결과에 요소별 곱
그래서 더해주고 다시한번 선형 변환 후 결과에 잔차 연결

이게 vim 블록


암튼 이걸 통과해서 블록개수만큼 반복

그리고 클래스 토큰 노말라이즈 해서 헤드 달고 클래시피 케이션 진행

loss 주고 학습