LLM : 논문 리뷰 : Mamba: Linear-Time Sequence Modeling with Selective State Spaces

용어 설명

• Structured State Space Models (SSMs): 연속적인 시스템을 이산적인 sequence 데이터에 맞게 변환하여 모델링하는 기법입니다. 기존의 RNN과 CNN의 장점을 결합하여 linear-time으로 동작하지만, 주로 연속적인 데이터(audio, vision)에 강점을 보였습니다.
• Linear Time Invariance (LTI): 시간에 따라 시스템의 dynamics(파라미터)가 변하지 않는 특성입니다. 기존의 모든 효율적인 SSM들(S4, H3 등)은 연산 효율성을 위해 LTI 제약을 가졌으며, 이로 인해 입력 데이터의 내용에 따라 동적으로 반응하지 못하는 한계가 있었습니다.
• Selective State Spaces (Selection Mechanism): 이 논문에서 제안한 핵심 개념으로, SSM의 파라미터가 입력 데이터(input)의 내용에 따라 동적으로 변하게 만드는 메커니즘입니다. 이를 통해 모델이 중요한 정보는 기억하고 불필요한 정보는 잊는 content-based reasoning이 가능해집니다.
• Hardware-aware Parallel Algorithm / Scan: Selection mechanism을 도입하면 기존 SSM이 사용하던 효율적인 convolution 연산을 더 이상 사용할 수 없습니다. 이를 해결하기 위해 GPU의 메모리 계층 구조(SRAM과 HBM)를 고려하여, 메모리 I/O를 최소화하고 빠른 병렬 처리를 가능하게 한 커널 퓨전(kernel fusion) 기반의 알고리즘입니다.
• Information-dense data: 언어(language)나 유전체(DNA)와 같이 특정 discrete token 하나하나가 고밀도의 정보를 담고 있어, 문맥(context)에 따른 선택적 정보 처리가 필수적인 데이터 모달리티를 뜻합니다.

---

Purpose of the Paper

이 논문은 Foundation Models의 표준으로 자리 잡은 Transformer의 근본적인 한계(sequence length에 대한 quadratic scaling 및 유한한 context window)를 극복하고자 합니다.
기존에도 이를 해결하기 위해 linear attention, gated CNN, LTI 기반의 SSMs 등 다양한 subquadratic-time 모델들이 제안되었으나, 이들은 language와 같은 information-dense data에서 Transformer의 성능에 미치지 못했습니다.
본 연구는 기존 효율적 모델들이 실패한 원인이 '입력 데이터의 내용에 따라 정보를 선택적으로 처리하는 능력(content-based reasoning)의 부재'에 있다고 정의합니다. 따라서, 효율성(linear scaling)을 유지하면서도 Transformer 수준의 문맥 이해 능력을 갖춘 새로운 시계열 모델 아키텍처(Mamba)를 제안하는 것이 주된 목적입니다.

---

Key Contributions and Novelty

• Selection Mechanism 도입을 통한 LTI 제약 극복
  • Contribution: SSM의 핵심 파라미터(Delta, B, C)를 input의 함수로 만들어, 모델이 sequence를 따라가며 특정 token의 정보를 선택적으로 전파(propagate)하거나 잊도록(forget) 만들었습니다.
  • Novelty: 연산 속도를 위해 반드시 LTI(시간 불변성)를 유지해야 한다는 기존 SSM 연구들의 고정관념을 깨고, time-varying(시간 가변성) 모델을 제안하여 discrete data에 대한 reasoning 능력을 획기적으로 부여했습니다.
• Hardware-aware Parallel Algorithm 설계
  • Contribution: Convolution 연산이 불가능해진 selective SSM을 위해, GPU SRAM에서 discretization과 recurrent scan을 융합(fusion)하여 처리하고 중간 상태를 HBM에 저장하지 않는 최적화 알고리즘을 개발했습니다.
  • Novelty: 이론적으로만 존재하던 time-varying SSM의 병목 현상을 GPU 메모리 계층 구조를 활용한 엔지니어링으로 해결하여, 최적화된 Attention 메커니즘(FlashAttention)보다 오히려 빠른 속도를 달성했습니다.
• 단순화된 Mamba Architecture 제안
  • Contribution: Attention이나 MLP block 없이, 기존 H3 block의 설계와 MLP block의 역할을 하나로 통합한 단순하고 균일한(homogenous) SSNN (State Space Neural Network) 아키텍처를 고안했습니다.
  • Novelty: 복잡한 하이브리드 구조(Attention + SSM) 없이, 오직 Selective SSM 단일 모듈만으로 Transformer를 대체할 수 있는 강력한 backbone을 제시했습니다.

---

Experimental Highlights

• Synthetic Tasks (능력 검증): Selective Copying 및 Induction Heads 태스크에서 Mamba는 완벽한 성능을 보였으며, 특히 훈련 시 본 적 없는 1M (백만) 길이의 sequence에서도 완벽하게 extrapolation에 성공했습니다. (기존 LTI 모델이나 Transformer는 실패)
• Language Modeling (Scaling Laws & Zero-shot):
  • 130M에서 1.3B 파라미터 규모의 scaling law 실험에서, Mamba는 강력한 최신 Transformer++ 레시피(LLaMa 구조 등 적용)와 동등 이상의 성능을 최초로 달성한 attention-free 모델입니다.
  • Mamba-3B 모델은 동일 크기의 Transformer를 압도하며, 2배 크기인 7B 모델(Pythia-7B)과 맞먹는 downstream zero-shot 성능을 기록했습니다.
  • Inference 성능: KV cache가 필요 없는 recurrent 구조 덕분에 Transformer 대비 5배 높은 generation throughput을 달성했습니다.
• DNA Modeling (Long Context): HyenaDNA 모델을 능가하며, sequence length를 1M까지 늘려도 성능(perplexity)이 지속적으로 향상되는 결과를 보였습니다. Great Apes classification 같은 고난이도 태스크에서도 SOTA를 달성했습니다.
• Audio Generation: SC09 speech generation 데이터셋에서 기존 SOTA였던 SaShiMi 모델뿐만 아니라, 훨씬 크기가 큰 GAN 및 Diffusion 기반 모델들보다 압도적으로 우수한 충실도(fidelity)를 달성했습니다.

---

Limitations and Future Work

• Limitations:
  • 평가 규모의 한계: 실험이 최대 3B 파라미터 규모에서 진행되었으므로, 현재 open-source LLM들의 주류인 7B 이상의 대규모 스케일에서도 Mamba가 Transformer를 능가하거나 필적할 수 있는지는 아직 실증되지 않았습니다.
  • Continuous-Discrete Spectrum의 Trade-off: Selection mechanism은 discrete data(Text, DNA)에 강력하지만, 완전히 연속적이고 부드러운 신호(순수 audio waveform 등)에서는 오히려 LTI 모델의 inductive bias가 더 유리할 수 있다는 결과가 나타났습니다.
• Future Work:
  • 대규모 확장 (Scaling): Mamba 아키텍처를 7B 파라미터 이상으로 확장(scaling)하여 대규모 Foundation Model로서의 역량을 입증해야 합니다.
  • Downstream Affordances 연구: Transformer 기반 LLM 생태계에 존재하는 다양한 기술들(instruction tuning, RLHF, prompting, adaptation 등)이 Mamba와 같은 SSM 기반 모델에도 동일하고 유효하게 적용될 수 있는지에 대한 심도 있는 연구가 필요합니다.

---

Overall Summary

이 논문은 기존 Sequence 모델들이 가지던 효율성과 표현력 사이의 딜레마를 해결하기 위해, 입력 데이터에 따라 중요도를 판별하는 Selection Mechanism을 결합한 Selective State Space Model인 Mamba를 제안했습니다. 연산 효율성 저하라는 기술적 한계는 하드웨어 친화적인 병렬 스캔 알고리즘(Hardware-aware Algorithm)을 통해 극복하였으며, 결과적으로 O(N)의 선형적인 메모리 및 연산 복잡도를 유지합니다. Mamba는 언어, 유전체, 오디오 등 다양한 모달리티에서 Transformer에 필적하거나 이를 뛰어넘는 SOTA 성능과 5배 빠른 추론 속도를 보여주었으며, 향후 대형 AI 모델의 핵심 backbone을 Transformer에서 Attention-free 아키텍처로 전환할 수 있는 매우 강력한 대안을 제시했다는 점에서 큰 의의를 가집니다.

---

쉬운 설명

Transformer가 문서를 읽는 방식은 **"책의 전체 페이지를 한 번에 사진 찍어두고 필요할 때마다 모든 단어들의 관계를 동시에 대조해 보는 방식"**과 같습니다. 성능은 뛰어나지만 책이 두꺼워질수록 메모리(KV cache)가 터지고 속도가 기하급수적으로 느려집니다. 반면 기존의 효율적인 모델(LTI SSM)들은 "단어를 하나씩 순서대로 읽는 컨베이어 벨트" 같아서 빠르긴 하지만, 중요한 주인공 이름이나 의미 없는 "음", "어" 같은 추임새를 모두 똑같은 비중으로 취급해버려 문맥을 제대로 파악하지 못했습니다.

Mamba는 이 둘의 장점을 합친 **"매우 똑똑한 속독가"**입니다. 단어를 하나씩 순서대로 읽어 속도가 빠르고 메모리를 적게 차지하면서도, 글을 읽으면서 **"이 단어는 중요하니까 기억(Select)하고, 이 단어는 쓸데없으니까 잊어버리자(Forget)"**라고 스스로 판단합니다. 게다가 이 속독가는 뇌의 단기 기억(GPU SRAM)을 기가 막히게 잘 활용해서, 읽는 도중에 굳이 가방(GPU HBM)을 열어 노트를 꺼내는 시간을 없애버렸기 때문에 엄청난 속도로 책을 읽고 이해할 수 있습니다.

 

 

 

 

 

 

 

 

 

B L D 형태의 입력이 필요함 그게 입력데이터

학습파라미터는 고정된 A랑 여러 리니어 프로젝션 가중치들

입력데이터를 리니어 프로젝션 이후 투배럭으로 돌림.
한쪽은 게이트쪽, 한쪽은 메인 ssm연산 통로

근데 아직 통과 안함.

메인 ssm으로 갈 데이터쪽에 1d 컨볼루션 하나 때림
지역적인 정보 좀 섞어줌

컨브 통과한 데이터를 3개의 리니어 레이어에 각각 때림. 그걸 B C 델타로 사용.
이때 B C 델타는 고정된 값이 아니라 현재 들어온 토큰의 내용에 따라 매번 달라짐 왜? 컨브 때리고 했으니까. 그것보단 리니어 레이어를 때리니까 데이터마다

그래서 이제 어쩌라는거냐? 여기서 델타를 써서 A랑 B에[ 곱해서 
A바랑 B바를 만듬

그래서 이제 ssm 쭉쭉 하듯 h업데이트 하면됨 유명한 공식으로
일단 근데 쭉쭉 안하고 h만들고 아까 만든 C를 곱해줌 그래서 y를 출력

아까 빼둔 게이트 통로 있지? 그거 통과한 값이랑 y를 요소별 곱해서 값을 뽑고 마지막에 리니어 프로젝션 한번 더 때림. 그래서 최종 아웃풋 텐서

그리고 쭉쭉 지나서 마지막에 loss 구해지면 쭉 역전파