2. Mamba 상세

한 줄 요약

"각 토큰마다 동적으로 변하는 '상태(state)'를 유지하며, 왼쪽에서 오른쪽으로 한 번에 흘려서 정보를 섞는 O(N) 연산."

어떻게 동작하나요?

Attention 이 "모든 과거 토큰을 한꺼번에 다시 보는" 방식이라면, Mamba 는 "작은 메모리(state) 하나를 들고 왼쪽에서 오른쪽으로 읽어가면서 계속 업데이트" 하는 방식입니다. 마치 RNN 처럼요. 하지만 Mamba 는 입력에 따라 이 state 업데이트 규칙이 동적으로 바뀝니다 (selective SSM) — 이 점이 고전 RNN/LSTM 보다 훨씬 강력합니다.

핵심 아이디어: 1. State Space Model (SSM): 연속 시간 미분방정식 h'(t) = A·h(t) + B·x(t) 를 이산화해서 h[k] = Ā·h[k-1] + B̄·x[k] 로 만든다. 2. Selective: A, B 가 고정이 아니라 입력 x 에 따라 달라진다 (selective scan). 3. Parallel scan: 왼쪽-오른쪽 의존성이 있지만 GPU 에서 병렬 scan 알고리즘으로 빠르게 계산. 4. Hardware-aware: SRAM 에 최적화된 CUDA 커널로 실제 속도도 Attention 못지 않게 빠름.

결과: 시퀀스 길이에 대해 O(N) 선형, state 는 고정 크기(d_state), 추론 시 KV 캐시 대신 작은 state 만 유지하면 됨.

장점

• 선형 복잡도: 시퀀스가 2배 길어져도 비용도 2배 (Attention 의 4배 대신).
• 작은 추론 메모리: 긴 문맥에서도 state 크기가 일정. KV 캐시처럼 폭발하지 않음.
• 매우 긴 시퀀스 처리: 32K, 128K, 1M 토큰까지 선형 확장.
• Selective mechanism: 중요한 정보는 기억하고 노이즈는 잊는 학습 가능.

단점

• In-context recall 약함: 작은 state 에 모든 정보를 압축하다 보니 "문맥 내 정확한 항목 찾기" 태스크에서는 Attention 에 밀린다. → 그래서 하이브리드(Mamba + Attention) 가 중요.
• 새로운 하드웨어 커널: FlashAttention 만큼 범용 최적화되지 않아 일부 GPU/환경에서 느릴 수 있음.
• 복잡한 구현: selective scan 의 CUDA 커널 구현이 까다로움.

Mamba1 vs Mamba2

• Mamba1 (Gu & Dao, 2023): 최초의 selective SSM. 성능 대비 효율 좋음.
• Mamba2 (Dao & Gu, 2024): SSD (structured state space duality) 로 Attention 과 연결. 더 큰 state 지원, 멀티헤드 적용, 하드웨어 효율 향상. EulerStack 기본값은 mamba2.

실제 적용 예

• Mamba (Gu & Dao, 2023) — SSM 단독 모델로 Transformer 와 동급 성능 증명
• Jamba-1.5 (AI21, 2024) — Mamba + Attention + MoE 로 398B/94B active 모델
• Zamba-2 (Zyphra, 2024) — Mamba 위주 + 주기적 attention 공유
• Samba (Microsoft, 2024) — Mamba + Sliding Window Attention
• Falcon-Mamba (TII, 2024) — 7B 순수 Mamba 모델
• Codestral Mamba (Mistral, 2024) — Mamba 기반 코드 생성 모델

어떤 상황에 유리한가?

EulerStack YAML 예시

layer_templates:
  mamba_layer:
    mixer:
      type: mamba
      mamba:
        variant: mamba2      # mamba1 | mamba2
        d_state: 128         # state 크기 (64/128/256 이 일반적)
        d_conv: 4            # 1D 컨볼루션 커널 크기
        expand: 2            # 내부 확장 비율
    ffn:
      type: gated_mlp
      activation: swiglu
    state:
      ssm_state: true        # 추론 시 SSM state 유지

하이브리드 스케줄 예시 (Jamba 스타일 3 mamba : 1 attention):

layer_schedule:
  - template: mamba_layer
    repeat: 3
  - template: attn_layer
    repeat: 1
  # ... 반복

논문 출처

• Gu & Dao, 2023. "Mamba: Linear-Time Sequence Modeling with Selective State Spaces."
• Dao & Gu, 2024. "Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality." (Mamba2)
• Lieber et al., 2024. "Jamba: A Hybrid Transformer-Mamba Language Model." (AI21)
• Glorioso et al., 2024. "Zamba: A Compact 7B SSM Hybrid Model." (Zyphra)