이 논문은 시퀀스 모델링의 핵심 딜레마 — 표현력 vs. 효율성 — 를 "선택적 압축(selective compression)"의 관점에서 본다. Transformer의 attention은 모든 맥락을 명시적으로 기억하기에 강력하지만 이차(quadratic) 비용을 치른다; Mamba는 입력에 따라 무엇을 기억하고 무엇을 잊을지를 동적으로 결정함으로써, 선형 시간 복잡도를 유지하면서 Transformer 수준의 성능에 처음으로 도달한다.
arXiv 2312.00752 →현대 딥러닝의 파운데이션 모델(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) 비용이 발생한다.
이런 한계를 극복하려는 시도로 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) 복잡도를 달성했다.
이 한계는 선택적 복사 과제(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)로 압축하는 것이라 주장한다.
이 관점에서 보면 기존 모델들의 위치가 명확해진다:
시퀀스의 각 위치를 클릭하면 해당 입력이 상태 업데이트에 미치는 영향 차이를 확인할 수 있습니다.
기존의 여러 접근법들이 효율성 문제를 해결하려 했다. linear attention, gated convolution 등이 있었지만, 이들은 content-based reasoning 능력을 포기했다. Mamba의 선택은 단순하지만 강력하다: "파라미터를 입력의 함수로 만들어라."
Mamba 아키텍처는 H3 블록과 MLP 블록을 하나의 단순화된 블록으로 통합한다. 이 블록을 균일하게(homogeneously) 반복하여 전체 모델을 구성한다.
각 블록을 클릭하면 자세한 설명이 나타납니다. 오렌지 강조 = 이 논문의 핵심 기여
| 단계 | 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² 파라미터를 갖는다.
| 변수 | 의미 | 비고 |
|---|---|---|
x(t) ∈ ℝ | 1차원 입력 시퀀스 | 시간 t에서의 입력값 |
h(t) ∈ ℝᴺ | N차원 잠재 상태(latent state) | 현재까지 압축된 맥락 |
y(t) ∈ ℝ | 1차원 출력 | 상태를 읽어낸 값 |
A ∈ ℝᴺˣᴺ | 상태 전이 행렬 | 상태가 어떻게 진화하는지 결정. 대각(diagonal) 구조로 단순화 |
B ∈ ℝᴺˣ¹ | 입력 투영 행렬 | 입력이 상태에 어떻게 영향을 미치는지 |
C ∈ ℝ¹ˣᴺ | 출력 투영 행렬 | 상태에서 출력을 읽어내는 방법 |
직관적 해설
수학적 유도
| 변수 | 의미 | 비고 |
|---|---|---|
Δ | 타임스텝 크기 (step size) | 연속 시스템을 이산 시스템으로 변환하는 샘플링 간격. Mamba에서는 입력의 함수! |
Ā | 이산화된 상태 전이 행렬 | exp(ΔA) 계산. Δ가 크면 상태가 빠르게 리셋됨 |
B̄ | 이산화된 입력 행렬 | 복잡해 보이지만 Δ→0이면 ΔB에 수렴 (Euler 근사) |
직관적 해설
수학적 유도
| 변수 | 의미 | 비고 |
|---|---|---|
h_t ∈ ℝᴺ | 타임스텝 t에서의 히든 상태 | 이전 상태 h_{t-1}과 현재 입력 x_t를 결합 |
x_t ∈ ℝ | 타임스텝 t에서의 입력 | LTI에서는 (Ā,B̄)가 모든 t에서 동일 → 선택 불가 |
직관적 해설
수학적 유도
선택 메커니즘의 핵심은 단순하다: Δ, B, C를 입력 x의 함수로 만든다.
입력: 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) ⟶ 재귀 또는 합성곱으로 계산 가능
입력: 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 게이팅과의 연결)에서 설명된다.
| 변수 | 의미 | 비고 |
|---|---|---|
g_t | 게이트 값 (0~1) | 0이면 이전 상태 유지, 1이면 현재 입력으로 완전 교체 |
σ | 시그모이드 함수 | ZOH 이산화로부터 자연스럽게 도출됨 |
직관적 해설
수학적 유도
선택 메커니즘을 추가하면 합성곱을 쓸 수 없어, 단순(naive) 재귀 구현은 느리다. Mamba는 세 가지 기법으로 이를 극복한다.
| 방법 | 학습 복잡도 | 추론 상태 크기 | 내용 인식 |
|---|---|---|---|
| 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) 상태 | ✓ 있음 |
| 항목 | 값 | 비고 |
|---|---|---|
| 확장 팩터 E | 2 | 내부 채널 수 = 2D |
| SSM 상태 차원 N | 16 | 효율성 벤치마크에 사용된 기본값 |
| A 행렬 구조 | 대각(diagonal) | 계산 효율을 위해 N개의 수로 표현 가능 |
| A 초기화 (실수) | S4D-Real: −(n+1) | HIPPO 이론 기반. n번째 원소 = -(n+1) |
| Δ 초기화 | τ_Δ⁻¹(Uniform[0.001, 0.1]) | bias 항에 적용 |
| 활성화 함수 | SiLU / Swish | SwiGLU 변형과 동일한 선택 |
| 정규화 레이어 | 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배 속도 향상을 보고한다.
| 모델 | Selective Copying 정확도 | 비고 |
|---|---|---|
| S4 (LTI) | 실패 (18%) | 정적 커널로는 변하는 간격 처리 불가 |
| H3 (LTI + gate) | 부분 성공 | 게이팅은 시퀀스 축 상호작용 없음 |
| Mamba (S6) Ours | 100% | 내용 인식으로 완벽 해결 |
"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 모델이다.
| 모델 | 파라미터 | Pile ppl ↓ | LAMBADA ppl ↓ | HellaSwag ↑ | 평균 ↑ |
|---|---|---|---|---|---|
| Pythia-160M | 160M | 29.64 | 38.10 | 30.2 | 40.6 |
| Mamba-130M Ours | 130M | 10.56 | 16.07 | 35.3 | 44.7 |
| Pythia-1.4B | 1.4B | 7.51 | 6.08 | 52.1 | 55.2 |
| Mamba-1.4B Ours | 1.4B | 6.80 | 5.04 | 59.1 | 59.7 |
| Pythia-2.8B | 2.8B | 6.73 | 5.04 | 59.3 | 59.1 |
| Mamba-2.8B Ours | 2.8B | 6.22 | 4.23 | 66.1 | 63.3 |
논문에 따르면 Mamba는 동일 크기 Transformer 대비 4-5배 높은 추론 처리량을 달성한다. 예: Mamba-6.9B (미학습)는 Transformer-1.3B보다 높은 처리량. KV 캐시가 필요 없기 때문에 훨씬 큰 배치 크기를 사용할 수 있기 때문이다.
"Replacing any of these with a selective SSM (S6) significantly improves performance, validating the motivation of Section 3."이들 중 어느 것이든 선택적 SSM(S6)으로 교체하면 성능이 크게 향상되어, 3장의 동기를 검증한다.
| 아키텍처 (블록) | 내부 레이어 | Perplexity ↓ |
|---|---|---|
| H3 | S4 (복소수) | ~9.0 (추정) |
| H3 | S4 (실수) | 유사 |
| Mamba | S6 (선택적) | 큰 개선 |
선택 메커니즘(S6)이 핵심 기여임을 에블레이션이 검증. LTI 모델 간 차이보다 LTI→Selective 전환의 효과가 훨씬 크다.
논문에 수록된 모든 주요 figure/table을 원본 그대로 보존합니다. 아래 설명을 참조하여 논문의 해당 위치를 찾을 수 있습니다.
주목할 포인트: Figure 1에서 Δ_t가 "Selection Mechanism"을 통해 입력 x_t에 의존함을 화살표로 명시하고 있다. 이것이 LTI와의 핵심 차이. 수식 (1)-(4)는 연속 시간 시스템 → 이산 재귀 → 합성곱의 세 가지 동등한 표현을 보여준다.
주목할 포인트: Figure 2에서 일반 Copying task(일정한 간격)는 LTI 모델이 쉽게 해결하지만, Selective Copying(무작위 간격)는 내용 인식이 필요하므로 LTI 모델이 실패한다. Induction Heads task는 LLM의 in-context learning 능력의 핵심으로, 맥락 기반 추론이 요구된다.
주목할 포인트: Mamba 블록은 H3에서 첫 번째 곱셈 게이트를 활성화 함수로 교체했고, MLP 블록에서 SSM을 메인 브랜치에 추가했다. 이 단순화가 균일한(homogeneous) 아키텍처를 가능하게 한다.
주목할 포인트: Mamba는 처음으로 Transformer++ (현재 표준이 된 강력한 Transformer 레시피)의 퍼플렉시티 곡선을 따라잡는 attention-free 모델이다. 특히 시퀀스 길이가 길어질수록 그 격차가 줄어든다.
주목할 포인트: Figure 5 오른쪽에서 HyenaDNA는 시퀀스 길이가 증가할수록 perplexity가 나빠지는 반면, Mamba는 1M 토큰까지 꾸준히 개선된다. 이것이 LTI 모델의 근본적 한계(컨텍스트 필터링 불가)를 실증적으로 보여주는 핵심 결과다.
주목할 포인트: Figure 8 왼쪽에서 표준 PyTorch scan 구현 대비 ~40배, FlashAttention-2 대비 시퀀스 길이 2K 이상에서 더 빠른 것을 볼 수 있다. 이는 CUDA 커널 융합의 효과다.
| 얻은 것 | 잃은 것 |
|---|---|
| ✓ 내용 인식(content-aware) 능력 | ✗ FFT 기반 합성곱 학습 효율성 |
| ✓ 선형 추론 복잡도 (KV 캐시 불필요) | ✗ 합성곱 모드의 단순함 |
| ✓ 1M 길이까지 성능 향상 | ✗ GPU tensor core 최적화 어려움 |
| ✓ 5x 높은 추론 처리량 | ✗ 초대형 규모 검증 미흡 |
Mamba는 다음 상황에서 특히 유용하다: (1) 매우 긴 시퀀스 처리가 필요한 경우 (게놈, 오디오, 긴 문서), (2) 추론 처리량이 중요한 배포 환경, (3) Transformer의 이차 복잡도가 병목인 엣지 디바이스/실시간 시스템.
논문은 다음 방향을 명시적으로 언급한다: (1) Selective SSM의 A 파라미터도 선택적으로 만드는 실험 (현재 고정), (2) 복소수 vs 실수 파라미터화의 데이터 모달리티별 최적 선택, (3) 다양한 도메인(강화학습, 시계열 등)으로의 확장.
Mamba(2023.12)는 SSM 연구의 전환점이 됐다. 이 논문 이후 VMamba(비전), Jamba(Mamba+Transformer 하이브리드), Mamba-2(행렬 구조 개선), RWKV-6 등이 등장했으며, 많은 연구가 선택적 SSM을 기반으로 한다.
"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 기반의 유용한 초기화도 그대로 쓸 수 있다.