Mamba: Linear-Time Sequence Modeling with Selective State Spaces

Introduction

Mamba는 기존 Transformer 기반 시퀀스 모델의 한계, 특히 긴 시퀀스에 대한 비효율적인 계산 복잡도와 문맥 기반 선택 추론 능력 부족를 해결하기 위해 제안된 새로운 선택적 상태공간 모델(Selective State Space Model, SSM)이다. 기존 SSM 방식의 단점으로 입력에 따라 유동적으로 정보를 선택하거나 무시하지 못하는 문제가 있었는데 mamba는 이를 해결하기 위해 파라미터를 입력에 따라 동적으로 생성하는 selective 메커니즘을 도입하고, 기존의 시간 불변 LTI 구조를 가변적 구조로 확장한다. 다만 이 선택성으로 인해 convolution 연산을 사용할 수 없어 계산 효율성이 떨어질 우려가 있었지만, Mamba는 병렬 scan 알고리즘, 커널 퓨전, 재계산 기법 등을 통해 Transformer 수준의 학습 및 추론 속도, 메모리 효율성을 유지했습니다. 자세한 내용은 아래에서 다루겠다.

 

실험 결과는 언어, DNA, 오디오 다양한 task에서 진행되었고 transformer를 능가하거나 동급의 성능을 보인다고 한다.

 

끄적끄적

	mamba에서 제안하는 S4( Structured state space sequence models)는 두 단계의 시퀀스 변환을 통해 수행된다고 한다. 1a,1b) 연속 시스템의 상태/출력 방정식 2a, 2b) 이산화된 RNN-style recurrence 계산 3a) 필터 커널 K 구성: 입력이 시간차를 두고 출력에 미치는 영향 3b) 입력 시퀀스 x에 커널 K를 컨볼루션해서 전체 출력 계산 이 구조 덕분에 SSM 기반 모델(Mamba 포함)은 RNN처럼 recurrence를 사용하지 않고, 전체 시퀀스를 convolution-like 연산으로 병렬 처리할 수 있다. 기존의 RNN보다 학습이 빠르고 길이도 잘 확장되는 이유는 여기에 있다.
	Discretization) - 이산화란 연속 시간에서 정의된 시스템의 파라미터 (Δ, A, B)를 이산 시간 시스템에서 사용할 수 있도록 변환하는 과정. 이때 ZOH방식을 사용한다고 하고, ZOH는 연속 시스템의 입력을 각 샘플 구간에서 일정하다고 가정하고, 그것에 대해 해석해 얻은 이산화 해. - 이산화를 하면 연속 시스템 이론에서 온 수학적 특성을 그대로 유지하기 때문에 resolution invariance(해상도 불변성), auto normalization(자동 정규화)과 같은 특성을 가질 수 있다고 함. - 이산화된 상태 업데이트 식은 (2a) 게이팅 구조 없는 단순 RNN 형태와 유사하다고 한다. Computation) - ZOH 이산화를 통해 연속 파라미터를 이산 파라미터로 변환한 후, 모델 계산은 2가지 방식 중 하나로 수행될 수 있다고 한다. 1. global convolution - 학습 시에는 global convolution 방법을 사용한다. 왜냐하면 학습은 전체 시퀀스를 한 번에 병렬로 처리하는 것이 효율적이기 때문. 2. linear recurrence - 추론 시에는 linear recurrence 방식을 사용. 실시간 예측(inference)에서는 하나의 입력이 들어올 때마다 순차적으로 계산해야 하므로, recurrence 구조가 효율적. Lineaer Time Invariance(LTI) - LTI는 '모델의 동역학이 시간에 따라 일정하다' 라는 의미인데, 즉 같은 입력이 들어오면 언제든 같은 방식으로 처리되는 시스템이라는 뜻. - 기존의 ssm방식들은 효율적인 병렬 계산과 간결한 구조 때문에 모두 lti 구조를 따랐다. - 하지만 LTI는 특정 유형의 데이터를 잘 모델링하지 못한다는 단점이 있다. - 그래서 Mamba는 LTI 제약을 제거하면서도 효율성 유지하는 전략을 제안한다고 한다.
	3. Selective State Space Models - 5개의 섹션으로 설명한다고 한다. 3-1) : selection mechanism의 필요성 3-2) : 상태 공간 모델을 통합하는 방법 3-3) : 하드웨어 친화적 알고리즘 설계 3-4) : conv,mlp 없는 단순한 구조 3-5) : 추가적인 논의
	3-1) selection mechanism의 필요성 - 저자는 기존 시퀀스 모델들이 완벽하지 않으며, 이를 보완하기 위해 Mamba가 왜 필요한지를 설명하고 있다. 시퀀스 모델의 핵심 과제는 컨텍스트(context)를 작은 상태(state)로 효과적으로 압축하는 것이라고 주장한다.   - 예를 들어, Transformer는 전체 컨텍스트를 모두 유지하기 때문에 성능은 뛰어나지만 계산량이 매우 크다.   -반면, RNN은 context를 작은 hidden state로 요약하지만, 압축이 너무 단순하고 유연성이 떨어진다. 이러한 문제를 직관적으로 보여주기 위해 Figure 2에서는 다음과 같은 합성 태스크 실험을 제시했는데, Copying Task: 단순 반복으로, 위치 기반 정보만으로 해결 가능. 기존 LTI 모델로도 충분. Selective Copying Task: 어떤 입력을 기억할지 내용 기반으로 선택해야 하므로, 단순 LTI 모델로는 어려움. Induction Heads Task: 앞선 문맥을 바탕으로 적절한 출력을 유도해야 하며, 문맥 기반 reasoning이 필요. 결론적으로, 이 실험들은 기존 모델들이 입력마다 다르게 반응하거나 필요한 정보만 선택적으로 기억하지 못하는 한계를 보여준다. 그래서 저자는 이 문제를 해결하려면 "selectivity(선택성)", 즉 입력에 따라 정보를 걸러내고 반영할 수 있는 메커니즘이 필요하다고 강조하며, 그것이 바로 Mamba의 핵심 설계 원칙이라고 주장한다.
	3-2) selection mechanism algorithm - 기존 SSM 모델(S4)은 파라미터(A, B, C, Δ)가 고정되어 있어 모든 시점에 동일하게 적용된다. 반면, Selection Mechanism이 적용된 SSM(Mamba)은 (B, C, Δ ) 파라미터를 입력 x에 따라 동적으로 생성되도록 설계되었다. 이로 인해 모델은 시퀀스의 각 시점마다 입력에 따라 다른 상태 업데이트를 수행할 수 있게 된다. 결과적으로, 기존에는 불가능했던 입력 기반 선택적 정보 반영(selectivity)이 가능해지며, Selective Copying, Induction Heads 같은 문맥 의존적 문제를 해결할 수 있는 구조적 유연성을 갖추게 된다.
	3-3) 하드웨어 친화적 설계 현대 하드웨어에서 효율적인 연산을 가능하게 하기 위해서 3가지 최적화 기법을 제안한다. a) kernel fusion(커널 결합) 여러 연산을 하나의 커널로 묶어 실행함으로써 메모리 IO를 최소화하고 처리 속도 향상 b) parallel scan(병렬 스캔) 원래 순차적으로 처리되는 recurrence를 병렬로 수행하여 연산 시간 단축 c) recomputation(재계산) 중간 상태를 저장하지 않고 역전파 시 다시 계산함으로써 메모리 사용량 절감
	3-4) Mamba Architecture - Mamba는 기존 SSM 구조(H3)를 바탕으로 설계되었고, H3는 Linear Attention과 MLP 블록을 번갈아 쌓는 구조로 구성되어있다. - 하지만 Mamba에서는 이 두 구성 요소를 하나의 통합 블록으로 결합하여 구조를 단순화하였다.  Residual Connection과 Layer Normalization을 포함하여 반복적으로 쌓아 전체 Mamba 아키텍처를 구성 하고, 활성화 함수로는 SiLU(Swish)를 사용하며, 이는 Gated MLP 구조에서 널리 사용되는 “SwiGLU” 형태로 작동한다. figure3는 요약하면 H3의 게이팅 → Swish로 대체 MLP에 SSM을 추가 정도가 될 거 같다.
	연산량이나 메모리 효율성의 측면에서 우수한 모습을 보인다.
마무리 개인생각 - pointmamba를 보다가 mamba가 뭔지 궁금해서 찾아보게 된 논문 - mamba가 어떤 구조로 동작하는지 알 수 있어서 좋았던 거 같다. 아쉬운 점은 아무래도 point cloud에 국한되어있지 않고 language,audio 등 다양한 task가 있다보니 아직은 완벽히 이해하진 못했고, 제어 이론이 들어가니 그 부분도 쉽지 않았던 거 같다. - pointmamba처럼 적용한 논문들을 보는게 좀 더 나을지도?