한줄 요약: 입력에 따라 SSM 파라미터를 동적으로 조절하는 "선택적 상태 공간"으로, Transformer의 O(n²) 어텐션 없이 동등한 언어 모델링 성능을 달성했다.
| 항목 | 내용 |
|---|---|
| 저자 | Albert Gu, Tri Dao |
| 소속 | Carnegie Mellon University, Princeton University |
| 발표 | ICML 2024 (Oral) |
| 링크 | arxiv.org/abs/2312.00752 |
| 키워드 | State Space Model, Selective Mechanism, Linear Complexity, Sequence Modeling |
1. 문제 정의
Transformer의 셀프 어텐션은 두 가지 근본적 한계를 가진다:
1. O(n²) 연산/메모리: 시퀀스 길이에 제곱 비례 → 긴 시퀀스에서 비용 폭증
2. KV Cache 선형 증가: 추론 시 이전 토큰의 Key/Value를 모두 저장 → 메모리 병목
기존 State Space Model(S4 등)은 O(n) 복잡도지만 입력에 무관하게 동일한 변환을 적용(Linear Time-Invariant, LTI)하므로 콘텐츠 기반 추론이 어렵다.
핵심 문제:
Transformer: 콘텐츠 인식 ✓, 효율성 ✗ (O(n²))
기존 SSM (S4): 효율성 ✓ (O(n)), 콘텐츠 인식 ✗ (LTI)
→ 두 장점을 결합할 수 있는가?2. 제안 방법
선택적 상태 공간 (Selective SSM)
기존 SSM (S4):
h(t) = Āh(t-1) + B̄x(t) ← A, B가 고정 (입력과 무관)
y(t) = Ch(t)
Mamba (Selective SSM):
h(t) = Ā(x)h(t-1) + B̄(x)x(t) ← B, C, Δ가 입력 x의 함수
y(t) = C(x)h(t)
B = Linear_B(x) ← 입력에 따라 어떤 정보를 상태에 저장할지 결정
C = Linear_C(x) ← 입력에 따라 어떤 정보를 상태에서 읽을지 결정
Δ = softplus(Linear_Δ(x)) ← 입력에 따라 이전 상태를 얼마나 유지할지 결정왜 "선택적"이 중요한가?
예시: "The cat sat on the mat. It was comfortable."
기존 SSM (LTI):
모든 토큰을 동일한 비율로 상태에 축적
→ "It"이 무엇을 가리키는지 구분하기 어려움
Mamba (Selective):
"cat" → Δ 크게 → 상태에 강하게 저장 (중요한 명사)
"the", "on" → Δ 작게 → 상태를 거의 유지 (기능어, 빠르게 잊음)
"It" → C가 "cat"에 해당하는 상태 차원을 강하게 읽음
→ 콘텐츠 기반의 선택적 정보 흐름하드웨어 인식 알고리즘
선택적 SSM은 입력 의존적이므로 S4의 FFT 기반 병렬 계산 불가. 대신:
문제: 선택적 SSM → 파라미터가 시퀀스마다 다름 → 컨볼루션으로 변환 불가
해결: 하드웨어 인식 스캔 알고리즘
1. HBM에서 Δ, A, B, C를 읽어 SRAM에 로드
2. SRAM 내에서 이산화(discretization) + 순차 스캔 수행
3. 중간 상태를 HBM에 저장하지 않음
4. 역전파 시 재계산 (Flash Attention의 recomputation과 유사)
→ IO 복잡도: O(n) — Flash Attention과 동일한 IO-awareness 철학아키텍처 (Mamba Block)
입력 x
├── Linear (확장) → Conv1d → SiLU → Selective SSM → ×
└── Linear (확장) → SiLU ────────────────────────── ×
↓
Linear (축소)
↓
출력
→ Transformer의 MHA + FFN 2개 서브레이어를 하나의 블록으로 통합
→ 파라미터 효율 향상3. 실험 결과
3.1 언어 모델링 (Pile 데이터셋)
| 모델 | 파라미터 | Perplexity ↓ |
|---|---|---|
| Transformer++ | 1.4B | 8.28 |
| H3 (SSM) | 1.4B | 8.96 |
| Hyena | 1.4B | 8.69 |
| RWKV-4 | 1.4B | 8.48 |
| Mamba | 1.4B | 8.14 |
→ SSM 최초로 Transformer++를 perplexity에서 초과
3.2 스케일링 법칙
300M → 1.4B → 2.8B 스케일링:
Mamba의 스케일링 기울기 ≈ Transformer++
→ 모델이 커져도 성능 격차 유지 또는 확대
Mamba 1.4B ≈ Transformer++ 2× 크기의 성능
→ 같은 성능 도달에 절반의 파라미터3.3 추론 속도
| 모델 | 시퀀스 길이 | 처리량 (tok/s) | 배율 |
|---|---|---|---|
| Transformer | 2K | 680 | 1.0x |
| Transformer | 8K | 198 | 1.0x |
| Mamba | 2K | 1120 | 1.6x |
| Mamba | 8K | 1080 | 5.5x |
| Mamba | 64K | 990 | ∞ (Transformer OOM) |
→ 시퀀스 길이가 길어져도 처리량이 거의 감소하지 않음 (고정 상태 크기)
3.4 Selective Copying & Induction Head (합성 태스크)
| 태스크 | S4 (LTI) | Hyena | Mamba |
|---|---|---|---|
| Selective Copying | 실패 | 실패 | 성공 |
| Induction Head | 실패 | 부분 성공 | 성공 |
→ 선택적 메커니즘이 콘텐츠 기반 추론을 가능하게 함
4. 한계점
- In-context learning 능력: Transformer의 ICL(few-shot) 능력과 비교한 체계적 분석 부족 — 후속 연구에서 ICL이 다소 약함 확인
- 긴 범위 의존성의 한계: 고정 크기 상태(state)가 무한한 과거 정보를 압축하므로, 매우 먼 토큰 간 정확한 참조가 필요한 태스크에서 제한
- 학습 시 병렬성: 순차 스캔은 어텐션의 행렬 곱셈보다 GPU 병렬화 효율이 낮을 수 있음 → 학습 속도는 Transformer와 비슷하거나 약간 느림
- 생태계 성숙도: Transformer 기반 도구(Flash Attention, vLLM, 양자화)와 비교해 최적화/배포 인프라 부족
- 어텐션 패턴 해석 불가: Transformer의 어텐션 맵 같은 해석 도구가 없어 디버깅/분석 어려움
- 검증된 스케일 한계: 논문 시점 최대 2.8B — 수십B 이상에서의 검증은 이후 Mamba-2, Jamba 등에서 진행
5. 의의와 영향
- "Transformer가 유일한 답이 아니다"를 실증한 첫 모델 — 아키텍처 다양성 부활
- 추론 시 일정한 메모리·일정한 시간 per token → 긴 시퀀스에서 압도적 효율
- Mamba-2 (2024): 선택적 SSM과 어텐션의 수학적 연결(SSD 프레임워크) 제시 → 통합 이론
- Jamba (AI21): Transformer + Mamba 하이브리드로 52B 모델 → 실전 검증
- Zamba, Hymba, Samba 등 후속 하이브리드 아키텍처의 핵심 빌딩 블록
- 하드웨어 인식 알고리즘 설계(Tri Dao)가 새 아키텍처에서도 핵심임을 재확인
6. 💬 리뷰어 코멘트
Mamba의 핵심 통찰은 "모든 토큰이 동등하지 않다"를 아키텍처 수준에서 구현한 것이다. Transformer의 어텐션은 모든 토큰 쌍을 비교하여 이를 구현하지만 O(n²) 비용을 지불한다. Mamba는 상태 전이 속도(Δ)를 입력에 따라 조절하여 "중요한 것은 기억하고 불필요한 것은 잊는" 메커니즘을 O(n)에 구현했다.
현실적으로 Mamba가 Transformer를 완전히 대체하기는 어려워 보인다. 하지만 Jamba(AI21)처럼 Transformer 레이어와 Mamba 레이어를 혼합하는 하이브리드 접근은 매우 유망하다. 어텐션이 필요한 곳(정확한 장거리 참조)에는 Transformer를, 효율이 중요한 곳에는 Mamba를 배치하는 전략이다.
Albert Gu(SSM 이론)와 Tri Dao(시스템 최적화)의 조합이 만든 시너지도 주목할 만하다. 같은 Tri Dao가 Flash Attention도 만들었다는 점에서, 결국 AI 발전은 알고리즘과 시스템의 교차점에서 이루어진다.
관련 논문: S4, Hyena, RWKV, Mamba-2, Jamba, Flash Attention
