Vision Mamba: Efficient Visual Representation Learning with Bidirectional State Space Model

[Github]

[arXiv](2024/01/17 version v1)

 

 

 

Abstract

Bidirectional Mamba block을 사용한 새로운 vision backbone인 Vision Mamba (Vim) 제안

 

 

Method

1. Preliminaries
2. Vision Mamba
3. Vim Block

 

Preliminaries

필독!!! Mamba

Mamba: Linear-Time Sequence Modeling with Selective State Spaces

 

Vision Mamba

 

Transformer block이 mamba block으로 바뀌었을 뿐 절차는 똑같다.

 

입력 이미지를 차원이 D인 J개의 토큰으로 projection 하고 position embedding.

최종 mamba block의 출력을 정규화하고 MLP를 통과하여 출력을 얻는다.

 

Vim Block

기존 Mamba와 유일한 차이점은 토큰 시퀀스를 역방향으로 모델링할 수 있는 backward branch의 유무이다.

E = 확장된 차원, N = SSM 내부 state의 차원

 

 

 

Experiment

솔직히 논문의 성능 비교는 믿을 게 못되지만 어쨌든 매우 빠르고 메모리 소모가 적다.