Transformer 없이 언어 모델 만들기 — Mamba(SSM) 완전 정복
AI/LLM 시리즈 : Mamba 아키텍처 완전 정복
2017년 Transformer가 등장한 이후, "Attention Is All You Need"는 거의 모든 언어 모델의 기반이 되었습니다. GPT, BERT, LLaMA, Claude... 이름만 들어도 알 만한 모델들이 전부 Transformer 위에 세워졌습니다.
그런데 Transformer에는 태생적인 한계가 하나 있습니다. Self-Attention의 계산 복잡도가 O(T²) 라는 것입니다. 시퀀스 길이가 2배가 되면 연산량은 4배가 됩니다. 컨텍스트 윈도우를 1M 토큰까지 늘리고 싶다면? 이론적으로는 거의 불가능합니다.
2023년 말, Albert Gu와 Tri Dao는 이 문제를 정면으로 돌파한 논문을 발표했습니다. 바로 Mamba입니다.
왜 Transformer의 Attention이 느릴까?
Attention의 핵심 연산을 다시 떠올려보겠습니다.
Attention(Q, K, V) = softmax(QK^T / √d) * V
여기서 QK^T를 계산할 때 T × T 행렬이 만들어집니다. T가 4,096이면 약 1,600만 개, T가 100,000이면 100억 개의 원소를 계산해야 합니다. 게다가 이 행렬을 GPU 메모리(HBM)에 올려야 하므로, 긴 시퀀스에서는 메모리 병목도 심각해집니다.
상태 공간 모델(SSM)이란?
Mamba를 이해하려면 먼저 상태 공간 모델(State Space Model, SSM) 부터 알아야 합니다.
SSM은 원래 제어 이론에서 나온 개념입니다. 입력 신호 x(t)를 받아 내부 상태 h(t)를 업데이트하고, 출력 y(t)를 내보내는 시스템을 수식으로 표현하면 이렇습니다.
h'(t) = A·h(t) + B·x(t) ← 상태 전이
y(t) = C·h(t) ← 출력
이를 이산화(discretize)하면 딥러닝에서 쓸 수 있는 형태가 됩니다.
h_t = Ā·h_{t-1} + B̄·x_t
y_t = C·h_t
이 구조를 보면 LSTM이나 GRU와 비슷해 보입니다. 맞습니다. SSM은 RNN의 현대화 버전이라고 볼 수 있습니다. 차이점은 행렬 A, B, C가 학습 가능한 파라미터라는 점입니다.
SSM의 강점: 선형 시간 추론
RNN처럼 상태 h를 하나씩 업데이트하면, 추론 시 복잡도가 O(T) 가 됩니다. 시퀀스가 길어져도 이전 상태 하나만 들고 있으면 됩니다. Transformer의 KV 캐시가 시퀀스 길이에 비례해서 커지는 것과 대조적입니다.
학습 시에는 병렬 컨볼루션 형태로 변환하여 GPU를 효율적으로 활용할 수 있습니다.
기존 SSM의 문제점
S4, S5 같은 초기 SSM 모델들도 있었지만, 언어 모델에서 Transformer를 이기지 못했습니다. 이유는 하나였습니다.
A, B, C 행렬이 입력과 무관하게 고정되어 있다.
자연어에서 "나는 은행에 갔다"는 문장에서 "은행"이 금융 기관인지, 강가인지는 문맥에 따라 결정됩니다. 하지만 고정된 행렬로는 이런 입력 의존적인 처리가 불가능합니다. Transformer의 Attention은 이걸 잘 하기 때문에 언어 모델에서 강세를 보였던 겁니다.
Mamba의 핵심 혁신: 선택적 상태 공간
Mamba가 이 문제를 해결한 방식은 단순하지만 강력합니다.
B, C, Δ(스텝 크기)를 입력 x의 함수로 만들자.
# 기존 SSM: 고정된 파라미터
B = 학습된_고정_행렬
C = 학습된_고정_행렬
# Mamba: 입력에 따라 동적으로 결정
B = Linear(x) # 입력에 따라 변함
C = Linear(x) # 입력에 따라 변함
Δ = softplus(Linear(x)) # 스텝 크기도 입력에 따라
이를 Selective SSM 이라고 부릅니다. 이제 모델이 어떤 정보를 상태에 기억할지, 어떤 정보를 흘려보낼지를 입력을 보고 동적으로 결정할 수 있게 됩니다.
직관적으로 설명하면 이렇습니다.
- 중요한 토큰 (예: "은행" 앞의 "강가에 있는"): Δ를 크게 → 상태를 천천히 업데이트 → 기억을 오래 유지
- 무관한 토큰 (예: 조사 "에"): Δ를 작게 → 상태를 빠르게 업데이트 → 흘려보냄
이것이 Attention의 "어디를 볼지 선택"과 유사한 역할을 합니다.
하드웨어 인식 병렬화: 실제로 왜 빠른가
선택적 SSM의 문제는 B, C가 입력마다 달라지면 병렬 컨볼루션 형태로 변환할 수 없다는 것입니다. 다시 순차적으로 계산해야 하고, 이러면 GPU를 제대로 활용하지 못합니다.
Mamba는 이를 Parallel Scan 알고리즘과 커널 퓨전으로 해결합니다.
일반적인 방식:
HBM ← 로드 → SRAM → 계산 → SRAM → 저장 → HBM
(매 스텝마다 메모리 왕복)
Mamba의 방식:
HBM → SRAM에 한 번에 로드
SRAM 내에서 모든 스텝 계산 (재구성 생략)
결과만 HBM에 저장
이렇게 하면 메모리 I/O 횟수를 크게 줄일 수 있습니다. FlashAttention이 Attention 계산을 가속한 것과 동일한 원리입니다.
Mamba 블록 구조
전체 Mamba 블록은 다음과 같이 구성됩니다.
입력 x
│
├─ Linear (확장)
│ │
│ ├─ Conv1d (로컬 컨텍스트)
│ │ │
│ │ └─ SiLU 활성화
│ │ │
│ │ └─ Selective SSM ← B, C, Δ 계산
│ │
│ └─ SiLU 활성화 (게이트)
│
└─ 두 경로를 곱함 (게이팅)
│
└─ Linear (축소)
│
└─ 출력 y
SSM 전에 Conv1d를 두는 이유는 SSM이 위치 정보를 직접 다루지 않기 때문입니다. 1D 컨볼루션이 인접한 토큰들 사이의 지역적 패턴을 먼저 처리해줍니다.
Transformer와의 성능 비교
논문에서 보고한 핵심 결과를 정리하면 아래와 같습니다.
| 모델 | 파라미터 | Pile PPL | 추론 처리량 (tok/s) |
|---|---|---|---|
| Transformer++ | 370M | 8.09 | 기준 |
| Mamba | 370M | 8.03 | 5× 빠름 |
| Transformer++ | 1.3B | 7.41 | 기준 |
| Mamba | 1.3B | 7.25 | 5× 빠름 |
같은 파라미터 수에서 품질은 비슷하거나 우세, 속도는 5배 빠릅니다. 특히 긴 시퀀스에서 격차가 더 커집니다.
한계와 현실
Mamba가 모든 것을 해결해주지는 않습니다.
정보 압축의 한계: 상태 h는 고정된 크기입니다. 아무리 선택적이어도, 매우 먼 과거의 정보를 완벽하게 보존하는 건 어렵습니다. Transformer는 KV 캐시에 모든 과거 토큰을 저장하지만, Mamba는 압축된 상태만 유지합니다.
In-context Learning 약점: 최근 연구들에 따르면 Mamba는 few-shot 학습처럼 프롬프트 내 패턴을 즉각적으로 활용하는 능력에서 Transformer에 다소 뒤처집니다. 이건 Attention이 "지금 이 프롬프트를 직접 보는" 구조이기 때문에 생기는 차이입니다.
Hybrid 모델의 등장: 이런 이유로 최근에는 Mamba 레이어와 Attention 레이어를 섞은 하이브리드 모델(Jamba, Zamba 등)이 실용적인 대안으로 주목받고 있습니다.
마무리
Mamba는 "Attention이 없어도 된다"는 것을 처음으로 대규모에서 설득력 있게 보여준 모델입니다. O(T²)의 벽을 선택적 상태 공간과 하드웨어 인식 병렬화로 우회한 접근은, LLM 서빙 비용이 점점 중요해지는 지금 시점에서 매우 실용적인 의미를 가집니다.
Transformer의 시대가 끝났다기보다는, 이제 우리에게 진짜 선택지가 생겼다는 게 맞는 표현일 것 같습니다.
참고 자료
- Gu & Dao (2023). "Mamba: Linear-Time Sequence Modeling with Selective State Spaces." arXiv:2312.00752
- GitHub: state-spaces/mamba