State Space Models (SSM)

상태 공간 모델(SSM)은 각 단계에서 업데이트하는 은닉 상태 벡터를 통해 입력 시퀀스를 출력 시퀀스로 변환하므로 시퀀스가 아무리 길어져도 메모리 크기는 일정하게 유지됩니다. 이 개념은 선형 시스템을 상태 업데이트 h' = Ah + Bx와 출력 y = Ch + Dx로 표현하는 제어 이론에서 비롯되었습니다. 칼만 필터는 1960년부터 노이즈가 있는 측정값으로 이러한 상태를 추정해 왔습니다. 딥러닝 SSM은 데이터로부터 A, B, C를 학습하고 이러한 계층을 여러 개 쌓습니다. 2021년 S4와 2023년 Mamba 이후 SSM은 텍스트, 오디오, 시계열 데이터의 긴 시퀀스에서 Transformer와 경쟁해 왔으며, 연산량은 시퀀스 길이에 선형으로 증가합니다.

SSM 작동 방식#

신경망 SSM 계층은 연속 방정식을 학습 가능한 스텝 크기 Δ를 사용하는 이산 재귀식으로 변환하며, 이때 영차 유지 규칙 Ā = exp(ΔA)를 사용합니다. 각 단계에서는 h_t = Ā h_(t-1) + B̄ x_t과 y_t = C h_t을 계산합니다. A, B, C가 모든 타임스텝에서 고정되어 있으면 재귀식을 전개했을 때 출력은 입력과 커널 (CB̄, CĀB̄, CĀ²B̄, ...)의 컨볼루션으로 나타납니다. 따라서 동일한 계층을 학습 중에는 GPU에서 전체 시퀀스를 병렬 처리하는 컨볼루션으로 실행하고, 생성 중에는 토큰마다 일정한 크기의 상태 업데이트를 수행하는 재귀식으로 실행할 수 있습니다. PyTorch에서 두 방식은 동일한 출력을 생성합니다:

import torch

torch.manual_seed(0)
L, N = 64, 16  # 시퀀스 길이, 상태 크기
x = torch.randn(L)

# 대각 연속 시간 SSM을 스텝 크기 delta(영차 유지)로 이산화
A = -(torch.rand(N) + 0.5)
B, C, delta = torch.randn(N), torch.randn(N), 0.1
A_bar = torch.exp(delta * A)
B_bar = (A_bar - 1) / A * B

# 재귀 방식: 생성 과정처럼 단계마다 상태 하나를 업데이트
h, y_rec = torch.zeros(N), []
for t in range(L):
    h = A_bar * h + B_bar * x[t]
    y_rec.append(C @ h)
y_rec = torch.stack(y_rec)

# 컨볼루션 방식: 학습 과정처럼 커널 K_k = C . (A_bar^k * B_bar)를 미리 계산
K = torch.stack([C @ (A_bar**k * B_bar) for k in range(L)])
y_conv = torch.stack([K[: t + 1].flip(0) @ x[: t + 1] for t in range(L)])

print("Same outputs:", torch.allclose(y_rec, y_conv, atol=1e-5))

Same outputs: True을 출력합니다. 재귀 방식은 각 단계에서 상태 값 16개만 유지하고, 미리 계산한 커널을 사용하면 학습 중 출력 64개를 모두 병렬로 계산할 수 있습니다.

S4에서 Mamba까지#

초기 딥러닝 SSM은 규모 있게 학습하기에는 속도가 너무 느리고 메모리를 많이 사용했습니다. Albert Gu, Karan Goel, Christopher Ré가 만든 S4는 컨볼루션 커널을 효율적으로 계산할 수 있도록 상태 행렬에 구조를 부여하고, 긴 과거 정보를 기억하도록 설계된 초기화를 적용했습니다. S4는 Long Range Arena 벤치마크의 모든 작업에서 최신 성능을 달성했고, 16K 스텝 시퀀스로 구성된 Path-X 작업을 해결한 최초의 모델이 되었으며, Transformer보다 60배 빠르게 생성했습니다.

고정 커널은 모든 입력을 똑같이 처리하므로 중요하지 않은 내용을 건너뛰고 핵심 사실을 유지할 수 없습니다. Albert Gu와 Tri Dao가 만든 Mamba는 Δ, B, C를 현재 토큰의 함수로 만들어 모델이 상태에 기록할 정보와 잊을 정보를 선택하도록 합니다. 이러한 선택적 매개변수는 컨볼루션 방식의 관점을 사용할 수 없게 하므로 Mamba는 하드웨어를 고려한 병렬 스캔으로 재귀 연산을 수행합니다. Mamba는 Transformer보다 추론 처리량이 5배 높았으며, Mamba-3B는 같은 크기의 Transformer보다 뛰어나고 크기가 두 배인 Transformer와 동등한 성능을 보였습니다. 각 Mamba 블록은 선택적 SSM을 짧은 컨볼루션 및 SiLU 게이트와 결합합니다:

Mamba-2는 선택적 SSM과 선형 어텐션의 한 형태가 같은 연산을 바라보는 서로 다른 관점임을 보였습니다. 저자들은 이 프레임워크를 상태 공간 쌍대성이라고 부르며, 이를 통해 핵심 계층을 2~8배 빠르게 만들었습니다.

SSM과 관련 아키텍처 비교#

• Transformer: 셀프 어텐션은 모든 토큰을 서로 비교하므로 이전의 어떤 토큰이든 정확히 복사할 수 있지만, 연산 비용은 제곱에 비례해 증가하고 KV 캐시는 컨텍스트 윈도우와 함께 커집니다. SSM은 과거 정보를 고정된 상태로 압축하므로 비용은 더 저렴하지만 정보 손실이 발생합니다. 2024년 논문 Repeat After Me는 긴 입력을 복사하고 그 안에서 정보를 검색하는 능력에서 Transformer가 고정 상태 SSM보다 뛰어남을 보였습니다.
• 순환 신경망(RNN): RNN이나 LSTM은 비선형 업데이트를 적용하므로 학습 중 시퀀스를 토큰 단위로 순차 처리해야 합니다. SSM 업데이트는 상태에 대해 선형이므로 컨볼루션이나 병렬 스캔을 사용할 수 있습니다.
• 은닉 마르코프 모델(HMM): HMM은 전이 확률로 연결된 유한한 개수의 이산 상태를 갖습니다. 칼만 필터의 기반인 선형 가우시안 상태 공간 모델은 이에 대응하는 연속 모델이며, 딥러닝 SSM도 이 연속 상태 벡터를 유지합니다.

SSM의 활용 분야#

프로덕션에서의 주요 활용은 정확한 회상을 위해 어텐션 계층을 일부 유지하는 하이브리드 대규모 언어 모델(LLM)입니다. AI21의 Jamba는 Transformer 계층과 Mamba 계층 사이에 전문가 혼합(MoE) 계층을 배치합니다. 단일 80 GB GPU에서 실행할 수 있으며 최대 256K 토큰의 컨텍스트에서 우수한 결과를 보고합니다. IBM의 하이브리드 Granite 4.0-H 모델은 Mamba-2 계층과 Transformer 계층을 9:1 비율로 결합합니다. Vision Mamba는 선택적 스캔을 이미지 패치에 적용하며, SSM은 전체 어텐션을 적용하기에는 너무 긴 원시 오디오와 DNA 시퀀스도 모델링합니다. Ultralytics YOLO26 추적에서는 고전적인 형태의 SSM을 사용합니다. 기본 TrackTrack 추적기는 ByteTrack 및 BoT-SORT와 마찬가지로 칼만 필터를 사용해 각 객체의 다음 바운딩 박스를 예측합니다. 8차원 상태에는 바운딩 박스의 중심과 크기, 그리고 그 속도가 포함됩니다. 이 상태는 일정 속도 전이 행렬과 네 개의 바운딩 박스 좌표를 읽는 선형 관측 행렬을 사용해 업데이트됩니다.