State Space Model은 무엇일까?

딥러닝보다 오래된 발상이야

State Space Model은 Kalman과 Bellman, 1960년대 LQR/LQG로 이어지는 고전 제어 이론에서 왔어. 핵심 발상은 신경망보다 반세기 먼저 등장했지. 매 시점마다 지금까지의 정보를 고정 크기 hidden state 벡터에 압축하고, 현재 입력을 받아 그 상태를 갱신해.

연속 시간에서는 서로 맞물린 두 미분 방정식으로 상태의 변화와 출력을 나타내. 상태 변화는 ḣ(t) = A·h(t) + B·x(t), 출력은 y(t) = C·h(t)야. 이를 이산화하면 h_t = Ā·h_{t-1} + B̄·x_t, y_t = C·h_t라는 순환식이 돼. A는 상태 전이 행렬, B는 입력 투영, C는 출력 투영이야. 상태 h의 차원을 N이라고 하면 시퀀스가 아무리 길어져도 N은 커지지 않아.

시퀀스 모델링에서 왜 주목받을까

핵심은 상태 차원이 고정된다는 점이야. 추론할 때 100 토큰을 처리하든 1억 토큰을 처리하든 크기 N인 h 벡터 하나만 유지해. 따라서 한 단계당 메모리는 O(1)이야. 각 단계에서 A와 B에 대한 행렬-벡터 곱을 한 번씩 하므로 전체 시퀀스 계산량도 O(n)이지. 어텐션의 O(n²) 계산량과 O(n) KV 캐시에 견주면 분명 매력적이야. Transformer의 품질만 따라잡을 수 있다면 말이지.

문제는 바로 그 조건이야. 끝을 알 수 없는 과거를 고정 크기 벡터에 압축하면 정의상 정보가 손실돼. 지난 5년 동안 SSM 연구가 답해 온 수학적 질문은 이거야. A, B, C와 이산화 방식을 어떻게 골라야 유용한 정보를 가장 적게 잃을까? S4, S4D, H3, Mamba, Mamba-2, Mamba-3는 모두 그 질문에 대한 답이야. 세대가 바뀔 때마다 중요한 작업에서 정보 손실을 줄여 왔어.

Code

단순한 SSM 순전파 — 수식을 옮긴 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

위의 단순한 SSM 순전파를 구현하고 d_state=64인 무작위 (1, 1024, 16) 입력에서 실행해 봐. 상태 행렬 A를 0.99 * eye로 초기화하면 안정적으로 시작할 수 있어. 정상 실행을 확인한 다음 시간을 재 봐. for 반복문이 순차적으로 돌아서 GPU에서도 절망적으로 느리다는 걸 발견할 거야. 다음 레슨에서는 SSM 연구자들이 순환 추론 형태를 잃지 않으면서 학습 단계의 병렬성을 어떻게 되찾았는지 설명해.

딥러닝보다 오래된 발상이야

시퀀스 모델링에서 왜 주목받을까

Code

External links

Exercise

Progress

댓글 0