[Paper Review] Mamba - Linear Time Sequence Modeling with Selective State Spaces 1

나온지 벌써 1년도 넘었지만 최신 논문 리뷰를 안 한지가 백만년 정도 된 것 같아서 한 번 읽어보는 mamba... 그리 정확한 리뷰는 아닐 수 있습니다. 사실 번역에 가깝고 좀 더 이해해보면서 내용 수정하겠습니다.

 

1. Introduction

최근들어 Structured State Space Sequence Models(SSMs) 가 시퀀스 모델링 구조 분야에서 유망한 클래스로 등장했다. 전통적인 state space models에 영감을 받아 CNN과 RNN의 통합을 조율한다(interpreted).
Mamba에서는 selective state space model의 새로운 종류를 제안한다. 시퀄스 길이에 따라 선형적으로 확장하면서 transformer의 모델링 파워를 따라잡기 위해서 몇몇의 axes(여기에서는 기준? 축.)에 대한 기존 연구를 발전시킬 것이다.

Selection Mechanism

• 기존 모델의 주요 한계는 데이터를 입력 의존적 방식으로, 효율적으로 select 하는 능력을 가진다는 점이다.
• Mamba에서는 선택적 복사와 헤드 유도와 같은 중요한 합성 작업을 바탕으로 한 직관을 기반으로 입력에 기반한 SSM 파라미터를 매개변수화하여 간단한 선택 매커니즘을 설계하였다.
• 이는 모델이 관련이 적은 정보를 필터링하고 관련있는 정보를 기억하도록 한다.

Hardware-aware Algorithm

Architecture

2023년 발표된 SSM Acritecture과 Transformer의 MLP block을 single block으로 하여 선택적 state space 를 통합한 단순하고 통합적인 아키텍쳐 디자인(Mamba)를 설계하였다.
Selective SSMs는 fully recurrent models로 시퀀스에서 동작하는 일반적인 foundation models의 백본이 되기 적합한 key properties가 있다.

1. high quality
2. fast training and inference
3. long context

이 논문에서는 몇가지 타입의 modalities와 설정에 따른 pretraining quality와 domain-specific한 성능 평가로 Mamba의 일반적인 시퀀스 FM backbone으로서의 잠재력을 평가한다.

2. SSM(State Space Model)

Structured state space sequence models(S4)는 RNN, CNN, 그리고 전통적인 state space model과 연관이 있는 딥러닝을 위한 최신 시퀀스 모델의 한 종류이다. 이는1차원 함수 혹은 시퀀스를 mapping하는, $x(t) \in R \to y(t) \in R$ implicit latent state $h(t) \in R^{N}$ 부분적 연속 시스템이다.

S4 모델은 sequence-to-sequence transformation을 두 단계로 정의하는 4가지 파라미터($\Delta, A, B, C$) 로 정의된다.

이산화 Discretization

연속 파라미터($\Delta, A, B$)를 이진 파라미터($\bar{A}, \bar{B}$)로 변경한다. 이때 변환 방법은 zero-order hold(ZOH)에 따른다.

이산화는 해상도 불변과 같은 추가적인 속성을 부여할 수 있고, 모델이 적절하게 정규화되었는지 자동으로 확인하는 등 연속 시간 시스템과 긴밀한 관계가 있다. 또한 RNN의 gating 메커니즘과도 연관이 있다. 그러나 기계적 관점에서 이산화는 단순히 ssm의 forward pass의 그래프 계산의 첫 단계로 보일 수 있다. SSM의 대체 버전은 이산화 단계 대신 파라미터 ($\bar{A}, \bar{B}$)를 직접 대체할 수 있는데, 이는 더 쉬운 방법이 될 수도 있다.

Computation

파라미터를 이진 파라미터로 변환한 뒤에, 모델은 두 가지 방법으로 계산될 수 있는데 바로 Linear recurrence*와 *global convolution 이다.

대체로 모델은 효율적인 병렬 학습을 위해 convolutional mode를 사용하다가 효율적인 회귀추론을 위해 recurrent 방식을 사용한다.

선형시불변 Linear Time Invariance(LTI)

여기에서는 $(\Delta, A, B, C)$와 $(\bar{A}, \bar{B})$가 모든 time steps에 고정된다. 이러한 특성을 Linear Time Invariance라고 한다. 이는 recurrence와 convolution에 깊은 연관이 있다. 편하게 LTI SSMs를 어떤 선형 회귀나 합성곱과 동등하게 생각할 수 있고, LTI를 이런 종류의 모델 클래스에 대한 포괄적 용어라고 생각해도 되겠다.

근본적인 효율의 제약으로 인해 모든 SSMs 구조는 LTI이다. 그러나 중요한 시사점은 LTI 모델들이 특정한 타입의 데이터에 대한 모델링에 대한 근본적인 한계가 있다는 것이다. 그리고 우리의 기술적 기여는 LTI 제약을 삭제하고 효율적인 병목현상을 극복에 관련이 있다...

Structure and Dimensions

가장 인기있는 structure 형태는 diagonal이다. 이 경우 $A \in \mathbb{R}^{N \times N}, B \in \mathbb{R}^{N \times 1}, C \in \mathbb{R}^{1 \times N}$ 행렬들은 모두 N 숫자로 표현될 수 있다. Batch size B, Length L, channel D의 input sequence x를 계산하기 위해서는 SSM이 각 channel에 독립적으로 적용된다. 이 경우에는 전체 hidden state는 입력마다 $DN$ 차원을 가지고, 모든 sequence 길이에 따른 계산은 $O(BLDN)$ 시간과 메모리가 필요하다.

General State Space Models

• Markov Decision Process(MDP)
• Dynamic Casual Modeling(DCM)
• Kalman Filters
• Hidden Markov Models(HMM)
• Linear Dynamical Systems(LDS)
• Recurrent models at large deep learning

SSM Architectures

• Linear attention
• H3
• Hyena
• RetNet
• RWKV