티스토리 뷰
[논문] Mamba: Linear-Time Sequence Modeling with Selective State Spaces
sikaro 2024. 7. 28. 16:24transformer 기반 백본은 NLP 모델에 아주 파격을 불러왔고, 지금도 많이 연구되고 있는 여러모로 핫한 논문이다.
여태까지 모델은 transformer 기반 구조로, self-attention 메커니즘을 이용해서 텍스트의 정밀도와 정확도를 예측해왔다. 하지만 본 논문에서는 다시 RNN과 CNN 같은 초심으로 돌아가, 게이트 메커니즘을 사용하면서도 아주 빠르게 선형적으로 처리할 수 있는 기반을 마련한다.
어떻게? Selective SSM(State Spaces Modeling)을 통해서!
상태 공간 모델링은 원래 제어에 쓰는 것이다. 본 블로그에서도 로봇이론 포스팅을 할 때 한 번 소개한 적이 있었다.
학부가 기계과인 나로서는 자주 보던 공식이라 상당히 흥미로웠다.
본 논문은 어려워보이지만, 핵심적으로 요약하면 공식이 간단하면서도 정확도가 잘 나온다는 것에 있다.
그 쉬운 공식은 아래와 같다.
모르는 분들은 아, 그렇구나, 하고 넘어가겠지만, 대학을 나온 분들은 아실 것이다.
위의 식이 미분방정식이라는 걸.
미분 방정식이라고 하면 어려워 보이지만, 위의 것들은 사실 답이 정해져 있는 미분 방정식이라 정말 어렵지 않다.
각각에 대한 표현은 다음과 같다.
- 는 상태 벡터.
- h′(t)는 상태 벡터의 시간에 대한 도함수(즉, 변화율).
- A는 상태 행렬로, 시스템의 동적 특성.
- B는 입력 행렬로, 외부 입력 x(t)가 상태 변화에 미치는 영향.
- x(t)는 시스템에 입력되는 외부 신호.
- C는 출력 행렬로, 상태 벡터 h(t)가 출력 y(t)에 미치는 영향.
- y(t)는 시스템의 출력.
수식은 어려워 보이지만, 가중치에 대한 개념을 알고 있다면 말로 하면 쉽다.
즉, h(t)라는 상태 벡터의 변화(h'(t))를 알고 싶은데, 그거는 결국 h(t)와 외부의 영향인 x(t)의 합으로 이루어지고, 적절히 조화되서 나온다는 뜻이다.
그리고 C는 그 상태의 h(t)에서 그냥 C를 곱해주는 거고.
즉, 선형 시스템이기 때문에 직관적으로 이해하기 쉽게 AC * h + Bx 정도로 알아들으면 된다.
여기까지 오면 잘 생각해보자. 어? 그러면 결국 n*n 구조인데, transformer랑 다를게 없지 않은가?
아니다. 여기서 중요한 건 이게 이산화가 가능한 구조라는 것에 있다.
내가 로봇 이론 포스팅에서 한번 설명한 적이 있을 것이다. exponential 함수는 극단적으로 빠르게 수렴하는 특징을 가지고 있다고. 그래서 제어 이론에서 많이 쓰이는 거고, 상태 방정식이 미분 방정식으로 표현되는 것도 그렇다.
그렇기에 똑같은 n*n 구조라도 exponential하게 수렴하므로, transformer보다 5배 빠르고, 중요한 정보도 놓치지 않는 그야말로 아주 적합한 백본의 구조가 되어버리는 것이다.
그리고 무엇보다, 선형 구조이므로 병렬 처리가 아주 원활해진다.
self-attention은 병렬 처리 하려면 코드를 아주 기가 막히게 짜야 하지만, 이거는 그냥 데이터를 쪼개기만 하면 된다.
그런데도 성능이 비슷하게 나왔다고 하니, 이게 아주 엄청난 반향이었던 것이다.
이것에서 파생되어서 VIT처럼, Vision mamba도 나오고, MoE 메커니즘에 대한 것도 연속적으로 나왔다.
2024년 상반기에 우후죽순으로 빠른 모델이 나온 것도 아마 이 mamba 모델의 영향이지 않을까 싶다.
논문에서 사용된 mamba 블록은 기본적인 MLP 모델에 H3를 더했다.
conv 레이어에 SSM으로 시퀀스 변환을 수행하고, 해당하는 것을 게이트(B)로 만들어서 비선형으로 출력해 아주 효율적이면서도 계산이 빠른 모델이 만들어졌다.
엄청나게 간단한게 가히 혁명적인 논문이라고 불릴만 하다. transformer 아키텍쳐 이후에 이런 백본 모델이 다시 나올줄은 예상하지 못했다.
정말 모델의 세계는 무궁무진하니, 앞으로도 논문 열심히 읽고 공부해야 따라갈 수 있을 듯 싶다.