State Space Model 이 뭐야?

아이디어는 deep learning 보다 오래됐어

State Space Model 은 classical control theory 에서 와 — Kalman, Bellman, 1960 년대 LQR/LQG 계보 전체. 핵심 아이디어는 neural network 보다 반세기 앞서: 매 timestep 마다 전체 history 를 fixed-size hidden state vector 로 압축하고, 현재 input 기준으로 그 state 를 update.

continuous-time 형태는 두 개의 coupled differential equation: state evolution 은 ḣ(t) = A·h(t) + B·x(t), output 은 y(t) = C·h(t). discretize 하면 recurrence h_t = Ā·h_{t-1} + B̄·x_t, y_t = C·h_t. 끝. A 는 state transition matrix, B 는 input projection, C 는 output projection. state h 의 차원 — N 이라고 부르자 — 은 sequence length 따라 안 자라.

Sequence modeling 에 왜 흥미로운가

Fixed state dimension 이 헤드라인 feature 야. 추론 시점에는 100 토큰을 처리했든 1 억 토큰을 처리했든 size N 의 h vector 하나 유지. step 당 O(1) memory. compute 는 sequence 전체 O(n) — 각 step 이 A 와 B 에 대해 matrix-vector multiply 하나씩 하니까. attention 의 O(n²) compute 와 O(n) KV-cache 와 비교하면 매력은 명백 — Transformer quality 만 따라잡을 수 있다면.

걸리는 건 그 "라면" 이야. 임의로 긴 history 를 fixed-size vector 로 압축하는 건 정의상 lossy. 지난 5 년 동안 SSM 연구가 답해온 수학 질문은: 어떤 A, B, C 선택, 어떤 discretization 이, 가장 적은 useful information 을 잃나? S4, S4D, H3, Mamba, Mamba-2, Mamba-3 다 그 질문에 대한 답이야 — 각각 중요 워크로드에서 이전보다 덜 lossy.

Code

naive SSM forward pass — 수식, PyTorch 8 줄·python

import torch

def ssm_naive(x, A, B, C):
    # x: (batch, seq_len, d_model)
    # A, B, C: (d_state, d_state), (d_state, d_model), (d_model, d_state)
    B_, T, _ = x.shape
    h = torch.zeros(B_, A.size(0), device=x.device)
    ys = []
    for t in range(T):
        h = h @ A.T + x[:, t] @ B.T   # state update
        ys.append(h @ C.T)             # output projection
    return torch.stack(ys, dim=1)

Exercise

위 naive SSM forward pass 구현하고 d_state=64 인 random (1, 1024, 16) input 에서 돌려봐. 안 터지는지 확인 (state matrix A 를 0.99 * eye 로 init 하면 stable starting point). 그 다음 시간 재. 이 naive 버전이 GPU 에서 절망적으로 느린 거 발견할 거야 — for-loop 가 sequential 이니까. 다음 lesson 이 SSM community 가 어떻게 recurrent inference 형태를 안 잃고 training-time parallelism 을 회복했는지 설명해.

아이디어는 deep learning 보다 오래됐어

Sequence modeling 에 왜 흥미로운가

Code

External links

Exercise

Progress

댓글 0