[paper] mamba

오역, 의역, 생략 및 제 맘대로 해석한 글임에 주의.

---

논문 정보 :

https://arxiv.org/pdf/2312.00752

 

---

Mamba: Linear Time Squence Modeling with Selective State Spaces

 

Abstract

 

 

1. Intoduction

현대 Machien Learning에서는, 대규모 데이터(광범위한 도메인)로 pretraining된 후, downstream task에(개별 도메인에 맞게) 적용하는 Foundation Model이 효과적인 패러다임으로 자리잡았다. 이러한 model의 핵심은 언어, 이미지 등의 sequential data 입력을 처리하는 Sequence Model이다.

현재 대부분의 Foundation Model은 Transformer 아키텍쳐와 그 핵심 구성요소인 Attention Layer를 기반으로 한다.

- 장점 : Self-Attention은 context window(model이 한 번에 볼 수 있는 텍스트의 길이)내에서 밀도 있게 정보를 routing할 수 있어, 복잡한 data를 모델링하느데 효과적이다.

- 단점 : (1) context window 밖은 볼 수 없다. (2) context window의 크기가 늘어남에 따라 계산량이 O(n^2)으로 늘어난다.

이를 해결하기 위한 연구들은 대부분 Attention의 효과를 희생시키는 방향이었다.

 

최근 Sequence Modeling의 대안으로, SSM (Structured State Space Sequence Model)이 조명되었다. SSM은 RNN과 CNN의 특징을 결합한 것으로, 1960년대 State Space Model로부터 영감을 받았다. 이는 매우 효율적으로, sequence를 O(n)으로 처리할 수 있다.

- 장점 : Audio나 Vision과 같은 sequential data에서는 우수한 성능을 보였다

- 단점 : 하지만 text와 같이 discrete하고 정보 밀도가 높은 data를 modeling하는데는 썩 좋지 못했다.

 

 

 

이 논문에서는, Transformer의 modeling 성능을 가지면서도 Sequence 길이에 따라 linear하게 늘어나는, Selected State Space Model을 제안한다.

 

Selective Mechanism

먼저 우리는 기존 모델들의 핵심적인 한계인, 입력에 따라 효율적으로 데이터를 선별(특정 입력에 집중하거나 무시)하는 능력의 부족함을 강조한다. selective copy와 induction head와 같은 중요한 직관들을 바탕으로, 우리는 입력에 따라 SSM parameter들을 parameter화함으로써, 간단한 selection mechanism을 설계한다. 이는 model이 불필요한 정보는 걸러내고, 중요한 정보는 무한히 기억할 수 있게 해 준다.

 

Hardware-aware Algorithm

이러한 간단한 방법은 문제가 하나 생긴다. 기존의 model들은 계산을 빠르게 하기 위해, 입력이 바뀌어도, 계산이 똑같아야 했다. 하지만 입력마다 계산이 달라지니 Convolution을 쓸 수 없게 되었다. 대신에 우리는 scan을 사용해 model을 순환적으로 계산하는 hardware 인식 병렬 알고리즘을 도입했다. 또한 GPU 메모리 계층의 서로 다른 level 사이에서 IO 접근을 피하기 위해, 바로바로 계산하면서 넘겨 버린다. 결과적으로 다른 SSM model에 비해 최대 3배 더 빨라졌다. 

 

Architecture

기존의 Deep Sequence Model을 단순화하여, 기존의 SSM 구조와 transformer의 MLP block을 하나의 block으로 결합함으로써(Attention은 안 씀), 간단하고 일관된 구조를 가진 selective state model인 Mamba를 만들었다.

 

 

 

Selective SSM, 그리고 이를 확장한 Mamba 아키텍처는, 일반적인 foundation model의 backbone으로 사용하기에 매우 적합한, fully recurrent model이다. 그 이유는 아래와 같다.

- High Quality : selectivity 덕에 언어나 유전체 분야에서 높은 성능을 보인다.

- Fast training and inference : 학습 시에는 계산량과 메모리가 O(n)으로 늘어나고, 추론 시에는 이전 요소에 대한 별도의 cache가 필요 없기 때문에, 매우 빠르다.

- Long Context : quality가 좋고 fast하기 때문에, 최대 100만 sequence 길이까지 처리할 수 있게 된다. 

 

 

 

2. State Space Models

생략

 

 

3. Selective State Space Models

3.1 : synthetic task에서 얻은 직관을 우리의 selection mechanism에 적용해 볼 것이다.

3.2 : 그리고 이 selection mechanism을 state space model에 통합해볼 것이다.

3.3 : 그 결과로 time-varying해지기 때문에, SSM에서 convolution을 사용할 수 없어지고, 이를 해결하기 위해 우리는 GPU의 메모리 계층 구졸르 이용한 hardware-aware algorithm으로 이를 극복한다.

3.4. 그 다음 attention이나 MLP조차 없는 단순한 SSM architecture를 설명한다.

3.5. selection mechanism의 몇 가지 추가적인 특지에 대해 논의해볼 것이다.

 

3.1.  

 

4. Empirical Evaluation

 

5. Discussion

6. Conclusion

---

Appendix