Mamba
Mamba
State Space Model 기반 선형 복잡도 시퀀스 모델. Transformer의 대안으로 주목.
Mamba
State Space Model 기반 선형 복잡도 시퀀스 모델. Transformer의 대안으로 주목.
Mamba는 State Space Model(SSM) 기반의 시퀀스 모델 아키텍처로, Transformer의 대안으로 주목받고 있습니다. 선형 시간 복잡도(O(n))로 동작하여 Transformer의 이차 복잡도(O(n^2)) 문제를 해결하면서도 동등하거나 우수한 성능을 달성합니다.
2023년 12월 Albert Gu와 Tri Dao가 발표했습니다. 기존 State Space Model(S4 등)의 한계를 극복하기 위해 Selective State Space를 도입했습니다. 이를 통해 입력에 따라 동적으로 상태 전이를 조절할 수 있어, Transformer처럼 맥락을 선택적으로 기억할 수 있게 되었습니다.
Mamba의 핵심 혁신은 두 가지입니다. 첫째, Selective SSM은 입력 데이터에 따라 A, B, C 행렬을 동적으로 변경합니다. 둘째, Hardware-aware Parallel Scan 알고리즘으로 GPU 메모리 계층을 최적화하여 실제 속도를 5배 향상시켰습니다.
실무에서 Mamba는 긴 시퀀스 처리가 필요한 분야에서 유망합니다. Mamba-3B는 동일 크기 Transformer를 능가하고, 2배 크기와 동등한 성능을 보입니다. 2024년 Mamba-2가 출시되어 더욱 효율적인 학습이 가능해졌고, IBM의 Bamba, NVIDIA의 하이브리드 모델 등 상용화가 진행 중입니다.
# Mamba 모델 사용 예제 (mamba-ssm 라이브러리)
# pip install mamba-ssm causal-conv1d
import torch
from mamba_ssm import Mamba
# Mamba 블록 초기화
batch_size = 2
seq_len = 1024
d_model = 768 # 모델 차원
mamba_block = Mamba(
d_model=d_model, # 모델 차원
d_state=16, # SSM 상태 차원
d_conv=4, # 컨볼루션 커널 크기
expand=2 # 확장 비율
).cuda()
# 입력 데이터 생성
x = torch.randn(batch_size, seq_len, d_model).cuda()
# Mamba 블록 순전파
y = mamba_block(x)
print(f"Input: {x.shape} -> Output: {y.shape}")
# Input: torch.Size([2, 1024, 768]) -> Output: torch.Size([2, 1024, 768])
# ============================================
# 사전학습된 Mamba 모델 사용 (HuggingFace)
# ============================================
from transformers import AutoTokenizer, MambaForCausalLM
model_name = "state-spaces/mamba-2.8b"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = MambaForCausalLM.from_pretrained(
model_name,
torch_dtype=torch.bfloat16,
device_map="auto"
)
# 텍스트 생성
prompt = "Mamba is a new architecture that"
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
outputs = model.generate(
**inputs,
max_new_tokens=50,
do_sample=True,
temperature=0.7,
top_p=0.9
)
generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
print(f"\n생성된 텍스트:\n{generated_text}")
# 추론 속도 비교 (Transformer 대비 5배 빠름)
import time
long_input = torch.randn(1, 16384, d_model).cuda() # 16K 시퀀스
start = time.time()
with torch.no_grad():
_ = mamba_block(long_input)
print(f"16K 시퀀스 처리 시간: {time.time()-start:.3f}초")
"긴 문서 분석에 Mamba 기반 모델을 테스트하고 있습니다. 128K 토큰 컨텍스트에서 Transformer보다 추론 속도가 5배 빠르고, 메모리 사용량은 절반입니다. 성능은 동등해서 비용 효율이 훨씬 좋습니다."
"Mamba는 Selective State Space Model을 사용합니다. 기존 SSM은 입력에 무관하게 동일한 상태 전이를 했지만, Mamba는 입력에 따라 A, B, C 행렬을 동적으로 조절해서 Transformer처럼 선택적 주의 기능을 구현했습니다."
"Mamba-2는 State Space Duality 개념으로 SSM과 Attention 모드를 모두 지원합니다. 하이브리드 아키텍처(Mamba + Attention 레이어)가 순수 Transformer나 순수 Mamba보다 성능이 좋다는 연구 결과도 있어요."
In-context Learning이나 복잡한 추론에서는 아직 Transformer가 우위입니다. Mamba는 긴 시퀀스 처리에 특히 강점이 있습니다.
DNA/유전체 분석, 오디오 처리, 긴 문서 요약 등 긴 시퀀스가 필요한 태스크에 Mamba가 적합합니다. 일반 NLP에는 하이브리드 모델을 고려하세요.