Mamba: 선택적 상태 공간을 통한 선형 시간 시퀀스 모델링

Digest: Mamba는 SSM 파라미터 B, C, Δ를 입력의 함수로 만드는 **선택 메커니즘(Selection Mechanism)**을 도입하여 기존 SSM의 시간 불변(LTI) 제약을 제거한 최초의 실용적 선형 시간 시퀀스 모델이다. 선택적 SSM은 합성곱 기반 병렬화가 불가능해지므로, GPU SRAM에서 이산화와 재귀를 완전히 수행하는 하드웨어 인식 병렬 스캔 알고리즘을 설계하여 표준 PyTorch 대비 40배 속도를 달성했다. 언어(Pile), DNA(HG38), 오디오(SC09) 도메인에서 동크기 Transformer를 능가하고 2배 크기 Transformer와 동등한 성능을 기록한 최초의 선형 시간 모델이다. Selective Copying에서 S4(18.3%) 대비 97.0%, Induction Heads에서 100만 토큰까지 99.8% 완벽 외삽을 달성하며 선택 메커니즘의 효과를 직접 증명했다. 추론 처리량은 Transformer 대비 5배로, KV 캐시 불필요가 배치 크기 확대로 이어진다.

📚 이 논문을 읽기 전에 — 사전 학습 로드맵

Mamba는 제어이론(SSM) + 신호처리(이산화/합성곱) + 딥러닝(RNN/Transformer) + GPU 시스템(FlashAttention) 의 4개 분야가 한 점에서 만나는 논문이라, 백지에서 바로 읽으면 거의 모든 문장이 막힌다. 다음 순서로 배경을 쌓는 것을 권장한다.

Tier 1 — 필수 (이거 없으면 본문 절반이 안 읽힘)

순서	자료	왜 먼저 읽나	예상 시간
1	Attention Is All You Need (Vaswani et al., 2017)	Mamba가 “이기려는 대상”이 누구인지 — O(L²)·KV 캐시·content-based addressing의 의미를 알아야 §1의 동기가 잡힌다	2-3h
2	RNN/LSTM/GRU 게이팅 직관 (Olah blog “Understanding LSTM Networks” 등)	Theorem 1이 “Δ = GRU의 망각 게이트”라고 주장하는 순간을 이해하려면 게이팅 RNN의 직관이 필수	1-2h
3	Linear State Space Models 입문 (제어이론 교과서 1챕터, 또는 Albert Gu의 “The Annotated S4” 블로그 §1)	$h^{'} (t) = A h (t) + B x (t)$ 가 무엇이며 왜 ODE인지, ZOH 이산화가 무엇인지 — 본문의 모든 수식이 여기서 출발	2h

Tier 2 — 강력 권장 (이게 있으면 본문이 직선으로 읽힘)

순서	자료	왜 읽나	예상 시간
4	HiPPO (Gu et al., NeurIPS 2020) — Recurrent Memory with Optimal Polynomial Projections	Mamba가 그대로 재사용하는 A 행렬의 출처. “왜 하필 이런 A?”의 답이 여기에 있음	3h (수식 위주)
5	S4 (Gu, Goel, Ré, ICLR 2022) — Efficiently Modeling Long Sequences with Structured State Spaces	Mamba의 직계 부모. Recurrent ↔ Convolutional 이중 표현, structured A (DPLR), Long Range Arena 결과 — Mamba가 무엇을 물려받고 무엇을 깨는지 분명해짐. S4_2021_StateSpace	4-5h
6	FlashAttention (Dao et al., 2022)	Mamba §3.3의 hardware-aware 알고리즘은 FlashAttention의 직접 후속작 (저자 Tri Dao 동일). SRAM/HBM tiling, kernel fusion, recomputation이 같은 철학	2-3h

Tier 3 — 맥락 보강 (시간 여유 있을 때)

순서	자료	왜 읽나
7	H3 (Fu et al., 2023) — Hungry Hungry Hippos	S4를 언어로 처음 끌고 온 시도. Mamba 블록이 H3 + Gated MLP를 단일 블록으로 합친 것임을 알 수 있음
8	Hyena (Poli et al., 2023)	같은 LTI 진영의 경쟁 — Mamba가 LTI를 깨야 했던 이유를 반대편에서 확인
9	RWKV / RetNet (2023)	Linear attention 계열의 다른 답. Mamba와의 설계 분기점 비교
10	Parallel Scan (Blelloch, 1990) — Prefix Sums and Their Applications	§3.3.2의 work-efficient parallel scan의 고전적 출처. 30분 분량 노트로 충분
11	Mamba-2 (Dao & Gu, 2024) — State Space Duality	Mamba 읽은 직후의 자연스러운 다음 단계. SSM ≡ 특수한 attention의 형태로 통합

빠른 트랙 (반나절 코스)

시간이 정말 없다면: Annotated S4 블로그 통독 (3h) → Mamba 본문 직행. S4의 LTI 한계만 체득해도 §3.1 “Selection이 왜 필요한가”가 즉시 와닿는다. HiPPO 디테일은 일단 블랙박스로 두고, “A는 과거를 최적으로 압축하는 학습된 행렬” 정도로만 이해해도 본문은 따라간다.

읽는 순서 결정 트리

graph TD
    START["Mamba를 읽고 싶다"] --> Q1{"Transformer를<br/>이미 잘 안다?"}
    Q1 -->|No| T1["1) AIAYN → 2) RNN 직관"]
    Q1 -->|Yes| Q2{"SSM/제어이론<br/>배경 있다?"}
    T1 --> Q2
    Q2 -->|No| T2["3) SSM 입문 → 5) S4 (Annotated)"]
    Q2 -->|Yes| Q3{"S4를 읽었다?"}
    T2 --> Q3
    Q3 -->|No| T3["5) S4 필독"]
    Q3 -->|Yes| Q4{"FlashAttention<br/>아이디어 안다?"}
    T3 --> Q4
    Q4 -->|No| T4["6) FlashAttention<br/>(§3 tiling만이라도)"]
    Q4 -->|Yes| GO["✅ Mamba 본문 진입 가능"]
    T4 --> GO

    style GO fill:#a8e6cf
    style START fill:#ffd3b6

섹션별 요약

Introduction

Transformer의 O(L²) 복잡도가 긴 시퀀스의 근본적 병목
기존 SSM(S4, H3, Hyena)의 **시간 불변성(LTI)**이 핵심 한계: 파라미터가 입력과 무관하게 고정
LTI 제약으로 selective copying, induction heads 같은 내용 기반 추론이 원천 불가
Mamba는 선택 메커니즘 + 하드웨어 인식 알고리즘으로 동시 해결

Methods

선택 메커니즘 (S6):

B = Linear_N(x) — 입력 의존적 입력 행렬
C = Linear_N(x) — 입력 의존적 출력 행렬
Δ = softplus(Parameter + Broadcast(Linear_1(x))) — 입력 의존적 타임스텝
Theorem 1: N=1, A=-1, B=1인 선택적 SSM은 GRU 게이트와 수학적으로 동치 → SSM과 RNN 통합

하드웨어 인식 알고리즘: 3가지 기법 결합

커널 융합: 이산화+재귀를 GPU SRAM에서 완전 수행
병렬 스캔: 순차 재귀를 work-efficient 알고리즘으로 병렬화
재계산: 역전파용 중간 상태 재계산(저장 불필요) → FlashAttention급 메모리 효율

Mamba 블록: 입력 투영(D→2ED) → 선택적 SSM 분기 + SiLU 게이트 분기 → 원소별 곱 → 출력 투영(ED→D). Attention 없는 단일 동차 아키텍처.

Results

모든 스케일(130M~2.8B)에서 동크기 Pythia(Transformer) 능가: PPL 7~10% 개선
Mamba-1.4B(PPL 6.80) ≈ Pythia-2.8B(PPL 6.73): “2배 크기 Transformer 품질”
합성 과제: Selective Copying 18.3%→97.0%, Induction Heads 1M 토큰 99.8%
DNA: 컨텍스트 증가 시 성능 단조 향상 (LTI 모델은 반대로 저하)
SC09 오디오: FID 0.67, 이전 최고 대비 28.7% 개선
추론: Transformer 대비 5배 처리량, 스캔 속도 40배

Discussion

선택성이 이산 모달리티(텍스트, DNA)에서 강력하나 연속 모달리티에서는 LTI 유리할 수 있음
3B까지의 평가 → 더 큰 스케일과 Transformer 생태계 기법(RLHF, 양자화) 적용 미검증
“선형 시간에서 진정으로 Transformer 품질을 달성한 최초의 모델”

Insights

Theorem 1의 의의: SSM 이산화 = RNN 게이팅의 수학적 기반. 수십 년간 분리된 SSM과 게이팅 RNN 계보가 “이산화”라는 단일 다리로 통합. Δ 하나가 망각/기억 게이트 역할
O(N) 유한 상태 vs. O(L) KV 캐시: “압축된 기억 vs. 정확한 검색”이라는 시퀀스 모델링의 근본 설계 공간
FlashAttention과의 연결: 공동 저자 Tri Dao가 두 연구 모두 참여 — 하드웨어 최적화 철학의 연속
시사점: KV 캐시 불필요로 배치 확대 가능(5× 처리량) → 비용 효율적 서빙
비판적 코멘트: “대체”보다 “보완”이 더 정확 — 후속 하이브리드(Jamba, StripedHyena)가 실용적으로 더 강력. 3B까지의 증거로 범용적 주장은 과도할 수 있음

Discussion Points

다중 홉 추론: O(N) 유한 상태가 긴 문서의 정확한 사실 검색에서 근본적 한계를 드러내는가?
스케일 검증: 7B+ 규모에서 Transformer 대비 우위 유지 여부
후속 연구: Mamba-2(State Space Duality), Jamba/MambaFormer(하이브리드), VMamba(비전)

BibTeX

@article{gu2023mamba,
  title={Mamba: Linear-Time Sequence Modeling with Selective State Spaces},
  author={Gu, Albert and Dao, Tri},
  journal={arXiv preprint arXiv:2312.00752},
  year={2023},
  url={https://arxiv.org/abs/2312.00752}
}

왜 이 연구를 하는가?

핵심 질문

SSM에 입력 의존적 선택 메커니즘을 부여하면, 선형 시간 복잡도를 유지하면서 Transformer의 내용 기반 추론 능력을 달성할 수 있는가?

기존 접근법의 한계

한계	설명
Transformer O(L²)	긴 시퀀스에서 연산·메모리 병목
S4/H3/Hyena의 LTI	파라미터 고정 → 내용 기반 선택 불능
LTI 추론 실패	Selective Copying 18.3%, Induction Heads 2배 길이에서 실패
기존 선형 모델	RWKV, RetNet, H3 모두 Transformer 품질 미달

핵심 통찰

선택성 = 내용 기반 게이팅: Δ가 크면 현재 입력 강조·과거 망각, 작으면 과거 유지
Theorem 1: 이산화 = 게이팅의 수학적 기반 → SSM과 RNN의 통합적 이해
하드웨어 효율: SRAM 기반 커널 융합 + 병렬 스캔으로 선택성 도입의 효율 손실 상쇄

방법 (Method)

🧭 그 전에 — SSM (State Space Model) 이란 무엇인가
Mamba 본문은 “당신은 SSM을 안다” 고 가정하고 시작한다. 모르면 모든 수식이 막히므로 여기서 한 번 정리한다.

1. 기원: 제어이론(1960s, Kalman filter)
SSM은 원래 딥러닝 개념이 아니라 연속 시간 동역학계 를 기술하는 제어공학의 표준 도구다. 항공기·로봇·통신 시스템의 “현재 상태가 미래 입력에 어떻게 반응하는가”를 모델링한다. Mamba는 이 고전적 도구를 GPU 위의 시퀀스 모델로 재해석한 것.

2. 연속 시간 정의 (가장 기본형, LTI)
입력 $x (t) \in R$ , 숨은 상태 $h (t) \in R^{N}$ , 출력 $y (t) \in R$ 에 대해
$h^{'} (t) = A h (t) + B x (t), y (t) = C h (t) + D x (t)$

$A \in R^{N \times N}$ : 상태 전이 행렬 — 과거가 어떻게 감쇠/유지되는지

$B \in R^{N \times 1}$ : 입력 → 상태 매핑

$C \in R^{1 \times N}$ : 상태 → 출력 매핑

$D$ : skip connection (보통 0 또는 1)

직관: $h (t)$ 는 “지금까지의 입력 전체를 N차원으로 압축한 요약 벡터” 이고, $A$ 는 그 요약을 매 순간 어떻게 업데이트할지 정한다. RNN의 hidden state와 본질적으로 같지만 연속 시간 ODE 형태 라는 점이 다르다.

3. 이산화(Discretization) — ODE를 시퀀스 모델로
우리가 다루는 데이터는 토큰 단위(이산)이므로 위 ODE를 timestep $Δ$ 로 이산화해야 한다. 표준 기법은 Zero-Order Hold (ZOH):
$\overset{ˉ}{A} = exp (Δ A), \overset{ˉ}{B} = (Δ A)^{- 1} (exp (Δ A) - I) \cdot Δ B$
결과:
$h_{t} = \overset{ˉ}{A} h_{t - 1} + \overset{ˉ}{B} x_{t}, y_{t} = C h_{t}$
이제 형식적으로는 선형 RNN 과 동일하다. $Δ$ 의 의미: 크면 한 스텝이 ODE 시간상 길게 진행 → 새 입력의 영향이 커지고 과거가 빠르게 감쇠 (= 망각 게이트). 작으면 반대 (= 기억 유지).

4. 세 가지 동등한 표현 (S4의 핵심 트릭)
LTI(파라미터가 t에 무관) 라는 가정 하에 같은 SSM을 세 가지 방식으로 계산할 수 있다:

표현 형태 언제 쓰나 복잡도
Recurrent $h_{t} = \overset{ˉ}{A} h_{t - 1} + \overset{ˉ}{B} x_{t}$ 추론 (autoregressive) O(N) per step
Convolutional $y = x * \overset{ˉ}{K}$ , where $\overset{ˉ}{K}_{t} = C \overset{ˉ}{A}^{t} \overset{ˉ}{B}$ 학습 (병렬) O(L log L) FFT
Continuous $h^{'} (t) = A h (t) + B x (t)$ 이론 분석 / 연속 신호 —

→ “학습은 합성곱으로 병렬, 추론은 재귀로 효율적” — 이게 S4의 마법이며, RNN의 약점(병렬화 불가)과 Transformer의 약점(추론 시 KV 캐시 폭증)을 동시에 우회.

5. HiPPO 행렬 — A를 어떻게 정하나
무작위 A는 학습이 안 된다. HiPPO 이론(Gu et al., 2020) 은 “과거 신호를 다항식 기저로 최적 압축하는 A”를 닫힌형으로 유도한다 (예: HiPPO-LegS는 Legendre 다항식 기저). Mamba도 이 A 구조를 그대로 상속한다. 직관: $h_{t}$ 는 단순한 hidden state가 아니라 “과거 전체의 N차 다항식 근사 계수” 다.

6. LTI의 치명적 한계 — Mamba가 깨야 했던 벽
위 세 표현이 등가이려면 $A, B, C, Δ$ 가 시간에 무관(=입력에 무관) 해야 한다. 그런데 이게 곧:

“현재 입력에 따라 무엇을 기억/망각할지 결정할 수 없음”

Selective Copying (“noise 사이에서 표시된 토큰만 골라 복사”) 같은 content-based 과제에서 원천 실패

Induction Heads (“[A][B]…[A]→[B]” 패턴 복사) 도 동일 이유로 실패

Mamba의 한 줄 요약은 ” $B, C, Δ$ 를 입력 $x_{t}$ 의 함수로 만들자(=LTI를 포기하자)” 이며, 그 대가로 잃은 합성곱 표현을 hardware-aware parallel scan 으로 보충한 것이다.

7. 시퀀스 모델 가족 안에서의 위치
RNN/LSTM ────┐
             │  hidden state로 과거를 압축하지만 학습 병렬화 불가
─────────────┤
S4/SSM (LTI)─┤  연속 ODE → 합성곱 병렬 학습 + 재귀 추론. 그러나 content selection 불가
             │
Transformer ─┤  완전한 content-based addressing. 그러나 O(L²) + KV 캐시
             │
Mamba (S6) ──┘  Selection으로 SSM에 content 의식 부여 + parallel scan으로 효율 보존
8. 한 줄로
SSM = “과거를 N차원 상태로 압축하는 선형 ODE를 이산화한 것”. LTI일 때는 합성곱처럼 굴고, 비-LTI(=Mamba) 가 되면 RNN처럼 굴지만 GPU에서는 여전히 병렬화 가능. 이게 본문 전체의 무대 장치다.

표현	형태	언제 쓰나	복잡도
Recurrent	$h_{t} = \overset{ˉ}{A} h_{t - 1} + \overset{ˉ}{B} x_{t}$	추론 (autoregressive)	O(N) per step
Convolutional	$y = x * \overset{ˉ}{K}$ , where $\overset{ˉ}{K}_{t} = C \overset{ˉ}{A}^{t} \overset{ˉ}{B}$	학습 (병렬)	O(L log L) FFT
Continuous	$h^{'} (t) = A h (t) + B x (t)$	이론 분석 / 연속 신호	—

프레임워크 개요

graph TD
    IN["입력 x (B, L, D)"] --> NORM["LayerNorm"]
    NORM --> PROJ["입력 투영\nD → 2ED"]
    PROJ --> BR1["게이트 분기 z\n(B, L, ED)"]
    PROJ --> BR2["SSM 분기 x'\n(B, L, ED)"]
    BR2 --> CONV["1D Depthwise Conv"]
    CONV --> ACT["SiLU 활성화"]
    ACT --> S6["선택 메커니즘 S6"]

    subgraph SEL["S6: 입력 의존적 파라미터"]
        DELTA["Delta = softplus(Param + Linear(x))"]
        B_MAT["B = Linear_N(x)"]
        C_MAT["C = Linear_N(x)"]
        A_FIX["A: 고정 대각행렬 (HiPPO)"]
    end

    S6 --> SEL
    SEL --> DISC["ZOH 이산화\nA-bar = exp(Delta*A)"]
    DISC --> SCAN["병렬 스캔 (GPU SRAM)\nh_t = A-bar*h_{t-1} + B-bar*x_t"]
    SCAN --> GATE["게이트 합성\nSSM출력 * SiLU(z)"]
    BR1 --> GATE
    GATE --> OUTPROJ["출력 투영 ED → D"]
    OUTPROJ --> RES["잔차 연결"]

선택 메커니즘 파라미터 비교

파라미터	LTI SSM (기존)	Selective SSM (Mamba)	역할
A	고정 학습	고정 학습 (유지)	상태 감쇠 구조
B	고정 (D, N)	`Linear_N(x)` → (B, L, N)	입력→상태
C	고정 (D, N)	`Linear_N(x)` → (B, L, N)	상태→출력
Δ	고정 스칼라	`softplus(Param + Linear_1(x))` → (B, L, D)	망각/기억 게이트

발견 (Findings)

언어 모델링 (Pile)

모델	PPL ↓	기준선 대비
Pythia-1.4B	7.51	기준
Mamba-1.4B	6.80	-9.5%
Pythia-2.8B	6.73	기준
Mamba-2.8B	6.22	-7.6%

합성 과제: 선택 메커니즘 효과

과제	LTI S4	Mamba (S6)	차이
Selective Copying	18.3%	97.0%	+78.7pp
Induction Heads (1M)	실패	99.8%	근본적 차이

제로샷 다운스트림 (6과제 평균)

모델	평균 정확도
Pythia-2.8B	59.1%
Mamba-2.8B	63.3%

효율성

지표	Mamba vs Transformer
추론 처리량	5×
스캔 속도 (vs PyTorch)	40×
KV 캐시	불필요

이론적 의의

SSM-RNN 통합 이론: Theorem 1로 이산화가 게이팅의 수학적 기반임을 증명. HiPPO→S4→Mamba 계보와 LSTM→GRU 계보를 단일 프레임워크로 통합
선형 시간 + 내용 추론 동시 달성: “선형 시간 ↔ 내용 기반 추론” 트레이드오프라는 가정을 깨뜨림
모달리티-아키텍처 연결: 이산 모달리티(텍스트, DNA)에서 선택성이 극적, 연속 모달리티에서는 LTI 유리 → 모달리티 특성과 최적 아키텍처의 원칙적 연결

핵심 용어 정리

용어	정의
SSM (State Space Model)	h’=Ah+Bx, y=Ch 형태의 선형 동역학계
LTI (Linear Time-Invariant)	파라미터가 시간 불변인 선형 시스템. 합성곱 병렬화 가능하나 내용 선택 불가
Selection Mechanism	B, C, Δ를 입력 함수로 만들어 내용 기반 선택적 상태 업데이트를 가능하게 한 핵심 혁신
Δ (타임스텝)	ZOH 이산화의 타임스텝. 크면 현재 강조·과거 망각, 작으면 과거 유지. GRU 게이트와 동치
S6 (Selective SSM)	B, C, Δ가 입력 의존적인 선택적 SSM. LTI S4의 일반화
Parallel Scan	순차 재귀를 work-efficient 알고리즘으로 병렬화하는 기법
Kernel Fusion	여러 연산을 하나의 GPU 커널로 통합하여 HBM 접근 최소화
Induction Heads	[A][B]…[A]→[B] 패턴 복사. Transformer in-context learning의 핵심 메커니즘
Selective Copying	노이즈 사이의 신호 토큰만 선택적으로 복사. LTI 모델로는 근본적 해결 불가

Juhyeon's Blog

탐색기

Mamba - Linear-Time Sequence Modeling with Selective State Spaces