Mamba/SSM 기초

State Space Models (상태 공간 모델)은 콘텐츠 인식 선택적 필터링 (content-aware selective filtering)을 통해 선형 시간 시퀀스 모델링 (linear-time sequence modeling)을 제공하며, 긴 문맥 추론 (long-context inference)에서 Transformer (트랜스포머)에 도전장을 내밀고 있습니다.

이것이 중요한 이유

State Space Models (SSMs, 상태 공간 모델)은 긴 시퀀스 모델링 (long-sequence modeling)을 위해 Transformer (트랜스포머)를 대체할 수 있는 원칙적인 대안을 제공합니다. 긴 문맥을 처리하는 프로덕션 시스템(예: 코드 생성, 유전체 분석)에서 Transformer (트랜스포머) 어텐션 (attention)의 이차 비용 (quadratic cost)은 병목 현상이 됩니다. Mamba는 상수 메모리 상태 (constant-memory state)로 선형 시간 추론 (linear-time inference)을 달성하여, 어텐션 기반 모델이 비용 문제로 사용하기 어려운 백만 토큰 문맥에서도 실행 가능하게 만듭니다.

핵심 아이디어

SSMs는 연속 시간 제어 이론 (continuous-time control theory)에서 유래되었습니다. 즉, 잠재 상태 (latent state)가 입력에 의해 시간에 따라 진화하며, 관측값 (observations)은 해당 상태의 선형 투영 (linear projections)입니다. Mamba의 핵심 혁신은 SSM 파라미터를 _입력 선택적 (input-selective)_으로 만든 것입니다. 모델은 어떤 정보가 상태로 들어오고 나갈지를 게이팅 (gating)하는 법을 학습하며, 이는 $O(n^2)$의 비용 없이 관련 토큰에 집중하는 어텐션 (attention)의 능력을 모방합니다.

기술적 세부 사항

연속 시간 SSM은 다음과 같이 정의됩니다:

x′(t) = Ax(t) + Bu(t), \quad y(t) = Cx(t) + Du(t)

여기서 $x(t) \in \mathbb{R}^N$은 잠재 상태 (latent state), $u(t)$는 입력 (input)이며, $A \in \mathbb{R}^{N \times N}$, $B \in \mathbb{R}^{N \times 1}$, $C \in \mathbb{R}^{1 \times N}$ 입니다. 단계 $\Delta$를 사용하는 zero-order hold 이산화 (discretization)를 사용하면:

\bar{A} = \exp(\Delta A), \quad \bar{B} = (\Delta A)^{-1}(\exp(\Delta A) - I) \cdot \Delta B

순환 업데이트 (recurrent update)는 다음과 같이 됩니다:

x_k = \bar{A}x_{k-1} + \bar{B}u_k, \quad y_k = Cx_k

**Mamba의 선택적 메커니즘 (selective mechanism)**은 $B$, $C$, 그리고 $\Delta$를 입력 의존적 (input-dependent)으로 만듭니다:

$B_k = \text{Linear}_B(x_k), \quad C_k = \text{Linear}C(x_k), \quad \Delta_k = \text{softplus}(\text{Linear}\Delta(x_k))$

병렬 스캔 (Parallel scan) 알고리즘은 훈련(training) 중에 이 재귀(recurrence)를 $O(n \log n)$ 시간 복잡도로 계산합니다. 추론 (Inference)은 고정된 상태 크기 $N$에 대해 토큰당 $O(1)$의 복잡도를 가지며, 시퀀스 길이에 관계없이 **상수 메모리 디코딩 (constant-memory decoding)**을 실현합니다.

작동 원리

1. 입력 투영 (Project input): 토큰 $u_k$를 확장된 차원 $D \cdot N$으로 매핑합니다.
2. 선택적 파라미터 생성 (Generate selective parameters): $u_k$로부터 입력 의존적 (input-dependent)인 $B_k, C_k, \Delta_k$를 계산합니다.
3. 이산화 (Discretize): $\Delta_k$를 사용하여 연속적 (continuous)인 $(A, B)$를 이산적 (discrete)인 $(\bar{A}_k, \bar{B}_k)$로 변환합니다.
4. 재귀적 스캔 (Recurrent scan): 은닉 상태 (hidden states) $x_k$를 계산하기 위해 병렬 스캔 (training 시) 또는 순차적 업데이트 (inference 시)를 적용합니다.
5. 출력 투영 (Output projection): $y_k = C_k x_k$를 계산한 다음, 게이팅 (gating, SiLU)을 통해 출력 차원으로 투영합니다.

핵심 통찰 (Key Insights)

• 선택성 (Selectivity)의 필수성: 비선택적 SSM (S4)은 인컨텍스트 검색 (in-context retrieval)을 수행할 수 없습니다. $B, C, \Delta$를 입력 의존적으로 만드는 것이 콘텐츠 인식 필터링 (content-aware filtering)을 가능하게 합니다.
• $A$의 대각 (Diagonal) + 저계수 (low-rank) 구조: $A$에 대한 대각 구조는 $O(n)$ 재귀를 가능하게 합니다. Mamba는 대각 $A$ 행렬만을 독점적으로 사용합니다.
• 하드웨어 인식 설계 (Hardware-aware design): 스캔 커널 (scan kernel)은 연산 제한 (compute-bound)이 아닌 I/O 제한 (IO-bound)을 받습니다. Mamba의 CUDA 커널은 메모리 읽기를 최소화하기 위해 이산화 (discretization), 스캔 (scan), 그리고 출력 투영 (output projection)을 융합 (fuse)합니다.
• 선형 디코딩 비용 (Linear decoding cost): 선형적으로 증가하는 KV 캐시 (KV-cache)와 달리, SSM 상태는 고정된 크기 $O(ND)$를 가지므로 생성 시 메모리가 상수로 유지됩니다.

출처 (Sources)

출처 (Sources)

• Gu, A. & Dao, T. "Mamba: Linear-Time Sequence Modeling with Selective State Spaces." arXiv:2312.00752 (2023). [https://arxiv.org/abs/2312.00752]
• Gu, A. & Dao, T. "Mamba-3: Improved Sequence Modeling using State Space Principles." arXiv:2603.15569 (2026). [https://arxiv.org/abs/2603.15569]