Mamba-2: 상태 공간 기반 시퀀스 모델

Mamba-2: SSM과 어텐션의 수학적 동치성을 증명한 SSD 프레임워크

Carnegie Mellon University / Princeton University · 2024-05-31 · SSM · Apache-2.0

개요

Mamba-2는 2024년 Albert Gu와 Tri Dao가 발표한 후속 연구로, SSM과 Attention 간의 수학적 동치성을 발견한 SSD(Structured State Space Duality) 프레임워크를 기반으로 한다. 이 연구의 핵심 기여는 이론적 통찰에 있다. 선택적 SSM이 특정 조건하에서 구조화된 마스크 어텐션(Structured Masked Attention)과 수학적으로 동일한 연산임을 증명한 것이다.

Mamba-1의 선택적 SSM을 상태 확장 차원 $N$이 헤드 차원과 분리된 구조로 재설계하여, SSD 알고리즘으로 더 빠르고 효율적인 학습이 가능해졌다. FlashAttention 스타일의 타일링 알고리즘을 SSM에 적용할 수 있게 되었으며, 동일 규모 Mamba-1 대비 훈련 효율 2~8배 향상을 달성했다.

이 이중성(duality)의 발견은 SSM과 Transformer라는 두 패러다임이 근본적으로 같은 연산의 서로 다른 관점임을 보여준다. 이론적으로는 두 분야의 알고리즘과 최적화 기법을 자유롭게 상호 차용할 수 있는 기반을 마련했으며, 실용적으로는 SSM 모델의 학습 효율을 극적으로 개선했다.

아키텍처 상세

SSD 프레임워크의 핵심은 선택적 SSM을 행렬 연산으로 재해석하는 것이다.

1-세미분리 행렬(1-Semiseparable Matrix)

선택적 SSM의 입출력 관계를 행렬 형태로 전개하면, 출력 $y_t={\sum }_{s\le t}{M}_{ts}{x}_s$ 형태가 된다. 여기서 마스크 행렬 $M$의 원소는 다음과 같다.

$M_{ts}=C_t^T\left({\prod }_{r=s+1}^t{\bar{A}}_r\right)B_s$

이 행렬 $M$은 1-세미분리(1-semiseparable) 구조를 가진다. 1-세미분리 행렬이란 하삼각(lower triangular) 부분의 모든 부분행렬의 rank가 최대 1인 행렬이다. 상태 차원 $N>1$이면 rank-$N$ 세미분리 행렬이 된다.

SSM-Attention 동치성

Transformer의 어텐션은 $Y=\text{softmax}(Q{K}^T)V$ 형태이고, SSM의 출력은 $Y=M\cdot X$ 형태이다. softmax를 제거하고 구조화된 마스크를 적용하면, 두 연산이 수학적으로 동일해진다.

$\underset{\text{SSM (recurrence)}}{\underbrace{Y=M\cdot X}}⟺\underset{\text{Structured Masked Attention}}{\underbrace{Y=(L\odot Q{K}^T)\cdot V}}$

여기서 $L$은 구조화된 하삼각 마스크이다. 이 동치성은 SSM의 상태 전이를 어텐션의 마스크 패턴으로, 어텐션의 QKV를 SSM의 B, C, $x$로 대응시킨다.

멀티-헤드 구조

Mamba-2 블록은 상태 차원 $N$을 헤드 수 $H$로 나누어 멀티-헤드 구조를 가진다. 각 헤드가 독립적인 SSM을 수행하며, Transformer의 Multi-Head Attention과 유사한 역할 분화가 일어난다. 이 설계로 텐서 병렬화(tensor parallelism)와 시퀀스 병렬화(sequence parallelism)가 Transformer처럼 자연스럽게 적용된다.

SSD 타일링 알고리즘

FlashAttention의 핵심 아이디어인 블록 단위 연산을 SSM에 적용했다. 시퀀스를 청크 크기 $C$로 분할하고, 청크 내부에서는 행렬 곱셈으로 병렬 계산, 청크 간에서는 순환적 상태 전파를 수행한다.

$\text{Intra-chunk: }{Y}_{\text{chunk}}=(L_{\text{chunk}}\odot Q_{\text{chunk}}{K}_{\text{chunk}}^T)V_{\text{chunk}}$ $\text{Inter-chunk: }{S}_{i+1}={\bar{A}}_{\text{chunk}}{S}_i+B_{\text{chunk}}^T{X}_{\text{chunk}}$

이 타일링은 GPU의 SRAM-HBM 메모리 계층을 최적화하여 실제 처리 속도를 크게 향상시킨다. 상태 크기를 Mamba-1 대비 8~16배 확장해도 효율적으로 동작한다.

핵심 혁신

Mamba-2의 핵심 혁신은 세 가지이다.

첫째, SSM-Attention 이중성 이론이다. 선택적 SSM과 구조화된 마스크 어텐션이 수학적으로 동일함을 증명하여, 두 분야의 알고리즘을 상호 차용할 수 있는 이론적 기반을 제공한다.

둘째, 확장된 상태 크기이다. Mamba-1의 $N=16$에서 $N=64\sim 256$까지 확장할 수 있게 되었으며, 상태 크기 증가가 성능 향상으로 직접 이어짐을 확인했다.

셋째, 텐서/시퀀스 병렬화 지원이다. 멀티-헤드 구조 덕분에 Transformer와 동일한 분산 학습 전략을 적용할 수 있다.

벤치마크/성능

모델 (2.7B)	Pile PPL↓	학습 속도	상태 크기
Mamba-2	6.18	1.5x Mamba-1	256
Mamba-1	6.22	1x	16
Transformer++	6.10	0.8x	N/A
Mamba-2 (N=64)	6.20	1.3x Mamba-1	64

모델	학습 알고리즘	멀티-헤드	상태 크기	이론적 기반
Mamba-2	SSD 타일링	예	256 (8~16x)	SSM-Attn 이중성
Mamba-1	Parallel scan	아니오	16	선택적 SSM
RetNet	청크 순환	예	고정	Retention
GLA	Chunkwise	예	행렬	게이트 선형 어텐션

학습

Pile 데이터셋으로 학습하며, GPT-NeoX 토크나이저를 사용한다. A100 80GB GPU 클러스터에서 SSD CUDA 커널로 메모리 효율을 최적화했다. Mamba-1과 동일한 학습 설정에서 2.7B 기준 훈련 속도 약 50% 향상을 달성했다. 상태 확장 크기 $N=64$를 기본값으로 사용한다.

다음은 SSD의 청크 단위 이중 계산(dual computation)을 보여주는 의사 코드이다.

import torch
 
def ssd_chunk_computation(Q, K, V, A_cumsum, chunk_size):
    """SSD 프레임워크의 intra-chunk 행렬 곱셈"""
    B, L, H, D = Q.shape
    n_chunks = L // chunk_size
    
    # 청크 단위로 분할
    Q_c = Q.reshape(B, n_chunks, chunk_size, H, D)
    K_c = K.reshape(B, n_chunks, chunk_size, H, D)
    V_c = V.reshape(B, n_chunks, chunk_size, H, D)
    
    # Intra-chunk: 행렬 곱셈 (어텐션 관점)
    attn = torch.einsum('bchd,bcjd->bchj', Q_c, K_c)  # QK^T
    
    # 구조화된 마스크 적용 (세미분리 행렬)
    mask = build_semiseparable_mask(A_cumsum, chunk_size)
    attn = attn * mask  # causal + decay mask
    
    # 출력 계산
    Y_intra = torch.einsum('bchj,bcjd->bchd', attn, V_c)
    
    # Inter-chunk: 순환 상태 전파
    states = compute_chunk_states(K_c, V_c, A_cumsum)
    Y_inter = apply_states(Q_c, states, A_cumsum)
    
    return Y_intra + Y_inter

관련 모델

Mamba-2는 state-spaces/mamba 저장소에서 Mamba-1과 동일한 인터페이스로 사용할 수 있다. SSD 프레임워크의 이론적 기여는 SSM과 어텐션을 통합하는 장기적 기반으로서 큰 영향을 미칠 것으로 예상된다. 순수 SSM의 in-context retrieval 한계는 여전히 존재하며, 이는 Mamba-3에서 소수의 어텐션 레이어를 추가하는 하이브리드 접근법으로 해결을 시도하고 있다.

참고 자료

• 논문: Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality
• 코드: state-spaces/mamba

관련 문서

• Mamba: Linear-Time Sequence Modeling with Selective State Spaces — 발전 기반
• Mamba-3 — 후속 모델