Mamba: Linear-Time Sequence Modeling with Selective State Spaces
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
요약
최근에 나온 딥러닝 아키텍처의 경우 대부분 Transformer 아키텍처와 그 핵심인 Attention 모듈을 기반으로 한다. 그러나 Transformer는 긴 시퀀스에 대해 계산 효율이 떨어진다는 단점이 있다. 이를 해결하기 위해 State Space Model이 제안되었다.
기존 모델의 문제점
- 긴 시퀀스에 대한 비효율적인 계산 문제를 해결하기 위해 다양한 아키텍처가 제안되었지만 Transformer를 능가하지 못함.
- Content based reasoning을 수행하는데 약점이 있음
제안 방법론
0. State Space Models
State Space Model(SSM)은 연속적인 시간 시퀀스를 모델링하기 위한 프레임워크로, 시퀀스 데이터의 효율적인 처리를 위해 사용된다. 이 모델은 RNN, CNN의 특성을 결합하여 긴 시퀀스 데이터에서도 효율적으로 동작할 수 있도록 한다.
개념 및 구성
- State 벡터 $h(t)$
- 시퀀스의 현재 정보를 요약
- Observation 벡터 $y(t)$
- 모델이 예측하는 출력값
- Input 벡터 $x(t)$
- 시퀀스의 입력 데이터
SSM의 수학적 정의
- State equation
- $h(t+1) = Ah(t) + Bx(t)
- $A$는 state 전이 행렬, $B$는 input 행렬
- Observation equation
- $y(t) = Ch(t)$
- $C$는 observation 행렬
1. Selection Mechanism
- 입력에 따라 SSM의 파라미터가 동적으로 변하도록 하여, 시퀀스의 각 토큰에 따라 정보를 선택적으로 전달하거나 잊을 수 있도록 하였다.
- 이산 모탈리티에서의 성능 개선
2. Hardware-aware Algorithm
- 효율적인 병렬 알고리즘을 설계하여, 순환모드에서의 계산 효율성을 극대화
- GPU 메모리 계층 구조를 활용해 메모리 접근을 최소화
3. Simple Architecture
- Selective SSM을 Attention이나 MLP 블록 없이도 동작할 수 있는 단순 end-to-end 신경망 아키텍처(Mamba)에 통합
- Mamba는 긴 시퀀스에서도 선형적으로 확장되며, Transformer 대비 5배 빠른 추론 속도를 보임
성능
- SSM + Selective mechanism을 도입했을 때 성능이 더 좋다는 것을 보여줌
- Inductive Heads의 경우에도 Multihead attention을 사용하는 모델보다 Mamba가 훨씬 나은 것을 보여줌
- 더 나은 scaling 성능을 보였으며 특히, 시퀀스의 길이가 길어질수록 Mamba의 성능 향상이 두드러짐
- Mamba 모델의 크기에 따른 비교
- 다양한 크게에서 주요 언어모델들을 능가하는 성능을 보여줌
- 적은 수의 파라미터로도 더 큰 모델과 유사하거나 더 나은 성능을 제공하며, zero-shot 설정에서 좋은 성능을 보임
결론
긴 시퀀스 데이터를 효율적으로 처리하기 위한 새로운 selecive state spaces 모델을 제안하며, Transformer의 한계를 극복하고 더 나은 성능을 제공하는 방법을 제시. Selective Mechanism, Hardware aware algorithm을 통해 다양한 모달리티에서 좋은 성능을 달성할 수 있음을 보여줌.
