[논문리뷰] Vision Mamba: Efficient Visual Representation Learning with BidirectionalState Space Model

https://arxiv.org/abs/2401.09417

Abstract

SSM이 Visual data를 사용하는 것에 문제가 있음

position-sensitivity 이해와 전체 맥락에 대한 이해가 필요함

기존 mamba의 문제점

단방향 모델링 및 위치 인식 부족

⇒ 전역 시각적 컨텍스트 모델링을 위한 양방향 SSM과 위치 인식을 위한 position embedding을 통합한 Vision Mamba(Vim) 모델 제안

먼저 input image를 patch로 분할하고 Vim에 선형 투영

• 이미지 패치는 Vim 블록에서 시퀀스 데이터로 취급
• 제안된 양방향 선택적 상태 공간을 통해 시각적 표현을 효율적으로 압축
• position embedding은 공간적인 정보를 제공

Vision Mamba

Preliminaries

Mamba는 연속 시스템에서 영감을 받아 개발됨

기본적으로 모델들은 연속적인 시간 축을 따라 시퀀스 데이터를 처리

1. 연속 시스템

• $x(t)$ 는 시간 $t$에 따라 변함
• hidden state $h(t)$로 변환 후 다시 출력 $y(t)$로 매핑
• 이떄, $A$는 상태의 진화를 결정하는 행렬, $B,C$는 입력과 출력을 상태로 투영(projection)하는 역할

$$h'(t)=Ah(t)+Bx(t)$$

$$y(t)=Ch(t)$$

2. 이산화

• 연속적인 시스템을 이산적으로 표현하여 컴퓨터에서 계산
• 시간 척도 파라미터 $Δ$ 를 사용해 이산적인 $A,B$로 변환
• Zero-Order Hold (ZOH) 방법을 사용

$$A=exp(ΔA)$$

$$B=(ΔA)−1(exp(ΔA)−I)⋅ΔB$$

3. 이산 시스템

• $t$라는 시간 스텝을 따라 진행되며, $h_t$와 $y_t$는 이전 상태와 입력에 의해 결정됨
• 마지막으로 출력은 전역 컨볼루션을 통해 계산
• 컨볼루션 커널 $K$는 입력 시퀀스의 길이 $M$에 따라 결정

$$K=(CB,CAB,...,CAM−1B)$$

$$y=x∗K$$

Vision Mamba

표준 Mamba는 1차원 시퀀스를 위해 설계됨

1. 이미지 패치 변환

• 2차원 이미지 $t∈R^{H×W×C}$ 를 2차원 패치로 변환 $xp∈R^{J×(P2⋅C)}$
• $(H,W)$ 입력 이미지의 크기
• $C$ 채널 수 / $P$ 이미지 패치의 크기

2. 선형 투영 및 위치 임베딩

• 평탄화된 패치 $x_p$는 학습 가능한 투영 행렬 $W$를 통해 크기 $D$의 벡터로 변환
• 모든 패치를 동일한 크기의 벡터로 변환하여 모델의 입력 크기를 일관되게 유지
• 그 후, 위치 임베딩 $E_{pos}$가 추가되어, 각 패치의 위치 정보가 인코딩

$$T0=[t_{cls};t_1pW;t_2pW;⋅⋅⋅;t_JpW]+E_{pos},$$

3. 토큰 시퀀스

• 패치 시퀀스 전체를 나타내기 위해 클래스 토큰(class token)을 사용 > $t_{cls}$

4. Vim 인코더

• 토큰 시퀀스 $T_{l-1}$는 Vim 인코더의 각 레이어에 입력되고, $T_l$을 생성

$$T_l=Vim(T_{l−1})+T_{l−1}$$

• 출력 클래스 토큰 $T^L_0$ 정규화

$$f=Norm(T^0_L)$$

• 다층 퍼셉트론(MLP) 헤드로 전달해 최종 예측 $\hat{p}$를 얻음

$$\hat{p}=MLP(f)$$

Vim Block

1. 정규화

• 입력 토큰 시퀀스 $T_{l-1}$ 는 먼저 정규화 레이어를 통해 정규화됨

T'_{l-1} : (B, M, D) ← Norm(T_{l-1})

2. 선형 투영

• 정규화된 시퀀스는 두 개의 벡터 $x$ 와 $z$로 선형 투영

x : (B, M, E) ← Linear_x(T'_{l-1})
z : (B, M, E) ← Linear_z(T'_{l-1})

3. 전방향 및 역방향 처리

• $x$는 전방향과 역방향으로 처리됨
• 양방향 시퀀스 모델링을 통해 데이터의 모든 방향에서 정보를 활용

for o in {forward, backward} do

4. 1차원 컨볼루션

• 각 방향에서 $x$에 대해 1차원 컨볼루션을 적용해 $x’_o$를 얻음

x'_{o} : (B, M, E) ← SiLU(Conv1d_o(x))

5. 선형 투영

• $x’_o$ 는 선형 투영되어 $B_o, C_o, \Delta_o$

B_o : (B, M, N) ← Linear_B_o(x'_{o})
C_o : (B, M, N) ← Linear_C_o(x'_{o})

6. 변환

• $\Delta_o$는 $B_o, C_o$ 를 변환하는 데 사용
• 상태 공간 모델(SSM)의 파라미터를 이산화(discretize)하여 처리 가능

A_o : (B, M, E, N) ← Δ_o @ Parameter_A_o
B_o : (B, M, E, N) ← Δ_o @ B_o

7. SSM 계산 & 게이팅 및 합산

• SSM을 통해 전방향 출력 $y_{forward}$와 역방향 출력 $y_{backward}$를 계산

y_o : (B, M, E) ← SSM(A_o, B_o, C_o)(x'_{o})

• $z$에 의해 게이팅

y'_{forward} : (B, M, E) ← y_{forward} ⊙ SiLU(z)
y'_{backward} : (B, M, E) ← y_{backward} ⊙ SiLU(z)

• 입력의 중요성을 조절하는 역할을 하며, 이 두 출력을 더하여 최종 출력 토큰 시퀀 $T_l$ 생성

T_l : (B, M, D) ← Linear_T(y'_{forward} + y'_{backward}) + T_{l-1}
Return: T_l

 

Vim 블록의 구조는 비전 작업에서의 공간 인식을 향상시키고, 양방향 시퀀스 처리를 통해 데이터를 효과적으로 분석하게 함