728x90
Mamba: Linear-Time Sequence Modeling with Selective State Spaces
요약
최근에 나온 딥러닝 아키텍처의 경우 대부분 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을 통해 다양한 모달리티에서 좋은 성능을 달성할 수 있음을 보여줌.
728x90
반응형