Aayush Mittal, Mamba: Redefining Sequence Modeling and Outforming Transformers Architecture

https://www.unite.ai/mamba-redefining-sequence-modeling-and-outforming-transformers-architecture/

Mamba: Redefining Sequence Modeling and Outforming Transformers Architecture

 

Mamba가 Sequence modeling 분야에서  state-space model (SSM)를 통해서 어떻게 진화를 이루어 냈는지 살펴보자.

자연어 처리, 유전학, 음향 분석 같은 다양한 분야에서selective state spaces를 통해서 혁신적인 linear-time sequence 모델링 성능 향상을 이룸

 

기존의 Transformer가 가지고 있던 긴 배열에서의 연산량의 문제를 Mambda가 어떻게 극복했는지 볼 것이다.

 

Mambda의 고유성 : 빠른 연산 능력, selective SSM layer, FlashAttention에 영향을 받은 HW 친화적인 디자인

 

Transformers vs Mamba 

 

Transformer는 attention을 통해서 전체 sequence를 동시에 다룬다. 

... transformer 동작 설명....

 

Transformer와는 반대로 Mamba는 다른 접근법을 가진다.

Mamba는 selective state spaces를 통해서 squence를 다룬다. 

 

 

Mamba

Mamba의 고유성은 전통적인 attention과 MLP block에서 시작한다.  이러한 간략화를 통해서

가볍고 빠르게 sequence의 scale을 linear하게 조절 할 수 있는 것이다. 

 

주요 특징으로는 

1. Selective SSMs: 필요없는 정보를 거르고 필요한 정보에만 집중한다. content-based reasoning에 중요한 부분이다.
2. Hardware-aware Algorithm: 병렬알고리즘을 통해서 HW에 최적화할 수 있다. 
3. Simplified Architecture: attention과 MLP block을 selective SSM으로 통합함을 통해서 더 단순하고 동질적인 구조를 가진다.   

Mambda 코드는 github에 공개

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

GitHub - state-spaces/mamba

Standard Copying tasks are simple for linear models. Selective Copying and Induction Heads require dynamic, content-aware memory for LLMs.

 

Structured State Space (S4)는 최근에 등장 RNN과 CNN 그리고 고전적인 state space 모델의 특정을 담은 촉망된 sequence modeling 방법이다.  sequence를 latent state를 통해 maping하는 방식에서 영향을 받았다. DL 관점에서는 sequence를 효율적이고 적응력 좋게 모델링하는 새로운 방법론이다

 

The Dynamics of S4 Models

SSM(S4) 은 기본적인 structured state space model이다. 

sequence x를 받아서 출력 y를 learned parameters A,B,C와 delay parameter delta를 통해 만든다. 

이 변환은 parameter를 discrtizing 하기, SSM 연산을 적용하기(time-invariant한)를 포함한다. 

The Significance of Discretization

연속 parameter를 고정된 함수를 통해서 이산화한다. 이를 통해 S4 모델이 continuous time system과 연결을 유지할 수 있게 한다. 모델은 이 방식을 통해서 resolution invariance나 적절한 noramlization, 안정성 및 성능향상을 얻는다. 

또한 병렬 gating mechanism을 도출한다. 

 

Linear Time Invariance (LTI)

S4 모델의 주요 특징은 line time invarinace이다. 모델의 역학이 시간이 지나도 유지한다는 것이다. 

LTI는 RNN과 CNN의 초석이고 강력하고도 간단한 sequence 모델을 가능하게 한다. 

 

Overcoming Fundamental Limitations

동적인 역학을 요구하는 분야에서  S4 framework가 LTI 특성으로 인해 제한되었다. 최근의 연구는 tme-varying parmeter를 통해서 LIT제약을 없앰으로 극복하려고 했다. 이를 통해 S4 모델이 더 다양한 sequence와 task를 다루게한다.

 

state space model이란 용어는 latent state를 다루는 재귀적 처리를 하는 모든 방식을 포괄하는 용어이다. 

DL의 맥락에서는 S4 model 또는 구조화된 SSM은 복잡한 sequence를 처리하기 위한 효율적 연산을 수행하는 모델을 지칭한다. 

 

S4 model은 E2E 신경망에 하나의 독립적인 sequence 변환으로 통합할 수 있다. 기존의 CNN에 sequene 모델링을 하는 backbone을 제공하는 것으로 간주할 수 있다. 

 

 

 

Motivation for Selectivity in Sequence Modeling

 

해당 논문은 sequence modeling의 본질이 context를 처리가능한 state로 압축하는 것에 있다고 보았다.  선택적으로 집중하거나 입력을 필터링하는 모델이 더 효율적으로 압축된 state를 유지하여, 더 효율적인 sequence modeling을 가능하게 한다. 이 선택성이 모델이 적응형으로 정보를 통제하는 데 필수적이다.

 

...

 

Performance Highlights of Mamba

 

...

 

 

Getting Started with Mamba

...

Mamba's Impact

....