Mamba 논문의 이름은 Mamba: Linear-Time Sequence Modeling with Selective State Spaces다. (링크)
저자는 Albert Gu, Tri Dao다.
Mamba는 State Space Model을 활용한 본격적인 NLP 모델이다.
Absract
현재 딥러닝 분야의 흥미로운 애플리케이션 대부분을 구동하는 Foundation 모델은 거의 대부분 Transformer 아키텍처와 그 핵심인 어텐션 모듈을 기반으로 한다. Linear Attention, gated convolution, recurrent models, structured state space model (SSM)과 같은 많은 subquadratic-time 아키텍처가 긴 시퀀스에서 Transformer의 계산 비효율성을 해결하기 위해 개발되었지만, 언어와 같은 중요한 모달리티에서는 어텐션만큼 좋은 성능을 보이지 못했다. 이러한 모델의 주요 약점은 내용 기반 추론을 수행할 수 없다는 점임을 확인하고 몇 가지 개선을 시도했다. 첫째, SSM 매개변수를 입력 함수로 설정함으로써 discrete modalities 이산 모달리티의 약점을 해결하고, 모델이 현재 토큰에 따라 시퀀스 길이 차원을 따라 정보를 선택적으로 propagate 전파하거나 forget 잊을 수 있도록 한다. 둘째, 이러한 변경으로 효율적인 컨볼루션을 사용할 수 없게 되더라도, recurrent 모드에서 하드웨어 인식 병렬 알고리즘을 설계한다. 이러한 selective SSM을 어텐션이나 MLP 블록 (Mamba) 없이 단순화된 end-to-end 신경망 아키텍처에 통합한다. Mamba는 빠른 추론 (Transformers보다 5배 높은 처리량)과 시퀀스 길이의 linear scaling 선형 확장을 제공하며, 최대 백만 길이의 시퀀스까지 실제 데이터에서 성능이 향상되었다. Mamba는 일반적인 시퀀스 모델 백본으로서 language, audio, and genomics 등 다양한 분야에서 최첨단 성능을 구현한다. 언어 모델링에서 Mamba-3B 모델은 사전 학습 및 다운스트림 평가 모두에서 동일 크기의 Transformer보다 성능이 뛰어나고, 두 배 크기의 Transformer와 유사하다.
3. Selective State Space Models
Selection mechanism 선택 메커니즘에 대한 동기를 synthetic tasks 합성 작업에 대한 직관 (섹션 3.1)을 통해 제시하고, 이 메커니즘을 state space model (섹션 3.2)에 어떻게 통합하는지 설명한다. 결과적으로 생성되는 time-varying 시간 가변 SSM은 컨볼루션을 사용할 수 없으므로 효율적인 계산 방법에 대한 기술적 과제가 발생한다. 저자들은 현대의 하드웨어의 메모리 계층을 활용하는 하드웨어 인식 알고리즘 (섹션 3.3)으로 이 문제를 해결한다. 그런 다음, 어텐션이나 MLP 블록조차 없는 간단한 SSM 아키텍처(섹션 3.4)를 설명한다. 마지막으로, 선택 메커니즘의 몇 가지 추가 속성(섹션 3.5)에 대해 논의한다.
3.1. Motivation: Selection as Means of Compression 압축 수단으로서의 선택
저자들은 시퀀스 모델링의 근본적인 문제가 컨텍스트를 더 작은 상태로 압축하는 것이기 때문이라 주장한다. 사실, 이 관점에서 인기 있는 시퀀스 모델들의 장단점을 살펴볼 수 있다. 예를 들어, 어텐션은 컨텍스트를 전혀 압축하지 않기 때문에 효과적이면서도 비효율적이다. 이는 autoregressive 자기회귀적 추론이 전체 컨텍스트를 명시적으로 저장해야 한다는 사실 (즉, KV 캐시)에서 알 수 있으며, 이는 트랜스포머의 느린 선형 시간 추론과 quadratic-time 훈련을 직접적으로 유발한다. 반면에, recurrent models 은 유한한 state 상태를 가지므로 효율적이며 이는 상수 시간 추론과 선형 시간 훈련을 의미한다. 그러나 그 효과는 상태가 컨텍스트를 얼마나 잘 압축했는지에 따라 제한된다.
이 원리를 이해하기 위해 두 가지 합성 작업 예시 (아래 Figure 2)에 초점을 맞춘다.
- The Selective Copying 선택적 복사 작업은 인기 있는 복사 작업(Arjovsky, Shah, and Bengio 2016)을 수정하여, 기억해야 할 토큰의 위치를 변화시킨다. 이 작업은 관련 토큰 (색깔이 있는)을 기억하고 관련 없는 토큰 (흰색)을 걸러내기 위해 content-aware reasoning이 필요하다.
- Induction Heads 작업은 LLM의 in-context learning (ICL) 능력의 대부분을 설명하는 것으로 추정되는 잘 알려진 메커니즘이다. 이 작업은 적절한 컨텍스트 (검은색)에서 올바른 출력을 언제 생성해야 하는지 알기 위해 context-aware reasoning이 필요하다.
이러한 작업들은 LTI (Liner Time Invariant) 시간 불변 모델의 실패 방식을 드러낸다. Recurrent 관점에서 볼 때, 고정된 dynamics (e.g., (2)의 ($\bar{A}, \bar{B}$) 변환)은 컨텍스트에서 올바른 정보를 선택하거나, 입력에 의존하는 방식으로 시퀀스를 따라 전달되는 은닉 상태에 영향을 줄 수 없다. 컨볼루션적 관점에서는, global convolutions이 오직 time-awareness만 요구하는 기본적인 복사 작업은 해결할 수 있지만, content-awareness이 부족하기 때문에 선택적 복사 작업에는 어려움을 겪는다는 것이 알려져 있다 (Figure 2). 더 구체적으로 말해, 입력과 출력 사이의 간격이 변하므로 정적인 컨볼루션 커널로는 모델링할 수 없다.
요약하자면, 시퀀스 모델의 효율성과 효과성 간의 장단점은 상태를 얼마나 잘 압축하는지에 의해 결정된다. 효율적인 모델은 작은 상태를 가져야 하고, 효과적인 모델은 컨텍스트로부터 필요한 모든 정보를 포함하는 상태를 가져야 한다. 따라서 시퀀스 모델을 구축하기 위한 근본적인 원리가 selectivity 선택성, 즉 순차적 상태로 들어오는 입력을 집중하거나 걸러내는 컨텍스트 인식 능력이라고 제안한다. 특히, 선택 메커니즘은 시퀀스 차원을 따라 정보가 어떻게 전파되거나 상호작용하는지를 제어하며 더 자세한 논의는 섹션 3.5을 참조하면 된다.
3.2 선택으로 SSM 개선
선택 메커니즘을 모델에 통합하는 한 가지 방법은 시퀀스를 따라 상호작용에 영향을 미치는 매개변수(예: RNN의 순환 dynamics 또는 CNN의 컨볼루션 커널)를 입력에 의존하도록 만드는 것이다.
위 알고리즘 1과 2는 우리가 사용하는 주요 선택 메커니즘을 보여준다. 주요 차이점은 단순히 여러 매개변수인 Δ,B,C를 입력의 함수로 만들고, 이에 따른 텐서 모양의 변화에 대한 전체적인 적용이다. 특히, 이 매개변수들이 길이 차원 $L$을 가지게 되어 모델이 time-invariant 시간 불변에서 time-varying 시간 가변으로 변경됨을 강조한다. 이는 컨볼루션과의 등가성 (3)을 상실하여 효율성에 영향을 미치는데, 이에 대해서는 다음 섹션에서 논의한다.
구체적으로는 $s_B(x)$ = $\text{Linear}_N (x)$, $s_C(x)$ = $\text{Linear}_N (x)$, $s_\Delta (x)$ = $ \text{Broadcast}_D \text{Linear}_1 (x)$, and $\tau_\Delta$ = softplus. 그리고 $\text{Linear}_d$는 dimension $d$로 매개변수화된 projection이다. $s_\Delta (x)$ 와 $\tau_\Delta$의 선택은 RNN의 게이팅 메커니즘과 관련있으며 섹션 3.5에서 설명한다.
HiPPO 논문을 전에 리뷰 (링크)한적 있는데, 기본적으로 State Space Model의 $A, B$ 행렬들은 모두 time invariant, 즉 시간에 따라서 변하는 내용들이 아니다. 그래서 이 구조를 sequence length를 이용해서 time variant로 변경해준다는 이야기다.
3.4 단순화된 SSM 아키텍처
Structured SSM처럼, 선택적 SSM은 신경망에 유연하게 통합될 수 있는 독립적인 시퀀스 변환입니다. 가장 잘 알려진 SSM 아키텍처 (섹션 2)의 기반이 되는 H3 아키텍처는 일반적으로 linear attention에서 영감을 받은 블록과 MLP 블록을 번갈아 가며 구성된다. 저자들은 이 두 구성 요소를 하나로 결합하여 아키텍처를 단순화하고, 이를 homogeneously 동질적으로 쌓아 올렸다(아래 Figure 3). 이는 어텐션에 대해 유사한 작업을 수행했던 GAU (gated attention unit)에서 영감을 받았다.
이 아키텍처는 모델 차원 $D$를 제어 가능한 확장 계수 $E$만큼 확장한다. 각 블록에서 대부분의 매개변수($3ED^2$)는 linear projection (input projectiobn을 위한 $2ED^2$, output projection을 위한 $ED^2$)에 있으며, 내부 SSM은 적은 양을 차지한다. SSM 매개변수 (Δ,B,C를 위한 projection과 matrix A)의 수는 상대적으로 훨씬 작다. 이 블록을 standard normalization 및 residual connections과 번갈아 가며 반복하여 Mamba 아키텍처를 형성한다. 실험에서는 항상 $E=2$로 고정하고, 블록을 두 겹으로 쌓아 트랜스포머의 MHA (멀티 헤드 어텐션)와 MLP 블록을 번갈아 사용하는 $12D^2$ 매개변수와 일치시킨다. 우리는 SiLU / Swish 활성화 함수를 사용하는데, 이는 Gated MLP가 널리 사용되는 “SwiGLU” 변형이다. 마지막으로, RetNet이 비슷한 위치에 정규화 레이어를 사용한 것에서 영감을 받아선택적인 정규화 레이어 (저자들은 LayerNorm을 선택함)를 추가로 사용한다.
3.5. 선택 메커니즘의 속성
선택 메커니즘은 더 넓은 개념으로, 전통적인 RNN이나 CNN에 적용되거나, 다른 파라미터 (예: 알고리즘 2의 $A$)에 적용되거나, 다른 변환 $s(x)$를 사용하는 등 다양한 방식으로 적용될 수 있다.
3.5.1 게이팅 메커니즘과의 연결
가장 중요한 연결점을 강조하자면, RNN의 고전적인 게이팅 메커니즘은 SSM을 위한 저자들의 선택 메커니즘의 예시 중 하나다. RNN 게이팅과 연속 시간 시스템의 이산화 사이의 연결은 잘 확립되어 있다. 실제로, Theorem 1 정리 1은 ZOH 이산화 및 input-dependent gatese 입력 의존적 게이트로 일반화한 Gu, Johnson, Goel, et al. 2021 논문의 Lemma 3.1 의 개선 사항이다 (증명은 Appendix C). 더 넓게 보자면, SSM의 Δ는 RNN 게이팅 메커니즘의 일반화된 역할을 한다고 볼 수 있다. 이전 연구와 마찬가지로, 저자들은 SSM의 이산화가 heuristic 휴리스틱 게이팅 메커니즘의 원칙적인 기반이라는 관점을 채택한다.
Theorem 1 정리 1.
$N$=1,$A$=−1,$B$=1,$s_\Delta$ = Linear($x$) 그리고 $\tau_\Delta = \text{softplus}$일 때, 선택적 SSM 재귀 (알고리즘 2)는 다음과 같은 형태를 취한다.
$g_t = \sigma$ (Linear($x_t$))
$h_t = (1 - g_t) h_{t-1} + g_t x_t$
섹션 3.2에서 언급되었듯이, 이러한 연결에서 비롯되어 $s_\Delta$와 $\tau_\Delta$에 대한 저자들의 특정한 선택이 이루어진다.
특히, 주어진 입력 $x_t$가 완전히 무시되어야 하는 경우 (합성 작업에서 필요하듯이), 모든 $D$ 채널이 이를 무시해야 하므로, 저자들은 Δ를 사용하여 반복 / 브로드캐스트하기 전에 입력을 1차원으로 투영한다는 점에 주목해야 한다.
3.5.2 선택 메커니즘의 해석
선택의 세 가지 특정 메커니즘적 효과에 대해 자세히 설명한다.
Variable Spacing. 선택성은 관심 있는 입력들 사이에 발생할 수 있는 관련 없는 noise tokens 잡음 토큰을 필터링할 수 있게 한다. 이는 선택적 복사 작업에서 잘 나타나지만, 일반적인 데이터 양식, 특히 이산 데이터에서 흔히 발생한다. 예를 들어, “um”과 같은 language fillers의 존재가 그렇다. 이 속성은 모델이 특정 입력 $x_t$ 를 메커니즘적으로 필터링할 수 있기 때문에 발생하며, 예를 들어 게이트 RNN의 경우 (정리 1) $g_t$ → 0일 때 발생한다.
Filtering Context. 엄밀히 말해서, 더 많은 컨텍스트는 더 나은 성능으로 이어져야 한다는 원칙에도 불구하고, 많은 시퀀스 모델은 컨텍스트가 길어질수록 성능이 향상되지 않는다는 것이 경험적으로 관찰되었다. 한 가지 설명은 많은 시퀀스 모델들이 필요할 때 관련 없는 컨텍스트를 효과적으로 무시할 수 없다는 것이다. 직관적인 예시는 글로벌 컨볼루션 (및 일반적인 LTI 모델)다. 반면에, 선택적 모델은 언제든지 상태를 간단히 재설정하여 불필요한 이력을 제거할 수 있으므로, 원칙적으로 그들의 성능은 컨텍스트 길이에 따라 monotonoicly improve 단조 향상된다 (예: 섹션 4.3.2).
Boundary Resetting. 여러 독립적인 시퀀스가 함께 엮여 있는 환경에서, 트랜스포머는 특정 어텐션 마스크를 인스턴스화하여 시퀀스를 분리할 수 있지만, LTI 모델은 시퀀스 간에 정보가 섞일 수 있다. 선택적 SSM도 경계에서 상태를 재설정할 수 있습니다 (예: $\Delta_t$ → $\infty$, 또는 정리 1에서 $g_t$ → 0 일때. 이러한 상황은 인위적으로 (예: 하드웨어 활용도를 높이기 위해 문서를 함께 묶는 것) 또는 자연적으로(예: 강화 학습의 에피소드 경계) 발생할 수 있다.
또한, 아래에서는 각 선택 매개변수의 효과에 대해 자세히 설명한다.
Δ의 해석. 일반적으로 Δ는 현재 입력 $x_t$ 에 얼마나 집중하거나 무시할지에 대한 균형을 제어힌다. 이는 RNN 게이트를 일반화한다(예: 정리 1의 $g_t$). 큰 Δ는 상태 $h$를 재설정하고 현재 입력 $x$에 집중하는 반면, 작은 Δ는 상태를 유지하고 현재 입력을 무시한다. SSM (1) - (2)는 timestep Δ에 의해 이산화된 연속 시스템으로 해석될 수 있으며, 이 맥락에서 큰 Δ → $\infty$는 시스템이 현재 입력에 더 오래 집중하는 것을 (따라서 "선택"하고 현재 상태를 잊어버리는 것) 나타내는 반면, 작은 Δ → 0는 무시되는 일시적인 (transient) 입력을 나타낸다는 직관을 가진다.
$A$의 해석. $A$ 파라미터 또한 선택적일 수 있지만, 궁극적으로는 $A$=exp(Δ$A$)(이산화 (4))를 통해 Δ와의 상호작용을 통해서만 모델에 영향을 미친다. 따라서 Δ의 선택성만으로도 $(A, B)$의 선택성을 보장하기에 충분하며, 이것이 개선의 주요 원인이다. 우리는 $A$를 Δ에 추가로 (또는 대신에) 선택적으로 만드는 것이 비슷한 성능을 가질 것이라고 가정하고, 단순함을 위해 $A$는 선택적이지 않게 두었다.
$B$와 $C$의 해석. 3.1절에서 논의했듯이, 선택성의 가장 중요한 속성은 관련 없는 정보를 걸러내어 시퀀스 모델의 컨텍스트를 효율적인 상태로 압축할 수 있게 한다는 점이다. SSM에서 $B$와 $C$를 선택적으로 수정하면 입력 $x_t$를 상태 $h_t$로 받아들일지, 또는 상태를 출력 $y_t$로 내보낼지에 대해 더 세밀한 제어가 가능하다. 이는 모델이 각각 콘텐츠 (입력)와 컨텍스트 (은닉 상태)를 기반으로 recurrent dynamics을 조절할 수 있도록 해석될 수 있다.
4. Empirical Evaluation
4.1에서 2가지 synthetic tasks, 4.2에서 언어 모델 사전 학습을, 4.3에서 DNA 시퀀스 모델 사전 학습을, 4.4에서는 오디오 웨이브폼 사전 학습을 수행한다. 4.5는 맘바의 학습과 추론에서의 계산적 효율성을 보여주며 4.6에서는 albation study를 보여준다.
4.1. Synthetic Tasks
자세한 내용은 Appendix E.1에 나와있다.
4.1.1. Selective Copying
Copying tasks는 본래 recurrent models의 기억 능력을 평가하기 위한 태스크다. LTI SSM(선형 재귀 및 글로벌 컨볼루션)은 데이터를 추론하는 대신 시간을 추적하는 것만으로 이 작업을 쉽게 해결할 수 있다. 그러나 선택적 복사 작업은 토큰 사이의 간격을 무작위로 변경하여 이러한 shortcut을 막는다. 이 작업은 이전에 Denoising 작업으로 소개되기도 했다.
4.1.2. Induction Heads
Mechanistic interpretability의 관점에서 나온 간단한 작업으로, LLM의 in-context learning 능력을 놀라울 정도로 잘 예측한다. 이 작업은 모델이 associative recall 및 copy를 수행하도록 요구한다. 예를 들어, 모델이 시퀀스에서 "Harry Potter"와 같은 2음절 (bigram)을 본 적이 있다면, 같은 시퀀스에 "Harry"가 다음에 나타날 때, 모델은 과거 기록에서 복사하여 "Potter"를 예측할 수 있어야 한다.
데이터셋. Induction heads 작업에 대해 256의 시퀀스 길이와 16의 vocab size를 가진 2 layers 모델을 훈련시켰다. 이는 이전 연구와 비슷하지만 더 긴 시퀀스를 사용한다. 또한, 테스트 시점에는 $2^6 = 64$부터 $2^20 = 1,048,576$ 까지 다양한 시퀀스 길이에 대해 평가하여 일반화 및 extrapolation 능력을 조사했다.
모델. 기존 연구를 따라, 우리는 2 laeyers 모델을 사용했다. 멀티 헤드 어텐션(8개 헤드, 다양한 위치 인코딩 사용)과 SSM 변형 모델들을 모두 테스트했다. 모델 차원 $D$는 Mamba의 경우 64, 다른 모델의 경우 128로 설정했다.
결과. Table 2는 Mamba, 더 정확히는 그 선택적 SSM 레이어가 관련 없는 모든 것을 무시하면서 관련 토큰을 선택적으로 기억하는 능력 덕분에 이 작업을 완벽하게 해결할 수 있음을 보여준다. Mamba는 훈련 중에 본 시퀀스 길이보다 4000배 더 긴 백만 길이의 시퀀스까지 완벽하게 일반화하는 반면, 다른 어떤 방법도 2배 이상으로 일반화하지 못했다.
Table 1을 보면 Mamba에 S6를 적용했을 때 S4나 H3 보다 더 나은 Selective Copying 결과를 달성함을 알 수 있다.
Table 2에서는 Induction Heads의 성능을 다루는데 갈색으로 표시된 Mamba의 성능이 월등함을 알 수 있다.
4.2. Language Modeling
표준적인 autoregressive language modeling을 이용해서 pretraining metrics (perplexity)와 zero-shot evaluations에 대해서 평가한다. Model size (width와 depth)를 GPT3의 특징을 반영하여 설정했다. The Pile 데이터셋을 사용했다. 자세한 내용은 Appendix E.2에 나와있다.
4.2.1. Scaling Laws
Figure 4를 보면 Mamba 역시 Scaling law를 따름을 알 수 있다.
Perplexity가 스케일이 커질수록 줄어드는데 Sequence length가 2048일 때와 훨씬 더 긴 8192일 때 모두 작동한다.
4.2.2. Downstream Evaluations
여러가지 다운스트림 태스크의 Zero-shot 성능에 대해서 다룬다.
Pile에 대해서는 perplexity를, Lambada에 대해서는 perplexity와 accuracy를 측정한다.
그외에도 HellaSwag, PIQA, ARC-E, ARC-C, WinoGrande에 대해서 측정했다.
Hybrid H3와 Pythia가 주된 비교 대상이며 OPT, GPT-Neo, GTP-J, RWKV와도 비교했다.
4.3. DNA Modeling
Genomics에서 역시 언어 모델과 유사하게 discrete tokens의 시퀀스 형태로 DNA를 분석할 수 있다.
여기서도 언어 모델과 마찬가지로 standard causal language modeling (next token prediction)의 개념으로 학습한다.
자세한 내용은 Appendix E.2에 나와있다.
4.3.1 and 4.3.2. Scaling Laws
Figure 5를 보면 Mamba의 perplexity가 HyenaDNA와 Transformer++보다 더 낮음을 알 수 있다.
그리고 오른쪽을 보면 sequecne length가 증가하면서 HyenaDNA와 다르게 perplexity가 감소함을 알 수 있다.
4.3.3. Synthetic Species Classification
Classification 성능을 Figure 6에서 보여주는데 Mamba 7M이 가장 좋은 성능임을 알 수 있다.
4.4. Audio Modeling and Generation
웨이브폼의 오디오를 다룸에 있어서 저자들은 SaShiMi 아키텍처와 학습 프로토콜을 가장 우선적으로 비교한다.
위 Figure 7을 보면 S4 + FFN (SaShiMi 구조) 과 Mamba 모두 sequence length에 대한 scaling law가 성립함을 알 수 있다.
4.4.2. Autoregressive Speech Generation
SC09는 speech generation 데이터셋으로 1-second clips들이며 16,000 Hz로 0 부터 9까지의 digits로 구성된다.
아래 Table 4는 Mamba-Unet을 SampleRNN, WaveNet, SaShiMi, DiffWave 등과 성능을 비교한다.
Table 5에서는 S4 + MLP와 Mamba 구조의 ablation study 결과를 보여준다.
Mamba 구조가 가장 좋은 성능을 달성했음을 알 수 있다.
4.5. Speed and Memory Benchmarks
속도와 메모리 효율성에 대해서 다루는 섹션이다.
저자들의 Scan 방식이 Flash Attention 2 보다 더 적게 걸린다고 주장하는 점이 놀랍다.
오른쪽에서는 유사한 크기의 모델에 대해서 Mamba가 Transformer 보다 훨씬 더 많은 throughput 처리량을 가짐을 보여준다.
이때 사용한 GPU는 A100 80GB with prompt length 2048이다.
4.6. Model Ablations
Mamba 구조에 대한 ablation study 섹션이다.
H3, Hyena, S4, 그리고 $A$를 어떻게 초기화하는지 그리고 complex와 real 중에서 어떤게 좋은지를 연구한 결과를 보여준다.
References:
https://minyoungxi.tistory.com/118
https://www.ibm.com/think/topics/mamba-model
'NLP > LLM' 카테고리의 다른 글
| 41가지 로컬 LLM 벤치마크를 실행 결과 (1) | 2025.09.02 |
|---|---|
| LLM의 컨텍스트 윈도우와 관련된 글 (1) | 2025.09.01 |
| 도메인 특화 LLM 리서치 (4) | 2025.08.12 |
| MUVERA와 Mercury 리서치 (1) | 2025.07.15 |
| Mixtral (2024) 논문 리뷰 (0) | 2025.06.24 |
