Mamba는 2024년 Albert Gu와 Tri Dao가 제안한 모델입니다.
[Mamba]
https://arxiv.org/abs/2312.00752
Mamba: Linear-Time Sequence Modeling with Selective State Spaces
Foundation models, now powering most of the exciting applications in deep learning, are almost universally based on the Transformer architecture and its core attention module. Many subquadratic-time architectures such as linear attention, gated convolution
arxiv.org
[LSSL]
https://arxiv.org/abs/2110.13985
Combining Recurrent, Convolutional, and Continuous-time Models with Linear State-Space Layers
Recurrent neural networks (RNNs), temporal convolutions, and neural differential equations (NDEs) are popular families of deep learning models for time-series data, each with unique strengths and tradeoffs in modeling power and computational efficiency. We
arxiv.org
[S4]
https://arxiv.org/abs/2111.00396
Efficiently Modeling Long Sequences with Structured State Spaces
A central goal of sequence modeling is designing a single principled model that can address sequence data across a range of modalities and tasks, particularly on long-range dependencies. Although conventional models including RNNs, CNNs, and Transformers h
arxiv.org
[S4D]
https://arxiv.org/abs/2206.11893
On the Parameterization and Initialization of Diagonal State Space Models
State space models (SSM) have recently been shown to be very effective as a deep learning layer as a promising alternative to sequence models such as RNNs, CNNs, or Transformers. The first version to show this potential was the S4 model, which is particula
arxiv.org
[H3]
https://arxiv.org/abs/2212.14052
Hungry Hungry Hippos: Towards Language Modeling with State Space Models
State space models (SSMs) have demonstrated state-of-the-art sequence modeling performance in some modalities, but underperform attention in language modeling. Moreover, despite scaling nearly linearly in sequence length instead of quadratically, SSMs are
arxiv.org
<Mamba>
* Reference
https://www.youtube.com/watch?v=JjxBNBzDbNk
https://www.youtube.com/watch?v=N6Piou4oYx8&t=25s
* Abstract
1. 기반 모델은 현재 딥러닝에서 가장 흥미로운 응용 분야를 지원하며, 거의 대부분이 Transformer 아키텍처와 그 핵심 구성 요소인 attention 모듈에 기반을 두고 있습니다. 긴 시퀀스에서 Transformer의 계산 비효율성을 해결하기 위해 linear attention, gated convolution 및 recurrent models, structured state space models(SSMs) 등 여러 subquadratic-time 아키텍처가 개발되었으나, 이러한 모델들은 언어와 같은 중요한 모달리티에서 어텐션만큼 좋은 성능을 보이지 못했습니다.
2. 우리는 이러한 모델들의 주요 약점이 내용 기반 추론(content-based reasoning)을 수행하지 못하는 점이라고 확인하고, 이를 개선하기 위해 몇 가지 혁신을 제안합니다.
- 입력에 따라 SSM 파라미터를 함수화하여 discrete 모달리티에서의 약점을 극복합니다. 이 접근은 모델이 현재 토큰에 따라 시퀀스 길이 차원에서 정보를 선택적으로 전달하거나 잊어버릴 수 있도록 만듭니다.
- 이 변화로 인해 효율적인 컨볼루션 사용은 제한되지만, 우리는 recurrent mode에서 하드웨어를 고려한 병렬 알고리즘을 설계합니다.
3. 우리는 이러한 선택적 SSM을 간단한 end-to-end 신경망 아키텍처(Mamba)에 통합하며, 어텐션이나 MLP 블록조차 사용하지 않습니다. Mamba는 Transformers보다 5배 높은 처리량과 시퀀스 길이에 선형적으로 확장되는 스케일링을 자랑하며, 실제 데이터에서 최대 백만 길이의 시퀀스까지 성능이 향상됩니다.
4. Mamba는 일반적인 시퀀스 모델 백본으로서 언어, 오디오, 유전체학 등 여러 모달리티에서 SOTA를 달성합니다. 특히 언어 모델링에서, Mamba-3B 모델은 같은 크기의 Transformer보다 우수한 성능을 보이며, 크기가 두 배인 Transformer와 동등한 성능을, pretraining과 downstream 평가 모두에서 보여줍니다.
* Transformer vs Recurrent Model
- Transformer: 이전까지의 sequence 정보를 따로 선별해서 저장 X → O(N^2)의 시간 및 메모리 복잡도
- Recurrent Model: 이전까지의 sequence 정보를 따로 선별해서 저장 O → 저장 성능 문제, 병렬 계산 불가
* What is SSM(State Space Model)?
- from 제어 공학 (https://www.youtube.com/watch?v=8Ps_CjKxYD8)
- 시간에 따라 변화하는 데이터를 모델링
- State Equation, Output Equation
x'(t): 시간에 따른 상태 x(t)의 변화율
A: 상태 전이 행렬 (State Transition Matrix)
B: 입력 행렬 (Input Matrix)
u(t): 외부 입력 (Input)
y(t): 출력 (Output)
C: 출력 행렬 (Output Matrix)
D: 입력-출력 간 직접적인 관계를 나타내는 행렬 (생략 가능)
- Simplified SSM, Discretized Equation (How to discretized: Eulter method, ZOH(Zero-Order-Hold), ...)
* LSSL → S4 → S4D → H3
1. LSSL(Learnable State-Space Layer): SSM을 딥러닝에 도입하려는 초기 시도
→ State Transition Matrix의 중요성
2. S4(Structured State Space Sequence Model): 병렬화와 효율성 증대 (O(N), O(NlogN))
→ by HIPPO matrix
3. S4D: State Transition Matrix의 대각화를 통한 계산 단순화 → bottleneck 완화
4. H3: Language model에서의 Structured SSM의 성능 향상을 위해 무엇이 부족한가?
→ Induction Head, Associative Recall task → Conv1D Layer 추가
* Mamba's Key Idea
- Selective Copying Task, Induction Heads Task
- Key Idea: Selectivity, Parallel Scan, Kernel Fusion
1. Selectivity
* Problem: 기존 LTI(Linear Time-Invariant, 선형 시간 불변 시스템) 모델의 한계
→ recurrent 관점에서: LTI 모델의 constant dynamics (ex. A_bar, B_bar)는 context에서 올바른 정보를 선택하거나 입력 의존적으로 은닉 상태를 변경할 수 없음
→ convolution 관점에서: Selective Copying Task에서는 입력과 출력 간 간격이 변화하여, 이를 정적 convolution kernel로 모델링 할 수 없음
* Solution
→ 각 token마다 B, C, delta가 다를 수 있게 하자.
2. Parallel Scan
* Problem: 더 이상 convolutionize 불가
* Solution
→ Parallel Scan: 연산의 Associativity 이용
→ O(N) 유지
3. Kernel Fusion
* Problem: B x L x D x N 차원의 tensor에 의해 발생하는 메모리 부담
→ B x L x D 차원으로 운용되는 Transformer에 비해 부담 증가
* Solution
→ 실제 SRAM에서의 연산보다 HBM에서 SRAM으로 복사하는 시간이 오래 걸리는 것
→ Discretization을 SRAM에서 바로 진행
* Mamba's Algorithm
- SSM에서는 B, C, delta가 learnable 했다면, SSM + Selection에서는 s_B, s_C, s_delta 가 learnable 하다.
* Mamba's Architecture
- 기존 SSM 아키텍처와 Transformer의 MLP 블록의 결합
* Empirical Evaluation
Mamba 모델의 성능 평가: 5가지 실험
- Synthetic Tasks – Mamba가 Selective SSM의 강점을 활용하여 시퀀스 데이터를 얼마나 잘 처리하는지 평가
- Language Modeling – GPT-3 수준의 Transformer와 비교하여 Mamba의 성능 평가
- DNA Modeling – 긴 시퀀스 길이에서의 성능과 일반화(Scaling) 평가
- Audio Modeling & Generation – 음성 및 음악 데이터에서 Transformer 기반 모델 대비 성능 비교
- Speed & Memory – Transformer와 비교하여 Mamba의 효율성 검증
1. Synthetic Task
→ Selective Copying: Mamba가 가장 좋은 성능을 보임 (타 아키텍처에서도 S6 layer 사용 시 성능 대폭 향상)
→ Induction Heads: Mamba는 훈련 데이터보다 4,000배 더 긴 시퀀스에서도 완벽하게 일반화 (100만 토큰 이상까지 확장 가능)
→ 기존 LTI 기반 SSM보다 훨씬 강력하며, Transformer보다 훨씬 긴 문맥에서도 강한 성능
2. Language Modeling
→ Scaling Laws: Transformer 수준의 성능을 달성한 최초의 Attention-Free 모델
→ Zero-Shot Evaluation: 동일 크기 모델에서 항상 가장 좋은 성능을 보임
3. DNA Modeling
→ Scaling-Model Size: Mamba는 40M 규모의 가장 큰 모델에서 Transformer++ 대비 3~4배 적은 파라미터로도 동일한 PPL 기록
→ Scaling Across Context Length: 1M 길이 시퀀스까지 긴 문맥 활용 가능, 문맥이 길어질수록 성능이 향상
→ Species Classification: HyenaDAN 대비 높은 정확도 기록
4. Audio Modeling & Generation
→ Waveform Modeling: BPB(Bits Per Byte, NLL Loss의 변형) 측정 결과, Mamba가 SaShiMi 대비 높은 성능
→ Autoregressive Speech Generation: 소형 Mamba 모델조차 대형 GAN 및 Diffusion 모델보다 더 높은 품질의 음질 생성 가능
5. Speed & Memory Efficiency
→ SSM Scan: FlashAttention-2 대비 시퀀스 길이 2K 이상에서 더 빠르고, PyTorch의 기본 Scan 연산 대비 20~40배 빠름
→ 추론 처리량: Transformer보다 4~5배 높은 처리량 달성
* Discussion
1. No Free Lunch: Continuous-Discrete Spectrum
Structured SSM은 원래 연속 시스템을 discretization한 형태로 정의되었으며, 오디오와 비디오와 같은 지각적 신호와 같은 연속 시간 데이터 모달리티에 강한 귀납적 편향(inductive bias)을 가지고 있습니다. 3.1절과 3.5절에서 논의된 바와 같이, selection mechanism은 텍스트와 DNA와 같은 discrete modalities에서의 약점을 극복합니다.
하지만, 이와 반대로 selection mechanism은 LTI SSM이 뛰어난 성능을 보이는 데이터에 대해서는 성능을 저해할 수 있습니다. 오디오 파형(audio waveforms)에 대한 우리의 실험적 분석(ablation study)은 이러한 트레이드오프를 더 자세히 탐구합니다.
2. Downstream Affordances
Transformer 기반의 파운데이션 모델(특히 대규모 언어 모델, LLM)은 미세 조정(fine-tuning), 적응(adaptation), 프롬프팅(prompting), 인컨텍스트 학습(in-context learning), 명령 튜닝(instruction tuning), RLHF(강화 학습을 통한 인적 피드백), 양자화(quantization) 등 다양한 속성과 상호작용 방식으로 구성된 풍부한 생태계를 보유하고 있습니다. 우리는 SSM과 같은 Transformer 대안 모델들이 이러한 속성과 기능을 얼마나 제공할 수 있는지에 특히 관심이 있습니다.
3. Scaling
우리의 실험적 평가는 소규모 모델 크기에 국한되어 있으며, 이는 Llama(Touvron et al., 2023)와 같은 강력한 오픈 소스 LLM이나 RWKV(B. Peng et al., 2023), RetNet(Y. Sun et al., 2023)과 같은 기타 순환 모델들이 평가된 7B(70억) 이상의 파라미터 스케일보다 낮은 수준입니다. Mamba가 이러한 더 큰 크기에서 여전히 경쟁력을 유지할 수 있는지는 추가적인 평가가 필요합니다. 또한, SSM의 확장은 이 논문에서 다루지 않은 추가적인 엔지니어링 과제와 모델 조정이 필요할 수 있음을 언급합니다.