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

Albert Gu, Tri Dao|arXiv (Cornell University)|2023. 12. 01.

Neural Networks and Applications인용 수 967

한 줄 요약

Mamba는 입력 의존적인 선택적 상태 공간 모델을 도입하여 Transformer보다 추론 처리량이 5× 더 높은 선형 시간 시퀀스 모델링을 가능하게 하며, 언어, 오디오, 게놈학 전반에서 강력한 성능을 보인다. 언어 영역에서 Mamba-3B는 동등한 규모의 Transformer보다 우수하고, 규모를 두 배로 키운 경우에도 Transformer와 동등한 성능을 달성한다.

ABSTRACT

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$ imes$ 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의 비효율성을 해결하기 위해 콘텐츠 기반의 선택적 상태 공간(SSM)을 도입한다.
• 어텐션이나 MLP 없이도 정확도를 유지하며 모든 모달리티에서 선형 시간 시퀀스 모델링을 가능하게 한다.
• 순환 방식에서도 효율성을 유지하는 하드웨어 친화적 병렬 알고리즘을 선택적 SSM에 대해 개발한다.

제안 방법

• SSM 매개변수를 입력 주소의 함수로 만들어 시퀀스를 따라 정보 흐름을 이산적이고 콘텐츠 기반으로 조절할 수 있도록 한다.
• 어텐션이나 MLP 없이 간단한 엔드 투 엔드 백본에 선택적 SSM을 통합한다(이를 Mamba 모델이라 한다).
• 표준 합성곱 최적화가 부재한 상황에서도 빠른 추론을 유지하기 위한 하드웨어 친화적 병렬 순환 알고리즘을 설계한다.
• 추론 중 시퀀스 길이에 따른 선형 확장성과 Transformer 대비 5× 처리량 향상을 입증한다.

실험 결과

연구 질문

• RQ1 콘텐츠 기반(주소 의존적) 선택적 SSM이 전통적인 하위 제곱 시간 모델과 비교하여 긴 시퀀스 추론을 개선할 수 있는가?
• RQ2 선택적 SSM이 언어, 오디오, 게놈학 전반에 걸쳐 경쟁력 있는 성능으로 선형 시간 시퀀스 모델링을 가능하게 하는가?
• RQ3 주의/MLP 없이 순환형 구조에서 하드웨어 친화적 병렬성을 유지하면서 처리량을 보장하는 것이 가능한가?
• RQ4 동일 규모의 Transformer 및 더 큰 규모의 Transformer에 비해 언어 모델링에서 Mamba의 성능은 어떤가?

주요 결과

• Mamba는 Transformer보다 5× 높은 처리량으로 빠른 추론을 달성한다.
• 모델은 시퀀스 길이에 따라 선형적으로 확장된다.
• Mamba는 언어, 오디오, 게놈학 모달리티 전반에서 최첨단 성능을 보여준다.
• 언어 모델링에서 Mamba-3B는 동일 규모의 Transformer보다 우수하고, 2배 규모의 Transformer와도 사전학습 및 다운스트림 평가 모두에서 동등한 성능을 달성한다.