Mamba Linear-Time Sequence Modeling with Selective State Spaces

🧐 나의 생각 / 비판 (My Thoughts / Critiques)

• transformer 모델 쓸때는 확실히 대용량 시계열은 부담스러웠다.
• Mamba-2 가 나왔다고 하니 그거를 더 봐야겠다.

---

요약

• 핵심 제안: 기존 상태 공간 모델(SSM)에 입력 데이터에 따라 동적으로 파라미터가 변하는 ‘선택 메커니즘’을 도입한 ‘선택적 SSM(Selective SSM)’ 과 이를 효율적으로 연산하는 하드웨어 친화적 알고리즘을 제안함.

• 주요 성과: 언어, 오디오, 유전체 등 여러 분야에서 트랜스포머와 대등하거나 능가하는 최고 수준(SOTA)의 성능을 달성했으며 , 시퀀스 길이에 선형적으로 확장되고 추론 속도는 5배 빠른 효율성을 입증함.

• 의의: 트랜스포머의 2차 복잡도 문제를 해결하면서도 높은 성능을 유지하여, 긴 시퀀스 처리가 필수적인 다양한 분야의 차세대 파운데이션 모델 백본(backbone)이 될 수 있는 강력한 대안을 제시함.

---

1. 목표 (Goal)

이 연구의 주된 목표는 기존 AI 모델의 주류인 ‘트랜스포머(Transformer)‘의 근본적인 비효율성 문제를 해결하는 새로운 모델 아키텍처, Mamba를 제안하는 것입니다. 트랜스포머는 시퀀스(데이터) 길이가 길어질수록 계산량이 제곱으로 증가하여 매우 긴 데이터를 처리하기 어렵습니다. 이 연구는 트랜스포머의 높은 성능은 유지하면서도, 시퀀스 길이에 따라 계산량이 선형적으로 증가하는 효율적인 모델을 개발하고자 합니다. 이를 위해 기존의 효율적인 모델들이 가졌던 ‘입력 내용에 따라 정보를 선별하는 능력’의 부재라는 약점을 극복하는 데 중점을 둡니다.

---

2. 핵심 아이디어 (Key Idea)

Mamba의 혁신은 ‘선택적 상태 공간 모델(Selective State Space Model, S6)’ 이라는 핵심 아이디어에 기반합니다.

• 선택 메커니즘 (Selection Mechanism): 기존의 상태 공간 모델(SSM)은 입력 데이터에 상관없이 항상 동일한 규칙(파라미터)을 적용하여(시점 불변, LTI) 문맥을 파악하는 능력이 부족했습니다. Mamba는 모델의 핵심 파라미터(A, B, C, Δ)가 입력되는 토큰에 따라 동적으로 변화하도록 설계했습니다. 이 덕분에 모델은 현재 입력된 정보의 중요도를 판단하여, 필요한 정보는 기억(전파)하고 불필요한 정보는 잊어버릴(필터링) 수 있습니다.

• 하드웨어 친화적 알고리즘 (Hardware-aware Algorithm): 위와 같은 ‘선택’ 기능을 도입하면 기존의 빠른 계산 방식(컨볼루션)을 사용할 수 없게 됩니다. 이를 해결하기 위해, GPU 메모리 계층 구조의 특성을 활용한 새로운 병렬 스캔 알고리즘을 개발했습니다. 이 알고리즘은 계산에 필요한 데이터를 GPU의 느린 메모리(HBM)에서 매우 빠른 메모리(SRAM)로 옮겨 한 번에 처리함으로써, 입출력 병목 현상을 최소화하고 높은 연산 속도를 달성합니다.

---

3. 모델 구조 (Model Architecture)

Mamba는 ‘선택적 상태 공간 모델’을 기반으로 매우 간결하고 균일한(homogeneous) 구조를 가집니다.

• 통합된 Mamba 블록: 트랜스포머가 어텐션 블록과 MLP 블록을 번갈아 쌓는 구조인 것과 달리, Mamba는 이 두 기능을 하나의 ‘Mamba 블록’ 으로 통합했습니다.

• 블록 구성: 각 Mamba 블록은 입력값을 받아 차원을 확장하고(Expansion), 1D 컨볼루션과 SiLU 활성화 함수를 거친 뒤, 핵심적인 선택적 SSM(S6) 레이어를 통과시켜 시퀀스 정보를 처리합니다. 마지막으로 다시 원래 차원으로 투영(projection)하는 구조를 가집니다.

• 전체 구조: 이 단순화된 Mamba 블록을 여러 층으로 반복해서 쌓고, 잔차 연결(residual connection)과 정규화를 추가하여 전체 네트워크를 구성합니다.

---

4. 주요 성과 (Key Achievements)

• 성과 1: 선형 확장성 및 빠른 추론 속도: Mamba는 훈련 및 추론 시 계산량과 메모리 사용량이 시퀀스 길이에 따라 선형적으로 증가합니다. 특히, 이전 대화 기록을 저장할 필요가 없어(No KV Cache), 비슷한 크기의 트랜스포머보다 5배 높은 추론 처리량을 달성했습니다.

• 성과 2: 언어 모델링에서 트랜스포머급 성능 달성: Mamba는 선형 시간 복잡도를 가진 모델 중 최초로 강력한 트랜스포머 모델과 대등한 성능을 보였습니다. Mamba-3B 모델은 사전 훈련 및 다운스트림 평가 모두에서 두 배 크기의 트랜스포머 모델과 맞먹는 성능을 기록했습니다.

• 성과 3: 다양한 분야(오디오, 유전체)에서 최고 수준 성능: Mamba는 언어뿐만 아니라 오디오, 유전체(DNA)와 같은 여러 데이터 유형에서 최고 수준(SOTA)의 성능을 달성했습니다. 특히, 백만 토큰 길이에 달하는 DNA 시퀀스에서 다른 모델과 달리 컨텍스트가 길어질수록 성능이 향상되는 모습을 보였습니다.

• 성과 4: 핵심 능력 입증: LLM의 핵심 능력으로 여겨지는 ‘귀납적 추론(Induction Heads)‘과 같은 합성 과제(synthetic task)에서, 훈련 시 본 것보다 수천 배 긴 시퀀스에 대해서도 완벽하게 일반화하는 능력을 보여주며 ‘선택’ 메커니즘의 우수성을 입증했습니다.