Mamba: Linear-Time Sequence Modeling with Selective State Spaces

Albert Gu; Tri Dao

논문 · arXiv 2312.00752

Mamba: 선택적 상태 공간을 활용한 선형 시간 시퀀스 모델링

Mamba: Linear-Time Sequence Modeling with Selective State Spaces

Albert Gu, Tri Dao

게시일 2023-12-01cs.LGcs.AI arXiv PDF

XYZNOT로 읽기 (전문 번역과 AI 해설)

세 줄 요약

🎯 Transformers는 시퀀스 길이에 따라 이차(quadratic)로 증가하는 비용과 자기회귀 추론 중 대규모 KV 캐시를 필요로 하는 반면, 기존 LTI 상태 공간 모델과 같은 준이차 모델들은 언어 모델링처럼 이산적인 내용 추론 작업에서 어려움을 겪습니다.
🔧 SSM 파라미터를 입력 의존적(선택적)으로 만들어 정보를 동적으로 필터링하거나 전파하며, 하드웨어 인지 융합 병렬 결합 스캔 및 재계산 알고리즘을 결합하여 확장된 은닉 상태가 GPU HBM에 구체화되는 것을 방지합니다.
📈 Mamba-3B는 다운스트림 제로샷 태스크에서 두 배 크기의 Transformers와 일치하거나 능가하는 성능(예: Pythia-2.8B의 59.1%, Pythia-6.9B의 61.7% 대비 평균 63.3%)을 보이는 동시에 최대 5배 높은 추론 처리량을 제공합니다.

초록

현재 딥러닝의 혁신적인 응용 분야를 주도하는 파운데이션 모델들은 거의 전적으로 Transformer 아키텍처와 그 핵심인 attention 모듈에 기반을 두고 있습니다. 선형 어텐션, 게이티드 합성곱 및 순환 모델, 구조화된 상태 공간 모델(SSM) 등 많은 준이차(subquadratic) 시간 아키텍처들이 긴 시퀀스에 대한 Transformer의 연산 비효율성을 해결하기 위해 개발되었지만, 언어와 같은 중요한 모달리티에서는 어텐션만큼 뛰어난 성능을 보이지 못했습니다. 우리는 이러한 모델들의 핵심 약점이 내용 기반 추론(content-based reasoning)을 수행하지 못한다는 점임을 밝혀내고 몇 가지 개선점을 제시합니다. 첫째, SSM 파라미터가 입력을 기반으로 계산되도록 함으로써 이산적 모달리티에서 나타나는 약점을 해결하여, 모델이 현재 토큰에 따라 시퀀스 길이를 따라 정보를 선택적으로 전파하거나 잊을 수 있도록 합니다. 둘째, 이러한 변화로 인해 효율적인 합성곱(convolution)을 사용할 수 없게 됨에도 불구하고, 우리는 순환 모드에서 하드웨어 인지적 병렬 알고리즘을 설계했습니다. 우리는 이러한 선택적 SSM을 어텐션이나 심지어 MLP 블록조차 없는 단순화된 엔드투엔드 신경망 아키텍처(Mamba)로 통합했습니다. Mamba는 빠른 추론(Transformer 대비 5배 높은 처리량)과 시퀀스 길이에 대한 선형적 확장성을 지니며, 실제 데이터에서 최대 백만 길이의 시퀀스까지 성능이 향상됩니다. 범용 시퀀스 모델 백본으로서 Mamba는 언어, 오디오, 유전체학과 같은 여러 모달리티 전반에서 최고 수준(state-of-the-art)의 성능을 달성합니다. 언어 모델링에서 Mamba-3B 모델은 동일한 크기의 Transformer를 능가하며, 사전 학습 및 다운스트림 평가 모두에서 두 배 크기의 Transformer와 대등한 성능을 보입니다.

영어 원문 초록

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 and recurrent models, and structured state space models (SSMs) have been developed to address Transformers' computational inefficiency on long sequences, but they have not performed as well as attention on important modalities such as language. We identify that a key weakness of such models is their inability to perform content-based reasoning, and make several improvements. First, simply letting the SSM parameters be functions of the input addresses their weakness with discrete modalities, allowing the model to selectively propagate or forget information along the sequence length dimension depending on the current token. Second, even though this change prevents the use of efficient convolutions, we design a hardware-aware parallel algorithm in recurrent mode. We integrate these selective SSMs into a simplified end-to-end neural network architecture without attention or even MLP blocks (Mamba). Mamba enjoys fast inference (5 $\times$ higher throughput than Transformers) and linear scaling in sequence length, and its performance improves on real data up to million-length sequences. As a general sequence model backbone, Mamba achieves state-of-the-art performance across several modalities such as language, audio, and genomics. On language modeling, our Mamba-3B model outperforms Transformers of the same size and matches Transformers twice its size, both in pretraining and downstream evaluation.

논문 해설

문제점 및 연구 배경

파운데이션 모델은 압도적으로 Transformer 아키텍처와 self-attention에 의존하고 있습니다. 조밀한(dense) 어텐션은 시퀀스 컨텍스트 윈도우 전반에 걸쳐 정보를 효과적으로 라우팅하지만, 두 가지 근본적인 한계를 수반합니다. 바로 시퀀스 길이 $L$ 에 대한 이차 연산 비용 $O (L^{2})$ 과 과거 키와 값(KV 캐시)을 저장해야 한다는 점이며, 이는 스텝당 선형 시간 추론과 메모리 병목 현상을 유발합니다.

선형 어텐션, 순환 신경망 아키텍처, S4와 같은 구조화된 상태 공간 모델(SSMs)과 같은 기존의 준이차 접근법들은 시간 측면에서 선형 또는 준선형으로 확장됩니다. 그러나 구조화된 SSM은 텍스트와 같이 이산적이고 정보 밀도가 높은 모달리티에서 역사적으로 Transformer의 품질을 따라잡지 못했습니다. 저자들은 그 근본 원인을 찾아냈습니다. 바로 기존 SSM이 선형 시불변(LTI) 시스템이라는 점입니다. 전이 역학이 입력 내용에 따라 변하지 않기 때문에, 문맥을 선택적으로 압축하거나, 불필요한 노이즈를 필터링하거나, 가변적인 간격에 걸쳐 기억을 조절할 수 없습니다.

핵심 아이디어

컨텍스트 압축으로서의 선택 메커니즘: 시퀀스 모델링은 과거 이력을 하나의 상태로 압축하는 과정입니다. 기존 LTI SSM은 고정된 전이 파라미터를 사용하므로 모든 타임스텝을 획일적으로 처리해야만 했습니다. 현재 입력 토큰을 기반으로 SSM 행렬을 매개변수화함으로써, 모델은 정보를 선택적으로 필터링하거나 기억하는 데 필요한 내용 인지 능력을 얻게 됩니다.
비합성곱(Non-Convolving) 병목 현상 극복: 파라미터를 입력 의존적으로 만들면 선형 순환(linear recurrence)과 전역 합성곱(global convolution) 사이의 수학적 등가성이 깨집니다. 합성곱을 사용할 수 없다면, 단순한 순환 구현은 $(B, L, D, N)$ 크기의 확장된 잠재 상태를 GPU 고대역폭 메모리(HBM)에 구체화(materialize)해야 합니다. Mamba는 이산화(discretization), 병렬 결합 스캔(parallel associative scan), 출력 투영(output projection)을 빠른 GPU 온칩 SRAM 내부로 융합하여 HBM 트래픽을 최소화함으로써 이 문제를 해결합니다.
통합 아키텍처: Mamba는 번갈아 나타나는 Multi-Head Attention (MHA) 및 MLP 블록을 H3의 게이트 설계와 SwiGLU 스타일의 게이티드 MLP를 결합한 단일 균일 블록으로 대체합니다.

방법론

연속 상태 공간 모델은 암시적 잠재 상태 $h (t) \in R^{N}$ 를 통해 입력 $x (t) \in R$ 를 $y (t) \in R$ 로 매핑합니다. $h^{'} (t) = A h (t) + B x (t), y (t) = C h (t)$

스텝 크기 $Δ$ 와 함께 영차 유지(ZOH) 이산화를 사용하면 연속 시스템이 이산 순환 파라미터로 변환됩니다. $\overline{A} = exp (Δ A), \overline{B} = (Δ A)^{- 1} (exp (Δ A) - I) \cdot Δ B$ $h_{t} = \overline{A} h_{t - 1} + \overline{B} x_{t}, y_{t} = C h_{t}$

선택 메커니즘 (S6)

시간에 따라 $(Δ, B, C)$ 를 상수로 유지하는 대신, Mamba는 이를 입력 의존적으로 만듭니다.

$B_{t} = s_{B} (x_{t}) = Linear_{N} (x_{t})$
$C_{t} = s_{C} (x_{t}) = Linear_{N} (x_{t})$
$Δ_{t} = τ_{Δ} (Parameter + s_{Δ} (x_{t}))$ , 이때 $τ_{Δ} = softplus$ 및 $s_{Δ} (x) = Broadcast_{D} (Linear_{1} (x))$

이는 고전적인 RNN 게이팅을 일반화합니다. $N = 1, A = - 1, B = 1$ 일 때, 순환 식은 게이티드 업데이트 $h_{t} = (1 - g_{t}) h_{t - 1} + g_{t} x_{t}$ 로 직접 단순화되며, 여기서 $g_{t} = σ (Linear (x_{t}))$ 입니다.

하드웨어 인지 상태 확장

$(B, L, D, N)$ 상태를 HBM에 구체화하지 않고 높은 처리량을 유지하기 위해 다음과 같은 과정을 거칩니다.

파라미터 $(Δ, A, B, C)$ 를 HBM에서 온칩 SRAM으로 직접 로드합니다.
$(\overline{A}, \overline{B})$ 로의 이산화와 작업 효율적인 병렬 결합 스캔을 SRAM 내부에서 전적으로 연산합니다.
HBM에 다시 쓰기 전에 잠재 상태를 $C$ 와 축약(contract)하여 $(B, L, D)$ 형태로 만듭니다.
역전파 중에는 중간 상태를 HBM에 저장하는 대신 SRAM에서 재계산합니다.

결과

합성 외삽(Synthetic Extrapolation): Induction Heads 태스크에서 시퀀스 길이 256으로 학습된 Mamba는 최대 $2^{20} = 1, 048, 576$ 길이( $4000 \times$ 시퀀스 길이)까지 완벽하게 외삽(정확도 100%)을 수행한 반면, 어텐션 베이스라인은 실패하거나 $2^{14}$ 를 넘기며 OOM이 발생했습니다. Selective Copying에서 Mamba는 H3의 57.0%, S4의 18.3% 대비 99.8%의 정확도를 달성했습니다.
언어 모델링: Pile 데이터셋에서 Chinchilla 스케일링에 따라 최대 1.3B 파라미터까지 평가했을 때, Mamba는 Transformer++(rotary embeddings, SwiGLU, RMSNorm)와 동등하거나 이를 능가했습니다. 다운스트림 제로샷 태스크(예: LAMBADA, PIQA, HellaSwag, WinoGrande, ARC)에서 Mamba-2.8B는 평균 63.3%의 정확도를 달성하여 Pythia-2.8B(59.1%), Pythia-6.9B(61.7%)를 능가하고 6B~7B 파라미터 모델과 일치하거나 그 이상의 성능을 보였습니다.
유전체학 및 오디오: 시퀀스 길이 1M( $2^{20}$ )의 유인원 DNA 분류에서 Mamba는 HyenaDNA의 54.87% 대비 각각 71.67%(1.4M 파라미터) 및 81.31%(7M 파라미터)의 정확도를 기록했습니다. SC09 음성 생성에서 Mamba는 FID를 1.99(SaShiMi)에서 0.94(6.1M 파라미터) 및 0.67(24.3M 파라미터)로 낮췄습니다.
효율성: 융합 스캔 커널은 표준 PyTorch 스캔 구현보다 20~40배 빠르고, 32K 컨텍스트에서 FlashAttention-2보다 최대 7배 빠르게 실행됩니다. 추론 처리량은 유사한 Transformers 대비 최대 5배 더 높습니다.

한계점 및 미해결 과제

연속 신호 성능: 입력 의존적 선택 메커니즘은 이산적 추론(텍스트, DNA)을 향상시키지만, 원시 오디오 파형 사전 학습에 대한 소거 연구에 따르면 순수 선형 시불변(LTI) 모델(예: 표준 S4)이 균일하게 샘플링된 연속 데이터에서 더 나은 성능을 보입니다.
모델 규모: 평가된 모델은 약 3B 파라미터와 300B 사전 학습 토큰으로 제한되어 있으며, 7B~70B+ 이상의 규모에서의 실증적 성능은 추가적인 연구가 필요합니다.
다운스트림 생태계: Transformers는 프롬프트 엔지니어링, 인컨텍스트 조작, KV 캐시 양자화, 어댑터 튜닝을 위한 광범위한 도구를 보유하고 있습니다. 이러한 상호 작용 방식에 SSM이 어떻게 적응하는지 이해하는 것은 여전히 활발한 탐색 과제입니다.

핵심 용어

Structured State Space Models (SSMs): 암시적 잠재 상태를 통해 1차원 입력을 매핑하는 연속 선형 동적 시스템에서 파생된 심층 시퀀스 모델의 한 종류입니다.
Linear Time Invariance (LTI): 시스템 전이 파라미터가 모든 타임스텝에 걸쳐 정적으로 유지되는 수학적 특성으로, 빠른 전역 합성곱 표현을 가능하게 하지만 내용 기반 추론을 제한합니다.
Selection Mechanism: SSM 전이 역학(Δ, B, C)을 입력 토큰의 함수로 매개변수화하여 모델이 시퀀스 길이를 따라 문맥을 동적으로 전파, 리셋 또는 무시할 수 있도록 하는 메커니즘입니다.
S6: 합성곱 대신 결합 병렬 스캔을 통해 계산되는 선택적 구조화 상태 공간 모델의 약칭입니다.
Parallel Associative Scan: 가속기 하드웨어에서 시퀀스 길이에 걸친 선형 순환을 병렬로 평가하는 작업 효율적인 접두사 연산(prefix computation) 알고리즘입니다.
Hardware-Aware State Expansion: 대규모 중간 상태가 고대역폭 메모리(HBM)를 오가는 것을 방지하기 위해 빠른 온칩 SRAM 내부에서 이산화와 스캔을 수행하는 GPU 커널 융합 및 재계산 전략입니다.
Induction Heads: 모델이 이전 문맥 토큰에 뒤이어 이전에 관찰된 토큰을 재현할 수 있는지를 평가하는 합성 연상 기억 메커니즘입니다.
Mamba Block: 별도의 어텐션이나 MLP 레이어 없이 입력 투영, 1D 합성곱, 선택적 SSM 및 곱셈 게이팅을 결합한 단순화된 단일 시퀀스 모델 블록입니다.

자주 묻는 질문

선택적 SSM은 왜 빠른 FFT 합성곱을 사용하여 학습할 수 없나요?

빠른 전역 FFT 합성곱은 시스템 역학(행렬 A, B, C)이 시간에 따라 일정해야 합니다(선형 시불변성). Mamba는 각 타임스텝에서 B, C, Δ를 입력 토큰의 함수로 만들기 때문에 모든 위치에서 합성곱 커널이 달라지며, 이로 인해 합성곱 등가성이 깨져 순환 병렬 스캔을 사용해야 합니다.

Mamba는 학습 중 순환 상태 확장에 따른 메모리 폭증을 어떻게 방지하나요?

상태 확장 차원 N을 갖는 단순 순환은 GPU HBM에 (B, L, D, N) 형태의 중간 상태를 구체화해야 합니다. Mamba는 커널 융합을 사용하여 저차원 입력만 SRAM으로 로드하고, 빠른 SRAM 내부에서 이산화와 병렬 스캔을 계산하며, 축약된 (B, L, D) 출력만 HBM에 다시 기록하고 역전파 중에는 중간 활성화를 재계산합니다.

Mamba의 Δ 파라미터와 순환 게이팅 사이에는 어떤 관계가 있나요?

스텝 크기 파라미터 Δ는 일반화된 게이팅 메커니즘 역할을 합니다. N=1이고 A=-1인 영차 유지(ZOH) 이산화 환경에서 선택적 순환은 표준 RNN 게이트로 수학적으로 단순화됩니다. 큰 Δ는 모델이 현재 입력에 집중하고 상태를 리셋하도록 만들며, 작은 Δ는 이전 상태를 유지하고 입력을 무시하도록 합니다.

Mamba는 텍스트 생성 중에 KV 캐시를 사용하나요?

아니요. Mamba는 생성 중에 자기회귀 상태 공간 순환으로 동작하므로, 모든 과거 토큰에 대한 키와 값을 캐싱하는 대신 고정된 크기의 은닉 상태(레이어당 D × N 차원)만 유지합니다. 따라서 스텝당 O(1)의 일정한 시간 복잡도로 추론을 수행합니다.

이 페이지의 해설은 AI가 논문을 바탕으로 작성했으며 오류가 있을 수 있습니다. 자세한 내용은 원문을 확인하세요.

XYZNOT로 읽기 (전문 번역과 AI 해설)