Mamba란 무엇일까 — 한 레슨으로 정리하는 선택성 · 어텐션 너머 퀘스트

바뀐 건 하나야. 매개변수가 입력을 따라 움직여

Mamba(Gu & Dao, 2023.12, arXiv:2312.00752)는 Selective State Space Model, 줄여서 S6야. 핵심은 한 문장으로 끝나. B, C, Δ가 고정값이 아니라 현재 입력 x_t에 따라 달라져. S4·S4D·H3와 갈라지는 지점이 바로 여기야.

B와 C는 x_t를 선형 투영해서 만들고, 이산화 단계의 크기인 Δ도 x_t에서 뽑아. 선형 투영 뒤에 softplus를 붙여 양수로 유지하지. Δ는 토큰 하나가 상태에 얼마나 강하게 들어갈지를 정하는 게이트라고 보면 돼. Δ가 크면 현재 토큰을 세게 받아들이고, 작으면 기존 상태를 더 오래 붙잡아. 선택성 게이트가 바로 Δ야. 이제 모델은 중요한 토큰은 기억하고 군더더기는 흘려보낼 수 있어.

표현력이 늘어난 이유도 분명해

LTI 제약은 교과서 속 수학적 한계로 끝나지 않았어. 실제 대가가 있었지. Mamba 이전 SSM은 내용을 보고 압축 강도를 바꾸지 못했어. 핵심 사실이든 불용어든 모두 같은 시간 불변 필터를 통과했거든. 선택성을 넣으면 정보가 많은 토큰은 상태에 깊이 담고, 중요하지 않은 토큰은 거의 흔적 없이 흘려보낼 수 있어.

LSTM과 GRU의 게이트가 하던 일과 닮았지. 한편으로 보면 SSM 계열이 S4·S4D·H3에서 단순화를 위해 덜어 냈던 내용 의존 게이트를 다시 찾아낸 셈이야. 다만 이번에는 제어 이론과 연결되는 SSM의 틀과 하드웨어 친화적인 구조를 지킨 채, LSTM이 오래전부터 쓰던 게이트를 얻었어.

대가는 합성곱 관점이야

공짜는 없어. B, C, Δ가 x_t에 따라 달라지면 순환식 전체를 하나의 전체 합성곱 커널로 펼칠 수 없어. 시점마다 커널이 달라지니 S4의 FFT 기반 병렬 학습 기법도 쓸 수 없지. Mamba는 학습을 병렬화할 새로운 길을 만들어야 했고, 그 답이 다음 레슨의 하드웨어 인지 selective scan이야.

Code

Mamba S6 한 단계 — 토큰별 갱신 의사 코드·python

# x_t: (batch, d_model) — timestep t 의 input
# A: (d_state, d_state) — fixed (HiPPO-style 로 init)
# h_{t-1}: (batch, d_state) — 이전 state

B_t = linear_B(x_t)            # (batch, d_state)
C_t = linear_C(x_t)            # (batch, d_state)
delta_t = softplus(linear_d(x_t))  # (batch, 1) — selectivity gate

A_bar = torch.exp(delta_t * A)         # discretized state transition
B_bar = delta_t * B_t                  # discretized input projection

h_t = A_bar @ h_{t-1} + B_bar * x_t    # state update (이제 input-dependent)
y_t = C_t * h_t                        # output

Exercise

mamba-ssm을 설치하고(CUDA 필요) Hugging Face의 Mamba-130M 체크포인트를 짧은 생성 작업에서 실행해 봐. 1K, 8K, 32K, 128K 토큰마다 프리필 시간을 재서 매개변수 수가 비슷한 Transformer인 Pythia-160M과 비교해. Mamba는 거의 선형으로, Pythia는 제곱으로 늘어나는지 확인해 봐. 선택성이 들어와도 이 점근적 차이는 그대로야.

Mamba란 무엇일까 — 한 레슨으로 정리하는 선택성

바뀐 건 하나야. 매개변수가 입력을 따라 움직여

표현력이 늘어난 이유도 분명해

대가는 합성곱 관점이야

Code

External links

Exercise

Progress

댓글 0