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

Mamba: Linear-Time Sequence Modeling with Selective State Spaces

December 1, 2023

저자: Albert Gu, Tri Dao

cs.AI

초록

딥러닝 분야에서 가장 흥미로운 애플리케이션들을 주도하고 있는 파운데이션 모델들은 거의 전적으로 트랜스포머(Transformer) 아키텍처와 그 핵심 구성 요소인 어텐션(attention) 모듈을 기반으로 합니다. 긴 시퀀스에서 트랜스포머의 계산 비효율성을 해결하기 위해 선형 어텐션(linear attention), 게이트 컨볼루션(gated convolution), 순환 모델(recurrent models), 그리고 구조화된 상태 공간 모델(Structured State Space Models, SSMs)과 같은 서브쿼드라틱 시간(subquadratic-time) 아키텍처들이 개발되었지만, 이러한 모델들은 언어와 같은 중요한 모달리티에서 어텐션만큼의 성능을 보이지 못했습니다. 우리는 이러한 모델들의 주요 약점이 콘텐츠 기반 추론(content-based reasoning)을 수행할 수 없다는 점임을 확인하고, 이를 개선하기 위해 몇 가지 접근을 시도했습니다. 첫째, SSM 매개변수를 입력의 함수로 설정함으로써 이산 모달리티(discrete modalities)에서의 약점을 해결하여, 모델이 현재 토큰에 따라 시퀀스 길이 차원을 따라 정보를 선택적으로 전파하거나 잊어버릴 수 있도록 했습니다. 둘째, 이러한 변경으로 인해 효율적인 컨볼루션 사용이 불가능해졌지만, 하드웨어를 고려한 병렬 알고리즘을 순환 모드에서 설계했습니다. 우리는 이러한 선택적 SSMs를 어텐션 또는 MLP 블록 없이도 동작하는 단순화된 엔드투엔드 신경망 아키텍처(Mamba)에 통합했습니다. Mamba는 빠른 추론 속도(트랜스포머 대비 5배 높은 처리량)와 시퀀스 길이에 대한 선형 스케일링을 자랑하며, 실제 데이터에서 최대 백만 길이의 시퀀스까지 성능이 향상됩니다. 일반적인 시퀀스 모델 백본으로서, Mamba는 언어, 오디오, 유전체학 등 여러 모달리티에서 최첨단 성능을 달성합니다. 언어 모델링에서, 우리의 Mamba-3B 모델은 동일한 크기의 트랜스포머를 능가하며, 크기가 두 배인 트랜스포머와도 프리트레이닝 및 다운스트림 평가에서 동등한 성능을 보입니다.

English

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 (5times 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.