Albert Gu · Tri Dao  ·  Carnegie Mellon University / Princeton University  ·  2023

Mamba: Linear-Time Sequence Modeling with Selective State Spaces

이 논문은 시퀀스 모델링의 핵심 딜레마 — 표현력 vs. 효율성 — 를 "선택적 압축(selective compression)"의 관점에서 본다. Transformer의 attention은 모든 맥락을 명시적으로 기억하기에 강력하지만 이차(quadratic) 비용을 치른다; Mamba는 입력에 따라 무엇을 기억하고 무엇을 잊을지를 동적으로 결정함으로써, 선형 시간 복잡도를 유지하면서 Transformer 수준의 성능에 처음으로 도달한다.

arXiv 2312.00752 →
state space model selective SSM linear-time sequence modeling Mamba hardware-aware algorithm foundation model efficient RNN

문제의 배경 — 기존 연구의 한계

현대 딥러닝의 파운데이션 모델(foundation model)은 언어, 이미지, 오디오, 유전체 등 다양한 시퀀스 데이터를 처리한다. 이 모델들의 거의 모든 백본(backbone)은 2017년에 등장한 Transformer와 그 핵심인 어텐션(attention) 레이어이다.

"The efficacy of self-attention is attributed to its ability to route information densely within a context window, allowing it to model complex data. However, this property brings fundamental drawbacks: an inability to model anything outside of a finite window, and quadratic scaling with respect to the window length."

셀프-어텐션의 효능은 컨텍스트 윈도우 안에서 정보를 밀도 있게 라우팅하는 능력에 있지만, 이 특성은 근본적인 단점을 수반한다: 유한한 윈도우 밖의 정보를 모델링할 수 없고, 윈도우 길이에 대해 이차(quadratic) 비용이 발생한다.

Transformer의 구조적 한계

한계 1 이차 복잡도 (quadratic complexity): 셀프-어텐션은 시퀀스 길이 L에 대해 O(L²) 시간·메모리가 필요하다. 길이 100만짜리 DNA 시퀀스나 오디오 파형을 처리하는 것은 Transformer에게 사실상 불가능하다.
한계 2 추론 시 KV 캐시 문제: 자기회귀(autoregressive) 추론 시 이전 모든 토큰의 key/value 쌍을 저장해야 하므로, 배치 크기가 커질수록 메모리가 선형적으로 증가한다. 이는 추론 처리량(throughput)을 근본적으로 제한한다.

대안적 시도들: RNN 계열과 SSM

이런 한계를 극복하려는 시도로 RNN구조적 상태 공간 모델(structured SSM)이 개발됐다. 특히 S4(Gu et al. 2022), H3(Dao et al. 2023), Hyena(Poli et al. 2023), RetNet(Sun et al. 2023), RWKV(Peng et al. 2023) 등이 시퀀스 길이에 선형(linear)이거나 준선형(quasi-linear) 복잡도를 달성했다.

공통 한계 선형 시간 불변성(LTI: Linear Time Invariance)의 함정: 이 모든 모델은 SSM 파라미터 (Δ, A, B, C)가 시간에 걸쳐 상수다. 즉, 입력 내용(content)에 따라 어떤 정보를 기억하고 무시할지를 동적으로 결정할 수 없다. 모든 입력에 동일한 필터가 적용된다.

이 한계는 선택적 복사 과제(Selective Copying task)로 극명하게 드러난다. 입력과 출력 사이의 간격이 고정된 일반 복사 과제는 LTI 모델(합성곱 기반)이 쉽게 해결할 수 있다 — 정확히 맞는 길이의 합성곱 커널을 만들면 된다. 그러나 간격이 무작위로 변하면, LTI 모델은 실패한다. 정적인 합성곱 커널은 변하는 간격을 처리할 수 없기 때문이다.

"We identify that a key weakness of such models is their inability to perform content-based reasoning."

우리는 이러한 모델들의 핵심 약점이 내용 기반 추론(content-based reasoning)을 수행할 수 없다는 점임을 확인했다.

결론: 기존 효율적 시퀀스 모델들은 효율성은 얻었지만 내용 인식 능력(content awareness)을 잃었다. 이것이 Mamba가 해결하려는 핵심 문제다.

이 논문의 선택 — 핵심 아이디어와 트레이드오프

핵심 통찰: 시퀀스 모델링은 곧 "압축"이다

"We argue that a fundamental problem of sequence modeling is compressing context into a smaller state."

시퀀스 모델링의 근본적인 문제는 맥락(context)을 더 작은 상태(state)로 압축하는 것이라 주장한다.

이 관점에서 보면 기존 모델들의 위치가 명확해진다:

핵심 아이디어 선택 메커니즘 (Selection Mechanism): SSM 파라미터 Δ, B, C를 입력 x의 함수로 만들어, 모델이 현재 토큰에 따라 상태를 어떻게 업데이트할지 동적으로 결정하게 한다. 이것이 LTI에서 time-varying SSM으로의 전환이다.

인터랙티브 비교: LTI SSM vs. Selective SSM (S6)

시퀀스의 각 위치를 클릭하면 해당 입력이 상태 업데이트에 미치는 영향 차이를 확인할 수 있습니다.

LTI SSM (S4)
파라미터 (A,B,C,Δ) = 상수 (입력 무관)
토큰을 클릭하세요
Selective SSM (S6 / Mamba)
Δ, B, C = s(x) — 입력에 따라 동적으로 결정됨
토큰을 클릭하세요

왜 이 방향을 선택했는가?

기존의 여러 접근법들이 효율성 문제를 해결하려 했다. linear attention, gated convolution 등이 있었지만, 이들은 content-based reasoning 능력을 포기했다. Mamba의 선택은 단순하지만 강력하다: "파라미터를 입력의 함수로 만들어라."

기여점 1 선택 메커니즘 (Section 3.2): Δ, B, C를 입력의 함수 s(x)로 설정. 이로써 모델은 각 토큰에 대해 "이것을 얼마나 중요하게 볼 것인가"를 동적으로 결정할 수 있다.
기여점 2 하드웨어 인식 알고리즘 (Section 3.3): 선택 메커니즘이 추가되면 합성곱(convolution)을 쓸 수 없어 naive한 구현은 느리다. GPU 메모리 계층(HBM vs SRAM)을 활용한 커널 융합(kernel fusion) + 병렬 스캔(parallel scan) + 재계산(recomputation)으로 이 문제를 해결한다.
기여점 3 단순화된 아키텍처 (Section 3.4): H3 블록과 MLP 블록을 하나의 균일한(homogeneous) Mamba 블록으로 통합. Attention도 MLP도 없는 순수한 SSM 아키텍처.

무엇을 포기했는가? (트레이드오프)

트레이드오프 학습 시 합성곱 모드 사용 불가: 기존 LTI SSM은 학습 시 FFT 기반의 글로벌 합성곱을 사용해 O(L log L)의 효율적인 병렬 학습이 가능했다. 선택 메커니즘은 파라미터가 시간에 따라 변하므로 이 합성곱 표현을 쓸 수 없다. 대신 병렬 스캔(parallel scan)으로 O(BL²DN)에서 O(BLDN)으로 줄인다.

방법론

4a. 전체 아키텍처: Mamba 블록

Mamba 아키텍처는 H3 블록과 MLP 블록을 하나의 단순화된 블록으로 통합한다. 이 블록을 균일하게(homogeneously) 반복하여 전체 모델을 구성한다.

전체 Mamba 블록 구조 (블록 클릭 → 상세 설명)

입력 x (B,L,D) Norm RMSNorm Linear(ED) SSM 경로 Linear(ED) 게이트 경로 Conv1d causal SiLU 활성화 SSM (S6 선택적) 핵심 기여 × SSM 경로 게이트 경로

각 블록을 클릭하면 자세한 설명이 나타납니다. 오렌지 강조 = 이 논문의 핵심 기여

텐서 Shape 변환 흐름

단계 Shape 연산 설명
입력(B, L, D)배치 B, 시퀀스 길이 L, 채널 D (예: D=768)
RMSNorm(B, L, D)정규화 (shape 유지)
Linear (SSM 경로)(B, L, ED)E=2배 확장, 파라미터 2ED²
Conv1d → SiLU(B, L, ED)로컬 특징 추출 + 비선형화
SSM (S6)(B, L, ED)선택적 상태 공간 처리. 내부 히든 상태: (B, L, ED, N) but SRAM에만 존재
× (게이팅)(B, L, ED)SSM 출력 × SiLU(gate) elementwise
Linear (출력)(B, L, D)원래 차원으로 투영, 파라미터 ED²
+ 잔차 연결(B, L, D)입력 x를 더함 (residual connection)

[보충] 파라미터 수 계산: 하나의 Mamba 블록에서 선형 투영이 3ED² (SSM 경로 2ED² + 출력 ED²), SSM 파라미터는 이에 비해 훨씬 적다. E=2로 고정하고 두 개의 블록을 쌓으면 Transformer의 MHA+MLP 블록과 유사한 12D² 파라미터를 갖는다.

4b. 핵심 수식 — SSM부터 선택적 SSM까지

연속 시간 상태 공간 모델 (Continuous-time SSM)

\[\begin{aligned} h'(t) &= Ah(t) + Bx(t) \quad &(1a) \\ y(t) &= Ch(t) \quad &(1b) \end{aligned}\]
변수의미비고
x(t) ∈ ℝ1차원 입력 시퀀스시간 t에서의 입력값
h(t) ∈ ℝᴺN차원 잠재 상태(latent state)현재까지 압축된 맥락
y(t) ∈ ℝ1차원 출력상태를 읽어낸 값
A ∈ ℝᴺˣᴺ상태 전이 행렬상태가 어떻게 진화하는지 결정. 대각(diagonal) 구조로 단순화
B ∈ ℝᴺˣ¹입력 투영 행렬입력이 상태에 어떻게 영향을 미치는지
C ∈ ℝ¹ˣᴺ출력 투영 행렬상태에서 출력을 읽어내는 방법

직관적 해설

이 식은 아날로그 회로나 물리 시스템(예: 스프링-댐퍼)의 수학적 표현과 동일하다. h(t)는 "현재 시스템의 기억"이고, A는 "기억이 스스로 얼마나 빨리 사라지는가"를 결정하며, B는 "새로운 입력이 기억에 얼마나 강하게 각인되는가"를 결정한다. C는 현재 기억에서 출력을 읽어내는 렌즈 역할을 한다. 핵심: 연속 시간이라 딥러닝에 직접 쓰기 어렵다 → 이산화가 필요하다.

수학적 유도

이는 1차 선형 상미분방정식(ODE) 시스템이다. (1a)의 해는 h(t) = e^(At)h(0) + ∫₀ᵗ e^(A(t-τ))Bx(τ)dτ. 이를 디지털 시스템에서 사용하려면 시간축을 이산화(discretize)해야 한다.

이산화 (Discretization) — ZOH 규칙

\[\bar{A} = \exp(\Delta A), \quad \bar{B} = (\Delta A)^{-1}(\exp(\Delta A) - I) \cdot \Delta B \quad (4)\]
변수의미비고
Δ타임스텝 크기 (step size)연속 시스템을 이산 시스템으로 변환하는 샘플링 간격. Mamba에서는 입력의 함수!
Ā이산화된 상태 전이 행렬exp(ΔA) 계산. Δ가 크면 상태가 빠르게 리셋됨
이산화된 입력 행렬복잡해 보이지만 Δ→0이면 ΔB에 수렴 (Euler 근사)

직관적 해설

ZOH(Zero-Order Hold) 이산화는 "구간 [t, t+Δ] 동안 입력이 x(t)로 일정하다고 가정"하는 방법이다. Δ(타임스텝)가 작으면 Ā ≈ I+ΔA (상태가 크게 변하지 않음), Δ가 크면 Ā ≈ 0 (이전 상태를 거의 무시하고 새 입력으로 리셋). 바로 이 Δ를 입력에 따라 동적으로 조절하는 것이 Mamba의 선택 메커니즘의 핵심이다!

수학적 유도

연속 시간 해 h(t+Δ) = e^(AΔ)h(t) + e^(AΔ)∫₀^Δ e^(-Aτ)Bx(t)dτ 에서, 입력이 구간 내 일정하다고 가정하면(ZOH): Ā = e^(AΔ), B̄ = A⁻¹(e^(AΔ)-I)B. 이를 계산하면 식 (4)가 도출된다.

이산 재귀 (Discrete Recurrence)

\[h_t = \bar{A} h_{t-1} + \bar{B} x_t \quad (2a), \qquad y_t = C h_t \quad (2b)\]
변수의미비고
h_t ∈ ℝᴺ타임스텝 t에서의 히든 상태이전 상태 h_{t-1}과 현재 입력 x_t를 결합
x_t ∈ ℝ타임스텝 t에서의 입력LTI에서는 (Ā,B̄)가 모든 t에서 동일 → 선택 불가

직관적 해설

RNN 업데이트 식과 완전히 동일하다! h_t = Āh_{t-1} + B̄x_t. 차이점은 Ā와 B̄의 값이다. LTI SSM에서는 Ā, B̄가 모든 타임스텝에서 동일하다. Mamba에서는 Ā_t = f(Δ_t, A), B̄_t = f(Δ_t, A, B)이고, Δ_t, B_t, C_t가 입력 x_t에 따라 변한다.

수학적 유도

이는 식 (4)의 이산화 공식을 (1)에 직접 적용하여 얻는다. 재귀 모드는 추론(inference) 시 효율적: O(N) 상태 업데이트만 필요. 합성곱 모드(학습 시): K = (CB, CĀB, ..., CĀᵏB, ...), y = x * K 로 표현되어 FFT를 이용한 병렬화가 가능하지만 — 선택 메커니즘으로 파라미터가 시간에 따라 변하면 이 합성곱 표현을 쓸 수 없다!

선택 메커니즘 (Selection Mechanism) — S4 vs S6

선택 메커니즘의 핵심은 단순하다: Δ, B, C를 입력 x의 함수로 만든다.

Algorithm 1: SSM (S4) — LTI
입력: x : (B, L, D)
출력: y : (B, L, D)

A : (D, N) ← 학습 파라미터
B : (D, N) ← 학습 파라미터  ← 상수!
C : (D, N) ← 학습 파라미터  ← 상수!
Δ : (D) ← τ_Δ(파라미터)    ← 상수!

Ā, B̄ : (D, N) ← discretize(Δ, A, B)
y ← SSM(Ā, B̄, C)(x)
  ⟶ 재귀 또는 합성곱으로 계산 가능
Algorithm 2: SSM+Selection (S6) — Mamba ★
입력: x : (B, L, D)
출력: y : (B, L, D)

A : (D, N) ← 학습 파라미터
B : (B, L, N) ← s_B(x)  ← 입력 함수!
C : (B, L, N) ← s_C(x)  ← 입력 함수!
Δ : (B, L, D) ← τ_Δ(파라미터+s_Δ(x)) ← 입력 함수!

Ā, B̄ : (B, L, D, N) ← discretize(Δ, A, B)
y ← SSM(Ā, B̄, C)(x)
  ⟶ 재귀(scan)만 가능 — 합성곱 불가!

구체적인 파라미터화: s_B(x) = Linear_N(x), s_C(x) = Linear_N(x), s_Δ(x) = Broadcast_D(Linear_1(x)), τ_Δ = softplus. Δ의 투영 차원이 1인 이유는 Theorem 1 (RNN 게이팅과의 연결)에서 설명된다.

Theorem 1 — RNN 게이팅과의 동치

\[g_t = \sigma(\text{Linear}(x_t)), \qquad h_t = (1-g_t)h_{t-1} + g_t x_t \quad (5)\]
변수의미비고
g_t게이트 값 (0~1)0이면 이전 상태 유지, 1이면 현재 입력으로 완전 교체
σ시그모이드 함수ZOH 이산화로부터 자연스럽게 도출됨

직관적 해설

이 식은 LSTM/GRU의 게이팅 메커니즘과 동일한 형태다! N=1, A=-1, B=1, s_Δ=Linear(x), τ_Δ=softplus 조건 하에 S6 알고리즘이 바로 이 형태로 단순화된다. 즉 Mamba의 선택 메커니즘은 RNN 게이팅의 일반화다. g_t가 크면(→1): 현재 입력을 강하게 기억, 이전 상태 리셋. g_t가 작으면(→0): 현재 입력 무시, 이전 상태 유지.

수학적 유도

A=-1, B=1일 때 ZOH 이산화: Ā_t = exp(-Δ_t) = σ(-Linear(x_t)) = 1-σ(Linear(x_t)). B̄_t = 1-Ā_t = σ(Linear(x_t)). 따라서 h_t = Ā_t·h_{t-1} + B̄_t·x_t = (1-g_t)h_{t-1} + g_t·x_t (g_t = σ(Linear(x_t))). □

4c. 하드웨어 인식 알고리즘 — GPU 메모리 계층 활용

선택 메커니즘을 추가하면 합성곱을 쓸 수 없어, 단순(naive) 재귀 구현은 느리다. Mamba는 세 가지 기법으로 이를 극복한다.

문제: Naive 구현의 메모리 병목

병목 Naive하게 구현하면 이산화된 상태 h의 전체 크기 (B, L, D, N)을 GPU HBM(고대역폭 메모리)에 저장해야 한다. N≈16이면 입력/출력보다 N배 더 큰 메모리가 필요하다. 또한 대부분의 연산은 memory-bandwidth bounded다 — 계산보다 메모리 읽기/쓰기가 병목.

하드웨어 인식 알고리즘 단계별 설명

스텝 1 / 4
GPU HBM (느린 메모리, ~900GB/s) Δ, A, B, C 출력 y (B,L,D) GPU SRAM (빠른 메모리, ~19TB/s, 작음) 이산화 (Δ→Ā,B̄) 병렬 스캔 (h 계산) C×h → y 상태 h (B,L,D,N) HBM에 저장 안 함!

효율성 비교 요약

방법 학습 복잡도 추론 상태 크기 내용 인식
Transformer (Attention)O(L²D)O(LD) KV 캐시✓ 완전
LTI SSM (S4)O(L log L · D)O(DN) 상태✗ 없음
Selective SSM (Mamba)O(L·DN)O(DN) 상태✓ 있음

4d. 구현 세부사항 — 실제로 학습시키는 방법

설정값 (하이퍼파라미터)

항목 비고
확장 팩터 E2내부 채널 수 = 2D
SSM 상태 차원 N16효율성 벤치마크에 사용된 기본값
A 행렬 구조대각(diagonal)계산 효율을 위해 N개의 수로 표현 가능
A 초기화 (실수)S4D-Real: −(n+1)HIPPO 이론 기반. n번째 원소 = -(n+1)
Δ 초기화τ_Δ⁻¹(Uniform[0.001, 0.1])bias 항에 적용
활성화 함수SiLU / SwishSwiGLU 변형과 동일한 선택
정규화 레이어LayerNorm (선택적)RetNet의 설계에서 동기. 블록 내부에 추가
언어 모델 학습률1e-3 (≤370M), 6e-4 (>370M)LLaMa 스타일 레시피 따름
배치 크기0.5M~2.5M 토큰모델 크기에 따라 조정
학습 데이터The Pile (300B 토큰)GPT-NeoX-20B 토크나이저
하드웨어A100 GPU 8대1.3B 모델: 약 7일 학습 (추정)
공개 코드github.com/state-spaces/mamba모델 코드 + 사전학습 체크포인트 공개

[평가] 논문은 복잡도 O(BLDN)으로 표기하지만, 실제로는 병렬 스캔의 추가 오버헤드와 kernel fusion으로 인해 이론적 FLOPs와 실제 속도 사이에 상수 인자 차이가 있다. 논문은 A100에서 표준 스캔 구현 대비 20-40배 속도 향상을 보고한다.

결과 — 실험으로 검증된 성능

추론 처리량 향상
Transformer 동일 크기 대비
동등 성능 모델 크기
Mamba-3B ≈ Transformer-6B
1M
컨텍스트 길이
DNA/오디오 모델링에서 달성
FID 감소 (음성 생성)
SC09 데이터셋에서

5.1 합성 과제 (Synthetic Tasks)

Selective Copying — LTI의 실패와 S6의 성공

모델 Selective Copying 정확도 비고
S4 (LTI)실패 (18%)정적 커널로는 변하는 간격 처리 불가
H3 (LTI + gate)부분 성공게이팅은 시퀀스 축 상호작용 없음
Mamba (S6) Ours100%내용 인식으로 완벽 해결

5.2 언어 모델링 (Language Modeling)

"Mamba is the first attention-free model to match the performance of a very strong Transformer recipe (Transformer++) that has now become standard, particularly as the sequence length grows."

Mamba는 현재 표준이 된 매우 강력한 Transformer 레시피(Transformer++)의 성능에 필적하는 최초의 attention-free 모델이다.

제로샷 다운스트림 평가 (Table 3 요약)

모델 파라미터 Pile ppl ↓ LAMBADA ppl ↓ HellaSwag ↑ 평균 ↑
Pythia-160M160M29.6438.1030.240.6
Mamba-130M Ours130M10.5616.0735.344.7
Pythia-1.4B1.4B7.516.0852.155.2
Mamba-1.4B Ours1.4B6.805.0459.159.7
Pythia-2.8B2.8B6.735.0459.359.1
Mamba-2.8B Ours2.8B6.224.2366.163.3

5.3 속도 및 메모리 벤치마크

논문에 따르면 Mamba는 동일 크기 Transformer 대비 4-5배 높은 추론 처리량을 달성한다. 예: Mamba-6.9B (미학습)는 Transformer-1.3B보다 높은 처리량. KV 캐시가 필요 없기 때문에 훨씬 큰 배치 크기를 사용할 수 있기 때문이다.

5.4 에블레이션 (Ablation Studies)

"Replacing any of these with a selective SSM (S6) significantly improves performance, validating the motivation of Section 3."

이들 중 어느 것이든 선택적 SSM(S6)으로 교체하면 성능이 크게 향상되어, 3장의 동기를 검증한다.

아키텍처 (블록) 내부 레이어 Perplexity ↓
H3S4 (복소수)~9.0 (추정)
H3S4 (실수)유사
MambaS6 (선택적)큰 개선

선택 메커니즘(S6)이 핵심 기여임을 에블레이션이 검증. LTI 모델 간 차이보다 LTI→Selective 전환의 효과가 훨씬 크다.

논문 원본 Figure / Table

논문에 수록된 모든 주요 figure/table을 원본 그대로 보존합니다. 아래 설명을 참조하여 논문의 해당 위치를 찾을 수 있습니다.

Figure 1 — Overview & SSM equations
페이지 3 (§2, §3.3): Figure 1(상단)과 핵심 SSM 수식(1)~(4). Figure 1은 선택적 SSM의 전체 개요를 보여준다 — 각 채널이 잠재 상태 h를 통해 x를 y로 변환하며, 선택 메커니즘이 Δ를 입력 의존적으로 만든다. GPU HBM(빠른 메모리)과 GPU SRAM(느린 메모리) 계층도 표시된다.

주목할 포인트: Figure 1에서 Δ_t가 "Selection Mechanism"을 통해 입력 x_t에 의존함을 화살표로 명시하고 있다. 이것이 LTI와의 핵심 차이. 수식 (1)-(4)는 연속 시간 시스템 → 이산 재귀 → 합성곱의 세 가지 동등한 표현을 보여준다.

Algorithm 1&2 + Figure 2 — Copying Tasks
페이지 4 (§3.1, §3.2): Algorithm 1(S4)과 Algorithm 2(S6) 비교 및 Figure 2(Copying vs Selective Copying tasks). 알고리즘 비교에서 S4의 (D,N) 상수 파라미터가 S6에서 (B,L,N) 형태의 입력 의존적 파라미터로 바뀌는 것을 명확히 볼 수 있다.

주목할 포인트: Figure 2에서 일반 Copying task(일정한 간격)는 LTI 모델이 쉽게 해결하지만, Selective Copying(무작위 간격)는 내용 인식이 필요하므로 LTI 모델이 실패한다. Induction Heads task는 LLM의 in-context learning 능력의 핵심으로, 맥락 기반 추론이 요구된다.

Figure 3 — Mamba Architecture
페이지 6 (§3.4): Figure 3 — H3 블록, Gated MLP, Mamba 블록 비교. H3와 MLP 블록이 어떻게 단순화된 Mamba 블록으로 통합되는지를 보여준다.

주목할 포인트: Mamba 블록은 H3에서 첫 번째 곱셈 게이트를 활성화 함수로 교체했고, MLP 블록에서 SSM을 메인 브랜치에 추가했다. 이 단순화가 균일한(homogeneous) 아키텍처를 가능하게 한다.

Figure 4 — Language Scaling Laws
페이지 8 (§4.2.1): Figure 4 — 언어 모델 스케일링 법칙(Chinchilla 프로토콜). 125M~1.3B 파라미터 범위에서 Mamba vs. Transformer++, RWKV, RetNet, Hyena, H3 비교.

주목할 포인트: Mamba는 처음으로 Transformer++ (현재 표준이 된 강력한 Transformer 레시피)의 퍼플렉시티 곡선을 따라잡는 attention-free 모델이다. 특히 시퀀스 길이가 길어질수록 그 격차가 줄어든다.

Figure 5 — DNA Scaling Laws
페이지 9 (§4.3): 제로샷 평가 Table 3(상단)과 Figure 5(DNA 스케일링 법칙, 하단). Figure 5 왼쪽: 모델 크기에 따른 스케일링 (HG38 데이터셋), 오른쪽: 시퀀스 길이에 따른 스케일링.

주목할 포인트: Figure 5 오른쪽에서 HyenaDNA는 시퀀스 길이가 증가할수록 perplexity가 나빠지는 반면, Mamba는 1M 토큰까지 꾸준히 개선된다. 이것이 LTI 모델의 근본적 한계(컨텍스트 필터링 불가)를 실증적으로 보여주는 핵심 결과다.

Pages 11-12 — Figures 6-7, Tables 4-5
페이지 11 (§4.3.3, §4.4): Figure 6(DNA 분류), Figure 7(오디오 프리트레이닝), Tables 4-5(SC09 음성 생성). DNA에서 최대 220=100만 길이 시퀀스에서 성능 향상, 음성 생성에서 FID 0.94 달성 (기존 최고 1.42 대비).
Pages 13-16 — Efficiency & Ablations
페이지 13 (§4.5, §4.6): Figure 8(효율성 벤치마크) — 왼쪽: 학습 속도(표준 구현 대비 40배), 오른쪽: 추론 처리량(Transformer 대비 5배). Tables 6-7(에블레이션: 아키텍처, 선택성, 초기화).

주목할 포인트: Figure 8 왼쪽에서 표준 PyTorch scan 구현 대비 ~40배, FlashAttention-2 대비 시퀀스 길이 2K 이상에서 더 빠른 것을 볼 수 있다. 이는 CUDA 커널 융합의 효과다.

한계점 & 트레이드오프

한계 1 학습 시 병렬화 효율: LTI SSM은 FFT 기반 글로벌 합성곱으로 매우 효율적인 병렬 학습이 가능했다. Mamba는 선택 메커니즘으로 인해 이 합성곱을 쓸 수 없고 병렬 스캔(parallel scan)을 사용한다. 병렬 스캔도 효율적이지만, GPU의 행렬 곱셈 유닛(tensor cores)을 최적으로 활용하기 어렵다.
연구 방향: Mamba-2 (structured matrix)   hardware-efficient SSM
한계 2 초대형 규모 검증 미흡: 1~10B 파라미터 범위까지만 검증되었다. 현재 주류인 70B, 100B+ 규모에서의 체계적인 비교가 없다.
연구 방향: large-scale Mamba
한계 3 복잡한 in-context learning: Mamba는 기본 제로샷 태스크에서 강하지만, 복잡한 few-shot in-context learning, instruction following에서 Transformer와의 체계적 비교가 부족하다.
연구 방향: recurrent model instruction tuning
한계 4 비전/멀티모달 미검증: 언어, DNA, 오디오에서 검증됐지만 비전(ViT 대체), 멀티모달에서는 별도의 후속 연구가 필요하다.
연구 방향: VMamba (비전)

트레이드오프 재정리

얻은 것 잃은 것
✓ 내용 인식(content-aware) 능력✗ FFT 기반 합성곱 학습 효율성
✓ 선형 추론 복잡도 (KV 캐시 불필요)✗ 합성곱 모드의 단순함
✓ 1M 길이까지 성능 향상✗ GPU tensor core 최적화 어려움
✓ 5x 높은 추론 처리량✗ 초대형 규모 검증 미흡

영향력 & 후속 연구

누구에게 도움이 되는가?

Mamba는 다음 상황에서 특히 유용하다: (1) 매우 긴 시퀀스 처리가 필요한 경우 (게놈, 오디오, 긴 문서), (2) 추론 처리량이 중요한 배포 환경, (3) Transformer의 이차 복잡도가 병목인 엣지 디바이스/실시간 시스템.

저자가 제안한 후속 과제

논문은 다음 방향을 명시적으로 언급한다: (1) Selective SSM의 A 파라미터도 선택적으로 만드는 실험 (현재 고정), (2) 복소수 vs 실수 파라미터화의 데이터 모달리티별 최적 선택, (3) 다양한 도메인(강화학습, 시계열 등)으로의 확장.

Mamba 이후 연구 흐름

Mamba(2023.12)는 SSM 연구의 전환점이 됐다. 이 논문 이후 VMamba(비전), Jamba(Mamba+Transformer 하이브리드), Mamba-2(행렬 구조 개선), RWKV-6 등이 등장했으며, 많은 연구가 선택적 SSM을 기반으로 한다.

이 분야를 더 탐색할 키워드

Q&A — 연구자의 고민과 독자의 질문

Q1. 왜 A를 입력 의존적으로 만들지 않았는가?
"We remark that while the A parameter could also be selective, it ultimately affects the model only through its interaction with Delta via A-bar = exp(DeltaA) (the discretization). Thus selectivity in Delta is enough to ensure selectivity in (A-bar, B-bar)."
A의 선택성은 이미 Delta를 통해 간접적으로 구현된다. exp(DeltaA)에서 Delta가 입력 의존적이면 A-bar도 입력 의존적이 된다. A를 상수로 유지하면 HIPPO 기반의 유용한 초기화도 그대로 쓸 수 있다.
Q2. 선택 메커니즘이 단순한 "게이팅"과 어떻게 다른가?
이 질문은 논문의 Appendix A에서 직접 다룬다. 핵심 구분: 게이팅(gating)이 일반적으로 의미하는 elementwise 곱셈은 시퀀스 축을 따른 상호작용이 없다. 예를 들어 GLU(y = sigmoid(Wx) 곱 x)는 기술적으로 "gated"이지만 사실 단순한 활성화 함수다. 반면 Mamba의 선택 메커니즘은 Delta를 통해 시퀀스 축을 따라 정보 전파 방식을 제어한다.
정의 선택(selection): 모델이 시퀀스 차원을 따라 입력을 선택하거나 무시하고, 데이터 간의 상호작용을 용이하게 하는 메커니즘. 단순 gate와 달리 sequential dynamics를 제어한다.
Q3. Transformer와 성능이 같다면 왜 Mamba를 써야 하는가?
추론 효율성이 결정적 차이다. 사전학습 품질이 비슷하더라도, 배포(inference) 상황에서는 Mamba가 큰 이점을 가진다: KV 캐시 불필요 -> 배치 크기를 Transformer보다 훨씬 크게 설정 가능 / 추론 처리량 4-5x 향상 / 시퀀스 길이 1M에서도 선형 메모리 사용.

[평가] 하지만 실제로는 Transformer가 5년 이상 최적화된 코드베이스, FlashAttention 같은 기법들을 갖추고 있어, Mamba의 이론적 이점이 실제로 항상 발휘되지는 않는다. 또한 Attention이 컨텍스트 내 정밀한 특정 정보 검색에 더 강력할 수 있다.
Q4. 병렬 스캔은 어떻게 재귀를 병렬화하는가?
직관적으로 보면 h_t = A-bar_t * h_{t-1} + B-bar_t * x_t는 이전 결과에 의존하므로 순서대로 계산해야 할 것 같다. 하지만 이 재귀는 결합 법칙(associativity)을 만족하는 연산이다. Blelloch(1990)의 병렬 prefix scan 알고리즘을 사용하면, O(L) 순차 단계를 O(log L) 병렬 단계로 줄일 수 있다. 예를 들어 8개 타임스텝을 처리할 때: 1단계에서 (h1,h2), (h3,h4), (h5,h6), (h7,h8)을 동시에, 2단계에서 (h1..h4), (h5..h8)을 동시에 계산하는 식이다. [보충] 이것이 Mamba가 GPU의 병렬성을 활용하면서도 순서를 지키는 비결이다.
Q5. 코드는 어디서 볼 수 있는가?
모델 코드와 사전학습 체크포인트가 공개되어 있다: 공식 구현: github.com/state-spaces/mamba | 설치: pip install mamba-ssm | Hugging Face: huggingface.co/state-spaces. CUDA 커널이 포함되어 있으므로 GPU 환경이 필요하다.
Q6. Mamba가 실제로 "기억"을 어떻게 관리하는지 직관적으로?
[보충] Delta(타임스텝)가 핵심 레버다. Delta_t가 크면(→무한): A-bar_t = exp(DeltaA) → 0, 즉 이전 상태를 거의 지운다. 현재 입력을 새로 "기억"하는 것. Delta_t가 작으면(→0): A-bar_t ≈ I, 이전 상태를 거의 그대로 유지한다. 현재 입력을 "무시"하는 것. 모델은 각 토큰에 대해 Delta_t를 학습하므로, 사실상 "이 토큰을 얼마나 중요하게 봐야 하는가"를 자동으로 학습한다. 이것이 RNN의 forget gate, input gate와 개념적으로 동일하지만, SSM 이론의 엄밀한 수식에서 자연스럽게 도출된다.