Background
State Space Model (SSM)은 원래 칼만 필터 (Kalman Filter)에서 시작된 고전적인 모델로 긴 시퀀스에서의 장기 의존성을 잘 처리하면서도 병렬 학습이 가능하다는 장점을 가집니다. 최근 몇 년 간 S4, DSS, S4D와 같은 다양한 SSM 변형 모델들이 제안되면서 텍스트 및 시계열 등 긴 시퀀스를 다루는 분야에서 Transformer의 대안으로 부상하고 있습니다. 그 중에서도 Mamba라는 모델은 입력에 따라 동적으로 변화하는 (Time-varying) 파라미터, GPU 친화적인 구현을 통해 긴 시퀀스를 매우 효율적으로 처리하는 언어 모델로 주목받고 있습니다. 하지만, 지금까지는 이미지 및 비디오와 같은 시각 데이터 전체를 순수 SSM으로 처리하는 범용 비전 백본 (Pure-SSM Backbone)은 거의 연구되지 않았습니다.
비전 분야에서는 지금까지 Vision Transformer가 사실상의 표준 백본 역할을 해왔습니다. ViT는 입력 이미지를 패치 단위로 쪼개어 시퀀스로 보고 self-attention을 적용하여 각 패치가 데이터에 따라 달라지는 전역 컨텍스트를 볼 수 있게 됩니다. 이를 통해, CNN처럼 2D 구조에 강하게 묶이지 않고 모달리티에 얽매이지 않는 시퀀스 모델링이 가능하게 되어 멀티모달 모델에서 특히 유리하게 되었습니다.
하지만, self-attention은 토큰 개수에 대해서 $\mathcal{O}(N^{2})$만큼의 연산 및 메모리 복잡도를 가진다는 단점이 존재합니다. 이로 인해 해상도가 높아질수록 속도와 GPU 메모리가 폭발적으로 증가하게 되죠. 이로 인해, 실제 고해상도 비전 작업에서는 ViT를 그대로 사용하지는 않고 윈도우 어텐션 또는 패치 피라미드 등 여러가지 트릭을 적용하여 학습을 하게 됩니다.
그림 1에서는 대표적인 Vision Transformer 계열인 DeiT와 본 논문에서 제안하는 Mamba 기반 모델인 Vim의 성능을 보여주고 있습니다. 특히, 고해상도 ($1248 \times 1248$)에서 Vim이 DeiT에 비해 2.8배 더 빠르고 GPU 메모리는 86.8%나 덜 사용하지만 분류, 검출, 분할 성능은 모두 DeiT를 크게 뛰어넘는 것을 볼 수 있습니다.
그렇다면 본 논문에서는 어떻게 이런 놀라운 성능을 가지는 모델을 만들 수 있었을까요? 저자들은 이를 위해 질문을 이렇게 던집니다.
좋은 비전 모델을 배우는 데 꼭 Self-attention이 필수일까? 만약 SSM만으로도 비전 백본을 만들 수 있다면?
하지만, 이러한 방식으로 바로 Mamba를 비전 분야로 가져오게 되면 두 가지 문제가 있습니다.
- 단방향 모델링 (Unidirectional Modeling) 문제: 기본적으로 언어 모델링에서 파생한 Mamba는 왼쪽에서 오른쪽으로의 순서가 자연스럽게 정의되지만 이미지에서는 위,아래,왼쪽,오른쪽 모두가 동등하게 중요합니다. 따라서, 한 방향으로만 읽는 SSM은 공간적 컨텍스트를 온전히 잡기 어렵다는 문제가 존재합니다.
- 위치정보부족(Positional Awareness) 문제: 픽셀 및 패치는 위치에 민감하지만 SSM은 기본적으로 "순서"만 있고 2D 위치 정보를 따로 제공하지 않으면 공간 구조를 잘 이해하기 어렵습니다.
이러한 두 가지 핵심 문제를 해결하기 위해 본 논문에서는 양방향 (Bi-directional) SSM과 위치 임베딩을 도입합니다. 이 과정에서 self-attention 없이 SSM 만으로 구성되어 ViT 수준의 표현력을 가지면서도 시퀀스 길이에 대해 선형적인 복잡도를 가져 고해상도 이미지에도 효율적으로 활용할 수 있게 되었죠. 저자들은 이를 영상 분류, 객체 검출, 객체 분할 등과 같은 핵심적인 비전 작업에 적용하여 모두 높은 성능을 달성함을 보여주었습니다.
Method
1) Preliminaries
S4나 Mamba와 같은 SSM 계열 모델은 원래 연속 시간 (Continous System) 관점에서 출발합니다. $x(t) \in \mathbb{R}$과 $y(t) \in \mathbb{R}$를 각각 입출력 1차원 연속 신호라고 가정하겠습니다. 그리고 숨겨진 상태 $h(t) \in \mathbb{R}^{N}$이 있다고 가정하겠습니다. 그러면 시스템은 세 개의 행렬 파라미터를 가집니다. $A \in \mathbb{R}^{N \times N}$은 상태가 시간에 따라 어떻게 변하는 지 결정하는 상태 전이 행렬, $B \in \mathbb{R}^{N \times 1}$은 입력 $x(t)$를 상태공간으로 투영하는 입력 투영 행렬, 그리고 $C \in \mathbb{R}^{1 \times N}$은 숨은 상태 $h(t)$를 출력 $y(t)$로 바꾸어주는 출력 투영 행렬입니다. 연속 시간 시스템에서의 SSM의 기본식은 이를 기반으로 아래와 같이 쓸 수 있습니다.
$$\begin{cases} h^{'}(t) &= Ah(t) + Bx(t) \\ y(t) &= Ch(t) \end{cases}$$
즉, 입력 $x(t)$를 받아 상태 $h(t)$를 업데이트하고 그 상태를 다시 출력 $y(t)$로 읽어내는 시스템이라고 볼 수 있습니다. 하지만, 딥 러닝에서는 연속 시간이 아니라 시퀀스 $(x_{1}, \dots, x_{M})$의 형태로 다루기 때문에 이 시스템을 이산 시간 버전으로 바꾸어야합니다. 이를 위해 시간 스케일링 파라미터 $\Delta$를 두고 연속 행렬 $A$와 $B$를 이산 행렬 $\overline{A}$와 $\overline{B}$로 변환해야합니다. 본 논문에서는 가장 흔히 쓰이는 Zero-Order Hold (ZOH) 변환을 사용했다고 하네요. 변환식은 다음과 같습니다.
$$\begin{cases} \overline{A} &= \text{exp} (\Delta A) \\ \overline{B} &= (\Delta A)^{-1} (\text{exp} (\Delta A) - I) \Delta B \end{cases}$$
이 과정을 거치면 연속 시간 시스템에서의 SSM은 다음과 같은 이산 시간 시스템에서의 SSM으로 변환할 수 있습니다.
$$\begin{cases} h_{t} &= \overline{A}h_{t - 1} + \overline{B}x_{t} \\ y_{t} &= Ch_{t} \end{cases}$$
여기서 $t = 1, 2, \dots, M$은 시퀀스의 타임스텝 인덱스를 의미합니다. 마지막으로 본 논문에서는 이산 SSM이 전역 1D 컨볼루션 연산으로 다시 쓸 수 있다는 점을 보여줍니다. 길이 $M$인 입력 시퀀스 $\mathbf{x} = (x_{1}, \dots, x_{M})$에 대해서 커널 $\mathbf{K} = (C\overline{A}, C\overline{A}\overline{B}, C\overline{A}^{2}\overline{B}, \dots, C\overline{A}^{M - 1}\overline{B}) \in \mathbb{R}^{M}$을 정의합니다. 그러면 전체 출력 $\mathbf{y} = \mathbf{x} * \mathbf{K}$와 같이 쓸 수 있죠. 즉, SSM이라는 것은 3개의 행렬 $A, B, C$로 정의된 구조화된 행렬 커널 (Structured Convolutional Kernel)을 학습하는 것과 동일하다는 것을 보여주고 있습니다. 이때, 커널들이는 시퀀스 길이 $M$까지 늘어날 수 있기 때문에 아주 긴 범위의 의존성을 커버하는 컨볼루션을 효율적으로 구현한다고 볼 수 있습니다.
2) Vision Mamba
그림 2는 Vim의 전체적인 구조를 보여주고 있습니다. 사실 기본적인 단계는 ViT와 거의 동일합니다. 먼저 입력 이미지 $\mathbf{t} \in \mathbb{R}^{H \times W \times C}$ 를 $P \times P$ 크기의 패치로 잘라서 총 $J$개의 패치로 만든 뒤 각 패치를 길이 $P^{2} \cdot C$인 벡터로 펼칩니다. 이를 모아서 $x_{p} \in \mathbb{R}^{J \times (P^{2} \cdot C)}$ 형태의 패치 행렬로 만들어주는 것이죠. 다음으로 패치를 토큰 벡터로 투영하기 위해 학습 가능한 선형 변환 행렬 $\mathbf{W} \in \mathbb{R}^{(P^{2} \cdot C) \times D}$를 사용해 각 패치를 $D$ 차원의 토큰으로 만들어줍니다. 마지막으로 클래스 토큰과 위치 임베딩을 추가하여 다음과 같이 초기 토큰들을 얻습니다.
$$\mathbf{T}_{0} = [\mathbf{t}_{cls}; \mathbf{t}^{1}_{p}\mathbf{W}, \dots, \mathbf{t}^{J}_{p}\mathbf{W}] + \mathbf{E}_{pos}$$
이렇게 만든 토큰 시퀀스 $\mathbf{T}_{0}$를 $L$개의 Vim 블록에 통과시킵니다. $l$-번째 레이어에서는 $\mathbf{T}_{l} = \text{Vim}(\mathbf{T}_{l - 1}) + \mathbf{T}_{l - 1}$ 과 같이 잔차 연결도 함께 적용합니다. 여기서 $\text{Vim}(\cdot)$이 본 논문에서 핵심적으로 제안하는 비전용 Mamba 블록입니다.
마지막 레이어에서의 출력 시퀀스 $\mathbf{T}_{L}$에서 클래스 토큰 위치 $\mathbf{T}^{0}_{L}$만 선택합니다. 이를 정규화하여 MLP 헤드를 거치면 클래스 확률을 얻게 되는 것이죠.
$$\begin{cases} \mathbf{f} &= \textbf{Norm}(\mathbf{T}^{0}_{L}) \\ \hat{p} &= \textbf{MLP} (\mathbf{f}) \end{cases}$$
3) Vim Block
기존의 Mamba 블록은 1D 시퀀스 데이터를 위해서 제안되었습니다. 이로 인해, 바로 Mamba를 비전 작업에 활용하게 되면 spatial awareness가 중요한 비전 분야에 단점으로 적용할 수 있죠. 이러한 문제를 해결하기 위해 본 논문에서는 기존의 Mamba 블록을 변형한 Vision Mamba (Vim) 블록을 제안합니다. 이는 그림 2의 오른쪽 그림에서 더 자세히 볼 수 있습니다.
또한, 위 알고리즘 1에서도 Vim 블록의 전체적인 프로세스를 요약해서 보여주고 있습니다. 입력 토큰 $\mathbf{T}_{l - 1}$은 정규화 계층에 의해서 정규화됩니다. 다음으로 정규화된 토큰들은 차원이 $E$인 $\mathbf{x}$와 $\mathbf{z}$로 선형 투영됩니다. 다음 단계에서는 $\mathbf{x}$를 처리하게 되는 데 여기서 중요한 점은 forward-backward로 두 가지 방향에서 처리를 한다는 점입니다. 각각의 처리 과정에 있어서 1D 합성곱 연산을 $\mathbf{x}$에 적용하여 $\mathbf{x}^{'}_{o}$를 얻습니다. 여기서 $o \in \{ \text{forward}, \text{backward} \}$입니다. 다음으로 $\mathbf{x}^{'}_{o}$에 한번 더 선형 투영을 적용하여 $B_{o}, C_{o}, \Delta_{o}$를 얻습니다. 이때, $\Delta_{o}$는 이산화된 행렬인 $\overline{A}_{o}$와 $\overline{B}_{o}$를 얻는데 사용되게 됩니다. 마지막으로 SSM을 통해 $\mathbf{y}_{o}$를 얻고 처음에 얻은 $\mathbf{z}$를 이용해서 gating해준 뒤 하나로 fusion하고 다음 Vim 블록으로 넘어가게 됩니다.
Experiment Results
1) Image Classification
2) Object Detection
3) Semantic Segmentation
4) Efficiency Analysis
5) Ablation Study
5-1) Bidirectional Design and Classification Design
5-2) Hierarchical Design
