https://arxiv.org/abs/2312.00752

 

 

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

by. 카네기멜론대학교

https://github.com/state-spaces/mamba

 

GitHub - state-spaces/mamba

Contribute to state-spaces/mamba development by creating an account on GitHub.

github.com

 

 

 

0. 핵심 & Main Contribution + 후기

SSM은 시계열 데이터를 효율적으로 처리할 수 있는 구조를 갖는다.

이 모델은 시간에 따른 데이터의 변화를 고려하여, 더 정확하고 빠른 예측을 할 수 있다.

또한, 계산 비용이 상대적으로 낮아서, 긴 시퀀스 데이터를 처리하는 데 효과적이다.

 

 

반면, Transformer는 '어텐션 메커니즘(Attention Mechanism)'을 사용하여 데이터의 관계를 학습.

강점: 문장 내 단어들 사이의 관계를 파악하여, 전체 문맥을 이해

약점: 하지만, Transformer는 많은 데이터와 복잡한 계산을 필요로 하며, 긴 문장을 처리하는 데 어려움이 있다. 

 

 

 🐍 Mamba...?


혁신적인 SSM(State Space Model)로 복잡한 Seq처리에 용이
∙ 선택적 접근방식으로 추론속도가 빨라짐.
∙ seq_len에 따른 선형확장이 가능해 처리량이 크게 확장

[Main Contribution]
① 빠른 처리능력
② 선택적 SSM Layer
Flash Attention에서 영감을 얻은 HW친화적설계

Transformer: 복잡한 Attention Mechanism을 활용:
∙ Input Processing
∙ Attention Mechanism & Multi-head Attention
∙ Encoder-Decoder 구조
∙ Point-wise FFNs


Mamba: 선택적 상태공간을 사용,  더 많은 computing을 제공.
seq_len에 따라 선형적으로 확장해 특히나 seq가 길어짐에 따라 더 효율적인 seq_modeling을 위한 새로운 패러다임을 설정하는데, 다음과 같다:
먼저, 가장 눈에 띄는 점은 만연하는 Attention 및 MLP Block에서의 탈피 이다.

[Selective SSMs]
∙ 관련없는 정보를 filtering  /  관련정보에 집중해 seq처리 향상
∙ 이런 선택성은 효율적인 내용기반추론에 매우 중요

[Hardware-Aware Algorithm]
∙ GPU에 최적화된 병렬알고리즘을 사용
∙ 더 빠른 계산 및 메모리요구사항 감소

[Simplified Architecture]
∙ Selective SSMs를 통합
∙ Attention 및 MLP블록을 제거해 확장성이 향상
Standard Copying tasks are simple for linear models. Selective Copying and Induction Heads require dynamic, content-aware memory for LLMs.


cf) 이산화의 중요성
이산화는 고정된 공식을 통해 연속 매개변수를 이산 매개변수로 변환하여 S4 모델이 연속시간 시스템과의 연결을 유지할 수 있도록 하는 핵심 프로세스로
모델에 해상도 불변성과 같은 추가 속성을 부여하고 적절한 정규화를 보장하여 모델 안정성과 성능을 향상시킨다.

또한 이산화는 네트워크를 통한 정보흐름을 관리하는 데 중요한 RNN의 게이팅 메커니즘과 유사하다.

 

 

 

 

 

 

 

 

 

 

 

 

 

1. Intro

Background)

✔️State Space :


제어이론 기반, 동역학에서 사용되는 방정식.

현재 상태를 입력으로 시스템의 미래상태를 예측하는데 사용.
2개의 방정식이 존재:

✔️State equation:  X' = AX + Bu

✔️Output equation:  Y = CX + Du

이때, A를 잘 도출하는 것이 핵심.
보통 D를 0으로 두고 진행하는 경우가 다수.



위의 상태방정식 및 상태공간모델을 기반, 새로운 layer를 구축.
목표: 긴 sequence를 효율적으로 modeling하는 것.
필요 조건: 
  1) A 행렬을 잘 정의해야함 (by SSM)
  2) Continuous 2 Discrete
1. continuous → discrete  2. A행렬 잘 정의



✔️SSM(State Space Models):

이전에는 SSM이 매우 낮은 성능을 기록.
긴 seq에서 gradient problem이 발생했을 것으로 추측.

이를 해결하기 위해, HiPPO논문에서는 고차 다항식을 이용한 projection연산으로 History를 기억하게 만든다.
특정행렬(A)를 잘 정의하는 것이 핵심!!
A는 충분한 정보를 갖는 과거의 History를 잘 압축(기억)하는 행렬(hidden_state)인 것!
Legendre's Equation
위의 르젠드르 방정식을 이용하면, 이전 모든 History를 근사시킬 수 있는데, 이는 위의 빨간 선과 같다.
cf) 파란선: Legendre's 요소에 대한 계수에 대한 값들.
즉, HiPPO행렬 A는 각 단계마다 위 계수를 update하는 역할을 한다.





✔️ ZOH Rule (Zero Order Hold Rule)

좌) Continuous Equation  우) Discrete Equation
Continuous한 상태를 Discrete한 상태로 만들어야 한다.
ZOH규칙으로 A와 B행렬을 훈련시킬 수 있으며, 이를 이용해 Discrete한 상태로 만들 수 있다.
cf) 오일러 방법으로도 이산화 가능.

RNN과 CNN은 input Data차이
token형태의 연산상태 --> RNN형태로 도출
Data전체를 볼 수 있는 연산상태 --> CNN형태로 도출


✔️Limitation:




Sol) Hardware Approach - Kernel Diffusion

B: Batch_size  //. L: sequence_Length. //. D: d_model. //. N: hidden_state_size


[보통의 연산순서]:
1. Tensor연산이 DRAM에서 SRAM으로 Tensor를 로드 << Too Slow
2. GPU(CUDA)에서 연산을 진행
3. DRAM에 결과를 저장

[Kernel Diffusion]:
1. DRAM에서 빠른 SRAM으로 직접 Δ, A, B, C 로드
2. SRAM에서 이산화작업 진행 --> A ̅, B ̅ 생성.
3. 최종 출력인 B,L,D 만 DRAM에 다시 씀.

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

2. Recomputation

🤔  Mamba: Recomputation

위와 같은 DNN 훈련 시, Backpropagation을 살펴보자.
Gradient값 계산을 위해 순전파 단계의 값을 caching하는데, 
순전파 동안의 활성화 값을 caching한 후, BP동안 이를 재사용하는 것을 착안, 이를 GPU에 도입한 것.

BP 도중, DRAM을 재복사해야하는 비효율적인 점이 존재.
순전파단계 출력값을 BP도중 재계산하자!
결과적으로 Fusion Selective Scale Layer에서 Flash Attention을 적용한 Transformer와 동일한 효과를 갖는다!


🤔  Mamba Architectures

Transformer와 비견될정도로 좋다~

 

 

 

 

 

+ Recent posts