SSM (State Space Model) | KAITRUST AI 백과사전

📖 상세 설명

SSM(State Space Model, 상태 공간 모델)은 제어 이론과 신호 처리에서 유래한 시퀀스 모델링 기법입니다. 숨겨진 상태(hidden state) h(t)를 통해 입력 시퀀스 u(t)를 출력 y(t)로 매핑하며, 연속 시간 미분방정식 h'(t) = Ah(t) + Bu(t), y(t) = Ch(t) + Du(t)를 이산화하여 사용합니다. A, B, C, D 행렬이 시스템의 동적 특성을 결정합니다.

2021년 Stanford의 S4(Structured State Spaces for Sequence Modeling) 논문에서 딥러닝에 효과적으로 적용되었습니다. HiPPO(High-order Polynomial Projection Operators) 초기화로 긴 의존성을 효과적으로 포착합니다. 2023년 Carnegie Mellon의 Mamba가 Selective SSM으로 Transformer에 필적하는 성능을 보이면서 큰 주목을 받았습니다. Transformer의 O(n^2) 복잡도를 O(n)으로 줄여 긴 시퀀스 처리에 획기적으로 유리합니다.

핵심 장점은 선형 시간 복잡도와 이론상 무한한 컨텍스트입니다. Transformer는 KV 캐시가 시퀀스 길이에 비례해 증가하지만, SSM은 고정 크기 상태(보통 16-64차원)로 과거 정보를 압축합니다. Mamba의 Selective SSM은 입력에 따라 A, B, C 파라미터가 동적으로 변해 Attention처럼 관련 정보를 선택적으로 기억합니다. 추론 시 토큰당 계산량이 일정하여 스트리밍 생성에 이상적입니다.

Mamba 2(2024)는 하드웨어 친화적 설계로 Mamba 1 대비 2-8배 빠른 학습을 달성했습니다. Jamba(AI21, Mamba+Attention 하이브리드), Zamba(Zyphra), Falcon Mamba 등 SSM 기반 대규모 모델들이 등장하고 있습니다. DNA/단백질 시퀀스(100K+ 길이), 오디오(16kHz x 수분), 긴 문서 처리, 실시간 스트리밍 등에서 강점을 보이며, Transformer를 보완하거나 대체할 차세대 아키텍처로 주목받고 있습니다.

💻 코드 예제

import torch
import torch.nn as nn
from mamba_ssm import Mamba

# Mamba 블록 사용 예시 (mamba-ssm 라이브러리)
class MambaLM(nn.Module):
    """Mamba 기반 언어 모델"""
    def __init__(
        self,
        vocab_size: int = 50257,
        d_model: int = 768,
        n_layer: int = 12,
        d_state: int = 16,
        expand: int = 2
    ):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, d_model)

        # Mamba 레이어 스택
        self.layers = nn.ModuleList([
            Mamba(
                d_model=d_model,
                d_state=d_state,  # SSM 상태 차원
                d_conv=4,         # 로컬 컨볼루션 크기
                expand=expand     # 확장 비율
            )
            for _ in range(n_layer)
        ])

        self.norm = nn.LayerNorm(d_model)
        self.lm_head = nn.Linear(d_model, vocab_size, bias=False)

    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
        x = self.embedding(input_ids)

        for layer in self.layers:
            x = layer(x) + x  # Residual connection

        x = self.norm(x)
        logits = self.lm_head(x)
        return logits

# 모델 생성 및 추론
model = MambaLM(vocab_size=50257, d_model=768, n_layer=12)
input_ids = torch.randint(0, 50257, (1, 2048))  # 배치 1, 길이 2048

with torch.no_grad():
    logits = model(input_ids)
    print(f"출력 shape: {logits.shape}")  # [1, 2048, 50257]

# Transformer 대비 장점:
# - 추론: O(n) vs O(n^2) - 긴 시퀀스에서 5-10배 빠름
# - 메모리: 고정 상태 크기 vs 증가하는 KV 캐시
# - 학습: FlashAttention 없이도 효율적

# === Mamba 2 (State Space Duality) ===
# Mamba 2는 SSM과 Attention의 이론적 연결성을 활용
# 더 하드웨어 친화적인 설계로 2-8배 빠른 학습

# === Jamba: Hybrid SSM + Attention ===
# AI21의 Jamba는 Mamba와 Attention을 교차 배치
# SSM의 효율성 + Attention의 in-context learning 능력

from transformers import AutoModelForCausalLM

# Jamba 모델 사용 예시 (52B 파라미터, 140K 컨텍스트)
jamba_model = AutoModelForCausalLM.from_pretrained(
    "ai21labs/Jamba-v0.1",
    device_map="auto",
    torch_dtype=torch.bfloat16
)

# === 복잡도 비교 ===
def complexity_comparison(seq_len):
    """Transformer vs SSM 복잡도 비교"""
    transformer_attention = seq_len ** 2  # O(n^2)
    transformer_kv_cache = seq_len * 2 * 768  # 2 * d_model per layer
    ssm_compute = seq_len  # O(n)
    ssm_state = 64  # 고정 상태 크기 (d_state)

    print(f"시퀀스 길이: {seq_len:,}")
    print(f"Transformer Attention 연산: {transformer_attention:,}")
    print(f"Transformer KV Cache (1 layer): {transformer_kv_cache:,} floats")
    print(f"SSM 연산 (토큰당): {ssm_compute:,}")
    print(f"SSM 상태 크기: {ssm_state} (고정)")
    print(f"효율성 비율: {transformer_attention / seq_len:.0f}x")

complexity_comparison(100000)  # 100K 토큰

# === HiPPO 초기화 (긴 의존성 학습) ===
# HiPPO(High-order Polynomial Projection Operators)는
# 긴 시퀀스 의존성을 효과적으로 포착하는 A 행렬 초기화

# === SSM vs RNN vs Transformer 비교 ===
# |           | 학습 복잡도 | 추론 복잡도 | 메모리 | 긴 의존성 |
# |-----------|------------|------------|--------|----------|
# | RNN       | O(n)       | O(1)/토큰  | O(1)   | 어려움   |
# | Transformer| O(n^2)    | O(n)/토큰  | O(n)   | 좋음     |
# | SSM       | O(n)       | O(1)/토큰  | O(1)   | 좋음     |

# SSM은 RNN의 효율성 + Transformer의 성능을 결합!

# === 실제 응용 분야 ===
# 1. DNA/단백질 시퀀스 (수십만 base pair)
# 2. 오디오 처리 (16kHz x 수 분 = 수백만 샘플)
# 3. 초장문 문서 (법률, 의료, 코드베이스)
# 4. 실시간 스트리밍 (일정한 지연시간 보장)

# === Mamba 스타일 Selective SSM 핵심 ===
# 기존 SSM: A, B, C가 입력과 무관한 고정값
# Selective SSM: A, B, C가 입력 x에 따라 동적으로 결정

# x -> Linear -> B, C, delta 생성
# delta로 A를 이산화 (다른 시간 해상도 적용)
# 이를 통해 입력에 따라 무엇을 기억/망각할지 선택

# === State Space Duality (Mamba 2) ===
# SSM과 Attention이 수학적으로 연결된다는 발견
# 특정 조건에서 SSM의 출력 = 구조화된 Attention과 동일
# 이를 활용해 더 하드웨어 친화적인 구현 가능

# === Zamba (Zyphra) ===
# Mamba + Shared Attention 하이브리드
# 일부 레이어만 Attention을 공유하여 효율성 극대화
# 7B 규모에서 Llama 2 13B 수준 성능

# === Falcon Mamba (Technology Innovation Institute) ===
# 순수 Mamba 아키텍처로 7B 모델
# Transformer 없이 경쟁력 있는 성능 달성

# === 주요 SSM 모델 비교 ===
# | 모델      | 크기  | 컨텍스트 | 특징               |
# |----------|-------|---------|-------------------|
# | Mamba    | 3B    | 무제한   | 순수 SSM          |
# | Mamba 2  | 2.8B  | 무제한   | 하드웨어 최적화    |
# | Jamba    | 52B   | 256K    | SSM + Attention   |
# | Zamba    | 7B    | 무제한   | Shared Attention  |
# | RecurrentGemma | 9B | 무제한 | RG-LRU 기반       |

# === SSM 학습 팁 ===
# 1. 초기화: HiPPO 초기화가 긴 의존성 학습에 필수
# 2. 정규화: LayerNorm보다 RMSNorm이 더 효과적
# 3. 잔차 연결: 모든 블록에 residual connection 적용
# 4. 혼합: Mamba + MLP 또는 Mamba + Attention 혼합 고려

# === 벤치마크 성능 (Mamba 3B vs Transformer 3B) ===
# | 태스크 | Mamba 3B | Transformer 3B | 비고 |
# |-------|----------|----------------|-----|
# | HellaSwag | 74.2 | 73.8 | 동등 |
# | WinoGrande | 68.5 | 68.1 | 동등 |
# | 추론 속도 | 5x faster | baseline | Mamba 우위 |

🗣️ 실무에서 이렇게 말하세요

💬 아키텍처 선택 회의에서

"100K 토큰 이상 긴 문서 처리가 필요하면 Mamba나 Jamba 검토해보세요. Transformer는 KV 캐시가 100K면 수십 GB 메모리가 필요한데, SSM은 상태 크기가 고정이라 1GB 미만으로 처리됩니다."

💬 면접에서

"SSM은 상태 공간 모델로, x'=Ax+Bu, y=Cx 미분방정식을 이산화해서 시퀀스를 처리합니다. Mamba는 Selective SSM으로 입력에 따라 A,B,C 파라미터가 동적으로 바뀌어 Attention처럼 컨텍스트 인지가 가능해졌어요. 복잡도는 O(n)입니다."

💬 성능 최적화 논의에서

"실시간 스트리밍 추론이라면 SSM이 확실히 유리해요. Transformer는 토큰마다 전체 KV 캐시를 참조하는데, Mamba는 작은 상태만 업데이트하면 돼서 토큰당 latency가 일정합니다."

⚠️ 흔한 실수 & 주의사항

❌

모든 태스크에 SSM이 좋다고 가정

In-context learning, 복잡한 reasoning에서는 아직 Transformer가 우세합니다. SSM은 긴 시퀀스 처리, 스트리밍에 강점이 있습니다.

❌

라이브러리 호환성 무시

Mamba는 CUDA 커널에 의존하여 특정 GPU(Ampere+)와 CUDA 버전이 필요합니다. 환경 설정을 먼저 확인하세요.

✅

올바른 방법

Hybrid 모델(Jamba처럼 Transformer + Mamba)을 고려하거나, 유스케이스에 맞게 선택하세요. DNA, 오디오, 긴 문서에서 SSM이 강합니다.

🔗 관련 용어

📚 더 배우기

📝 Mamba: Linear-Time Sequence Modeling 📝 S4: Structured State Spaces 📦 Mamba GitHub Repository