Mamba (SSM) - 인공지능 > 딥러닝 | AI Insight Note

Mamba는 2023년 Gu와 Dao가 제안한 State Space Model(SSM) 기반 시퀀스 모델로, Transformer의 O(N²) 어텐션 복잡도 없이 O(N) 선형 복잡도로 긴 시퀀스를 처리한다.

Mamba vs Transformer

항목	Mamba	Transformer
시간 복잡도	O(N)	O(N²)
공간 복잡도	O(N)	O(N²)
긴 시퀀스	효율적	비효율
선택적 상태	입력 의존적 (핵심)	어텐션으로 구현
추론 속도	RNN처럼 일정	길이에 따라 느려짐
병렬 학습	특수 알고리즘 필요	용이

SSM 수식

연속 시간 SSM:
  h'(t) = Ah(t) + Bx(t)
  y(t)  = Ch(t) + Dx(t)

이산화 (ZOH):
  h_t = Ā h_{t-1} + B̄ x_t
  y_t = C h_t

Mamba의 선택적 SSM (S6):
  B, C, Δ = 입력 x_t의 함수 (선택성 부여)
  → 관련 없는 정보 필터링 가능

Python 예시 (mamba-ssm 패키지)

python

from mamba_ssm import Mamba
import torch

# d_model: 모델 차원, d_state: SSM 상태 차원
model = Mamba(
    d_model=256,
    d_state=16,
    d_conv=4,
    expand=2,
).cuda()

batch, length, dim = 2, 1024, 256
x = torch.randn(batch, length, dim).cuda()
y = model(x)  # (2, 1024, 256), O(N) 메모리

Mamba-2와 State Space Duality

Mamba-2(2024)는 SSM과 어텐션의 수학적 등가성(State Space Duality)을 발견하여 더욱 효율적인 하드웨어 최적화 구현을 제공한다. 멀티헤드 구조로 Transformer와의 통합이 용이해졌다.

Mamba (SSM)Mamba

Mamba vs Transformer

SSM 수식

Python 예시 (mamba-ssm 패키지)

Mamba-2와 State Space Duality

관련 노트

생성적 적대 신경망Generative Adversarial Network

합성곱 신경망Convolutional Neural Network

완전 연결 신경망Fully Connected Network