[논문리뷰] Mamba: Selective State Space Model로 Transformer의 한계를 돌파하다

들어가며

2023년 12월, Albert Gu와 Tri Dao가 발표한 “Mamba: Linear-Time Sequence Modeling with Selective State Spaces” 논문은 딥러닝 시퀀스 모델링 분야에 큰 파장을 일으켰습니다. Transformer가 NLP, 비전, 오디오 등 거의 모든 영역을 지배하던 시점에서, 선형 시간 복잡도로 동작하면서도 Transformer에 필적하거나 능가하는 성능을 보여주는 새로운 아키텍처를 제안했기 때문입니다.

이 논문은 기존 Structured State Space Model(S4)의 핵심 한계였던 콘텐츠 기반 추론(content-based reasoning) 능력 부족을 선택적 메커니즘(selection mechanism)으로 해결하면서, 동시에 GPU 하드웨어에 최적화된 효율적 구현까지 제공합니다.

핵심 한 줄 요약: Mamba는 입력에 따라 SSM 파라미터를 동적으로 변화시키는 “선택적 SSM”을 제안하여, Attention 없이도 Transformer 수준의 성능을 선형 시간 복잡도로 달성한 모델입니다.

논문 기본 정보

배경: 왜 Transformer를 넘어서야 하는가?

Transformer의 근본적 한계

Transformer는 Self-Attention 메커니즘을 통해 시퀀스 내 모든 토큰 간의 관계를 직접 모델링합니다. 이 접근법은 강력하지만 두 가지 근본적인 문제를 안고 있습니다:

1. 이차 시간 복잡도: Attention의 계산 복잡도는 시퀀스 길이 $L$에 대해 $O(L^2)$입니다. 시퀀스가 길어질수록 연산량이 급격히 증가합니다.
2. 선형 메모리 증가: KV Cache를 사용하는 추론 시, 시퀀스 길이에 비례하여 메모리가 증가합니다.

이로 인해 긴 시퀀스(수만~수십만 토큰)를 처리할 때 실질적인 병목이 발생합니다.

기존 SSM(S4)의 한계

Structured State Space Model(S4)은 연속 시간 시스템을 이산화하여 시퀀스를 모델링하는 접근법으로, 선형 시간 복잡도와 긴 시퀀스 처리 능력을 보여주었습니다. S4의 핵심 수식은 다음과 같습니다:

$$h'(t) = \mathbf{A}h(t) + \mathbf{B}x(t)$$
$$y(t) = \mathbf{C}h(t)$$

여기서:
– $h(t) \in \mathbb{R}^N$: 은닉 상태(hidden state)
– $x(t) \in \mathbb{R}$: 입력 신호
– $y(t) \in \mathbb{R}$: 출력 신호
– $\mathbf{A} \in \mathbb{R}^{N \times N}$: 상태 전이 행렬
– $\mathbf{B} \in \mathbb{R}^{N \times 1}$, $\mathbf{C} \in \mathbb{R}^{1 \times N}$: 입력/출력 투영 행렬

이를 이산화하면 (Zero-Order Hold 기준):

$$h_t = \bar{\mathbf{A}} h_{t-1} + \bar{\mathbf{B}} x_t$$
$$y_t = \mathbf{C} h_t$$

여기서 $\bar{\mathbf{A}} = \exp(\Delta \mathbf{A})$, $\bar{\mathbf{B}} = (\Delta \mathbf{A})^{-1}(\exp(\Delta \mathbf{A}) – \mathbf{I}) \cdot \Delta \mathbf{B}$이고, $\Delta$는 이산화 스텝 크기입니다.

하지만 기존 S4에서 $\mathbf{A}$, $\mathbf{B}$, $\mathbf{C}$, $\Delta$는 모두 입력과 무관한 고정 파라미터입니다. 이것이 핵심 한계입니다.

문제: 고정된 파라미터는 입력 내용에 따라 어떤 정보를 기억하고 어떤 정보를 잊을지 동적으로 결정할 수 없습니다. 이는 Selective Copying이나 Induction Heads 같은 콘텐츠 기반 추론 태스크에서 치명적입니다.

핵심 기여 (Key Contributions)

Mamba의 핵심 기여는 크게 세 가지로 정리할 수 있습니다:

1. Selection Mechanism (선택적 메커니즘)

기존 SSM의 파라미터 $\mathbf{B}$, $\mathbf{C}$, $\Delta$를 입력에 의존하는 함수로 만듭니다. 즉, 시퀀스의 각 토큰마다 파라미터가 달라집니다.

2. Hardware-Aware Algorithm (하드웨어 인식 알고리즘)

입력 의존적 파라미터로 인해 효율적인 convolution 연산을 사용할 수 없게 되었지만, GPU 메모리 계층(SRAM/HBM)을 활용한 새로운 병렬 스캔 알고리즘으로 이를 보완합니다.

3. Simplified Architecture (간소화된 아키텍처)

Attention 블록과 MLP 블록을 하나로 통합한 간결한 아키텍처를 제안하여, 모델 전체를 동질적인 Mamba 블록의 스택으로 구성합니다.

방법론 상세 분석

Selection Mechanism: SSM을 입력 의존적으로 만들기

Mamba의 핵심 아이디어는 놀라울 정도로 단순합니다. SSM의 파라미터를 입력의 함수로 만드는 것입니다:

$$\mathbf{B}_t = \text{Linear}_N(x_t)$$
$$\mathbf{C}_t = \text{Linear}_N(x_t)$$
$$\Delta_t = \text{softplus}(\text{Linear}_1(x_t) + \text{Parameter})$$

여기서:
– $\text{Linear}_N$: 입력 $x_t$를 $N$차원으로 투영하는 선형 변환
– $\text{softplus}$: $\log(1 + e^x)$ 함수로 \Delta_t > 0을 보장
– $N$: 상태 차원 (state dimension)

핵심 직관은 $\Delta_t$의 역할에 있습니다:

• $\Delta_t$가 클 때: $\bar{\mathbf{A}}_t \approx \mathbf{I}$ (이전 상태를 거의 그대로 유지), $\bar{\mathbf{B}}_t$도 커져서 현재 입력을 강하게 반영 → 현재 입력에 집중
• $\Delta_t$가 작을 때: $\bar{\mathbf{A}}_t$가 감쇠 효과를 가지며, $\bar{\mathbf{B}}_t$도 작아져서 현재 입력을 무시 → 현재 입력을 건너뜀

직관적 해석: Selection Mechanism은 모델이 시퀀스를 처리하면서 “이 토큰은 중요하니 기억하자” vs “이 토큰은 무시하자”를 동적으로 결정할 수 있게 해줍니다. 이는 Attention이 하는 역할과 본질적으로 유사하지만, 고정 크기의 은닉 상태를 통해 압축된 형태로 수행합니다.

Selective Copying 태스크로 보는 Selection의 효과

논문은 Selective Copying 태스크로 Selection Mechanism의 필요성을 명확히 보여줍니다:

• Copying 태스크: 입력 시퀀스의 특정 토큰들을 그대로 복사하여 출력 (토큰 간 간격이 고정)
• Selective Copying 태스크: 토큰 간 간격이 가변적 (무작위로 배치)

시간 불변 SSM은 간격이 고정된 패턴만 학습할 수 있지만, 선택적 SSM은 입력 내용을 보고 동적으로 대응할 수 있습니다.

Hardware-Aware Parallel Scan Algorithm

선택적 SSM에서 파라미터가 입력 의존적이 되면, 기존 S4의 핵심 효율성 기법인 FFT 기반 convolution을 사용할 수 없게 됩니다. 왜냐하면 convolution은 시간 불변 시스템에서만 적용 가능하기 때문입니다.

Mamba는 이 문제를 세 가지 전략으로 해결합니다:

1) Kernel Fusion

이산화($\Delta \rightarrow \bar{\mathbf{A}}, \bar{\mathbf{B}}$), 상태 업데이트, 출력 계산을 하나의 GPU 커널로 합칩니다. 이를 통해 중간 결과를 HBM(High Bandwidth Memory, 느림)에 쓰지 않고 SRAM(빠름)에서 모두 처리합니다.

2) Parallel Scan

순환 연산 $h_t = \bar{\mathbf{A}}<em>t h</em>{t-1} + \bar{\mathbf{B}}_t x_t$을 parallel associative scan 알고리즘으로 병렬화합니다. 이는 prefix sum의 일반화로, $O(L)$ 연산을 $O(\log L)$ 단계로 병렬 처리합니다.

3) Recomputation

역전파 시 중간 상태를 저장하지 않고 재계산합니다. 이는 메모리 사용량을 크게 줄이며, FlashAttention에서 사용한 것과 같은 전략입니다.

# Mamba의 핵심 선택적 SSM 연산 (의사 코드)
def selective_ssm(x, A, D):
    """
    x: (batch, length, d_model) - 입력 시퀀스
    A: (d_inner, d_state) - 상태 전이 행렬 (학습 가능)
    D: (d_inner,) - skip connection 파라미터
    """
    batch, length, d_model = x.shape

    # 입력 의존적 파라미터 생성
    B = linear_projection_B(x)  # (batch, length, d_state)
    C = linear_projection_C(x)  # (batch, length, d_state)
    delta = softplus(linear_projection_delta(x))  # (batch, length, d_inner)

    # 이산화
    A_bar = torch.exp(delta.unsqueeze(-1) * A)  # (batch, length, d_inner, d_state)
    B_bar = delta.unsqueeze(-1) * B.unsqueeze(2)  # (batch, length, d_inner, d_state)

    # Selective Scan (parallel scan으로 구현)
    h = torch.zeros(batch, d_inner, d_state)  # 초기 은닉 상태
    outputs = []
    for t in range(length):  # 실제로는 parallel scan으로 병렬화
        h = A_bar[:, t] * h + B_bar[:, t] * x[:, t].unsqueeze(-1)
        y_t = (h * C[:, t].unsqueeze(1)).sum(-1)  # (batch, d_inner)
        outputs.append(y_t)

    y = torch.stack(outputs, dim=1)  # (batch, length, d_inner)
    y = y + x * D  # Skip connection
    return y

Mamba Block Architecture

Mamba는 Transformer의 Attention + MLP 구조를 하나의 블록으로 통합합니다:

입력 x
├── Branch 1: Linear → Conv1d → SiLU → Selective SSM
└── Branch 2: Linear → SiLU
    ↓
    Element-wise Multiply (Gate)
    ↓
    Linear → 출력

각 구성 요소의 역할:

아키텍처의 핵심 특징: Mamba는 Attention 블록 없이 동질적(homogeneous) 블록만으로 전체 네트워크를 구성합니다. H3 아키텍처(SSM + Attention 혼합)와 달리 순수 SSM만 사용하면서도 더 나은 성능을 달성합니다.

이론적 분석: Selection과 기존 메커니즘의 연결

논문은 Selection Mechanism이 기존의 여러 메커니즘과 연결됨을 보여줍니다:

1. Variable-Length Gating (가변 길이 게이팅)

$\Delta$의 역할은 RNN의 게이트와 유사합니다:

$$h_t = (1 – \sigma(z_t)) \cdot h_{t-1} + \sigma(z_t) \cdot x_t \quad \text{(GRU 단순화)}$$

$$h_t = \bar{\mathbf{A}}<em>t h</em>{t-1} + \bar{\mathbf{B}}_t x_t \quad \text{(Selective SSM)}$$

GRU에서 $\sigma(z_t)$가 하는 역할을 Selective SSM에서 $\Delta_t$가 수행합니다. 하지만 SSM은 다차원 상태 공간($N$차원)과 연속-이산 변환 구조를 통해 더 풍부한 표현력을 가집니다.

2. Attention과의 관계

Attention은 모든 이전 토큰에 직접 접근하는 반면, Selective SSM은 압축된 은닉 상태를 통해 간접적으로 접근합니다:

실험 결과 분석

1. 언어 모델링 (Language Modeling)

가장 중요한 실험인 언어 모델링에서 Mamba는 인상적인 결과를 보여줍니다.

The Pile 벤치마크 Perplexity 비교

핵심 발견: Mamba는 모든 스케일에서 동일 파라미터 수의 Transformer++를 능가합니다. 특히 Mamba-1.4B가 Transformer++-1.3B보다 낮은 Perplexity를 기록하며, 이는 Mamba-1.4B가 약 2배 큰 Transformer++-2.8B와 비슷한 수준입니다.

Downstream Zero-Shot 평가

Mamba는 거의 모든 다운스트림 태스크에서 Transformer++와 RWKV를 앞서며, 특히 LAMBADA와 HellaSwag에서 큰 차이를 보입니다.

2. DNA 모델링 (Genomics)

유전체 시퀀스는 매우 긴 시퀀스(수십만 bp)를 다뤄야 하므로 SSM의 강점이 두드러집니다.

주목할 점: Transformer는 메모리 부족(OOM)으로 긴 시퀀스를 처리하지 못하는 반면, Mamba는 시퀀스가 길어질수록 성능이 크게 향상됩니다. 이는 긴 범위의 의존성을 효과적으로 캡처함을 보여줍니다.

3. 오디오 모델링 및 생성

SpeechCommands (음성 분류) 태스크에서도 Mamba는 경쟁력 있는 결과를 보입니다.

Ablation Study 분석

1. Selection Mechanism의 영향

논문의 가장 중요한 Ablation은 어떤 파라미터를 선택적으로 만들 것인가에 대한 분석입니다:

인사이트: $\Delta$의 선택성이 가장 큰 영향을 미칩니다. $\Delta$는 입력을 “얼마나 강하게 반영할지”를 제어하는 게이트 역할을 하기 때문입니다. $\mathbf{B}$, $\mathbf{C}$를 추가로 선택적으로 만들면 성능이 더 향상됩니다.

2. 상태 차원 $N$의 영향

$N = 16$에서 대부분의 이점을 얻으며, 그 이상에서는 수확체감이 나타납니다.

3. 추론 속도 벤치마크

Mamba의 추론 속도는 시퀀스가 길어질수록 Transformer 대비 이점이 커지며, 생성(generation) 시 처리량이 최대 5배 높습니다. 이는 Mamba가 RNN처럼 고정 크기 상태만으로 추론하기 때문입니다.

Mamba의 추론 효율성: 왜 빠른가?

Mamba의 효율성을 이해하려면 학습과 추론의 이중 모드를 이해해야 합니다:

학습 (Training)

• Parallel Scan을 사용하여 전체 시퀀스를 병렬 처리
• 시간 복잡도: $O(L \cdot D \cdot N)$ (시퀀스 길이 $L$, 모델 차원 $D$, 상태 차원 $N$)
• GPU 메모리 계층을 활용한 kernel fusion으로 실제 wallclock time 최적화

추론 (Inference / Generation)

• 순환(Recurrent) 모드로 토큰 하나씩 생성
• 상태 크기: $O(D \cdot N)$ — 시퀀스 길이와 무관
• KV Cache 불필요 → 메모리 사용량 일정

# Mamba 추론 모드 (토큰 단위 생성)
def mamba_inference_step(x_t, h_prev, model_params):
    """
    x_t: 현재 토큰 임베딩 (d_model,)
    h_prev: 이전 은닉 상태 (d_inner, d_state) - 고정 크기!
    """
    # 1. 입력 의존적 파라미터 계산
    B_t = model_params.linear_B(x_t)  # (d_state,)
    C_t = model_params.linear_C(x_t)  # (d_state,)
    delta_t = softplus(model_params.linear_delta(x_t))  # (d_inner,)

    # 2. 이산화
    A_bar_t = torch.exp(delta_t.unsqueeze(-1) * model_params.A)  # (d_inner, d_state)
    B_bar_t = delta_t.unsqueeze(-1) * B_t.unsqueeze(0)  # (d_inner, d_state)

    # 3. 상태 업데이트 (O(1) 연산!)
    h_new = A_bar_t * h_prev + B_bar_t * x_t.unsqueeze(-1)

    # 4. 출력 계산
    y_t = (h_new * C_t.unsqueeze(0)).sum(-1)  # (d_inner,)

    return y_t, h_new  # 상태 크기는 시퀀스 길이와 무관

이 이중 모드 덕분에 Mamba는 학습 시 GPU 병렬성을 최대한 활용하면서, 추론 시에는 RNN의 효율성을 누립니다.

기존 방법론과의 종합 비교

강점 분석

1. 이론적 우아함

Selection Mechanism은 기존 SSM의 가장 근본적인 한계를 정확히 짚고, 최소한의 변경으로 해결합니다. $\mathbf{B}$, $\mathbf{C}$, $\Delta$를 입력 의존적으로 만드는 것만으로 질적 도약을 이루어낸 것은 연구의 깊이를 보여줍니다.

2. 실용적 효율성

FlashAttention의 공저자인 Tri Dao가 참여한 만큼, 하드웨어 최적화가 단순한 이론이 아닌 실제 구현으로 검증되었습니다. Kernel fusion, parallel scan, recomputation 전략은 이론적 복잡도를 실제 속도로 변환하는 데 핵심적입니다.

3. 범용성

언어, DNA, 오디오 등 다양한 모달리티에서 일관된 성능 향상을 보여주며, Foundation Model의 backbone으로서의 가능성을 입증합니다.

4. 확장성 (Scaling)

125M부터 1.4B까지의 스케일에서 일관된 개선을 보여주며, 더 큰 스케일에서도 성능이 유지될 것으로 기대됩니다.

한계점 분석

1. In-Context Learning 능력의 불확실성

Mamba는 압축된 은닉 상태를 통해 정보를 전달하므로, 명시적 검색(retrieval)이 필요한 태스크에서 Attention 대비 불리할 수 있습니다. 예를 들어 “문서에서 특정 사실을 찾아 답하라”는 유형의 태스크에서 Attention은 관련 토큰에 직접 접근하지만, Mamba는 압축된 상태에서 해당 정보가 보존되어야 합니다.

2. 실험 스케일의 제한

가장 큰 모델이 1.4B 파라미터로, 현대 LLM 기준(7B, 13B, 70B)으로는 작은 편입니다. 대규모에서의 행동을 예측하기 어렵습니다.

3. 커스텀 CUDA 커널 의존성

효율적 구현을 위해 전용 CUDA 커널이 필요하며, 이는 PyTorch 네이티브 구현 대비 접근성이 떨어집니다. 새로운 하드웨어(TPU 등)로의 이식에도 추가 작업이 필요합니다.

4. Attention과의 상보성 미탐구

논문은 순수 Mamba 아키텍처에 집중하지만, Attention과 SSM을 결합한 하이브리드 접근법이 더 효과적일 수 있다는 가능성은 충분히 탐구되지 않았습니다. (이후 Jamba 등의 후속 연구에서 이 방향이 탐구됨)

5. 양방향(Bidirectional) 처리의 제한

Mamba는 기본적으로 단방향(causal) 모델입니다. BERT처럼 양방향 정보를 활용해야 하는 태스크(분류, NLU 등)에서는 추가적인 설계가 필요합니다.

후속 연구 방향

Mamba 이후 SSM 분야는 폭발적으로 성장하고 있습니다:

1. Mamba-2 (2024)

Albert Gu와 Tri Dao는 후속 연구에서 Structured State Space Duality (SSD)를 제안합니다. 이는 SSM과 Attention의 수학적 동치 관계를 밝히고, 더 효율적인 알고리즘과 2~8배 빠른 학습 속도를 달성합니다.

2. 하이브리드 아키텍처

• Jamba (AI21 Labs): Mamba + Attention + MoE를 결합한 하이브리드 모델
• Zamba (Zyphra): Mamba 블록과 공유 Attention 블록을 교대 배치

이러한 하이브리드 모델들은 Mamba의 효율성과 Attention의 명시적 검색 능력을 동시에 활용합니다.

3. Vision Mamba

• Vim (Vision Mamba): ViT를 대체하는 비전 백본
• VMamba: 2D 선택적 스캔을 적용한 비전 모델
• U-Mamba: 의료 영상 세그멘테이션에 Mamba 적용

4. 대규모 학습

• Falcon Mamba (TII): 순수 Mamba 기반 7B 모델
• NVIDIA의 SSM 연구: 대규모 학습에서의 SSM 행동 분석

마무리

Mamba는 Transformer 지배 시대에 근본적인 대안을 제시한 논문입니다. 핵심 내용을 정리하면:

1. Selection Mechanism은 SSM의 파라미터를 입력 의존적으로 만들어, 콘텐츠 기반 추론 능력을 부여합니다. $\Delta$, $\mathbf{B}$, $\mathbf{C}$를 입력의 함수로 만드는 단순하지만 효과적인 아이디어입니다.

2. 하드웨어 인식 알고리즘으로 이론적 효율성을 실제 속도 향상으로 전환합니다. Kernel fusion, parallel scan, recomputation 전략이 핵심입니다.

3. 선형 시간 복잡도와 고정 크기 상태로 추론 시 Transformer 대비 최대 5배 높은 처리량을 달성하며, 긴 시퀀스에서 특히 강점을 보입니다.

4. 언어, DNA, 오디오 등 다양한 모달리티에서 Transformer에 필적하거나 능가하는 성능을 입증하여, 범용 시퀀스 모델로서의 가능성을 보여주었습니다.

5. 이후 Mamba-2, Jamba, Vision Mamba 등 후속 연구의 토대가 되어, SSM 기반 모델의 새로운 생태계를 열었습니다.

이 논문이 중요한 이유: Attention이 아닌 다른 메커니즘으로도 Foundation Model을 구축할 수 있다는 것을 실증적으로 증명했습니다. Transformer가 유일한 정답이 아니라는 사실은, 앞으로의 아키텍처 연구에 더 넓은 가능성을 열어줍니다.