Mamba와 Selective State

Mamba는 SSM을 LLM에 맞게 다시 설계한 대표 구조다.

핵심은 selective state다.

모든 token을 같은 방식으로 state에 반영하지 않고,
입력 token에 따라 state update 방식을 바꾼다.

왜 selective가 필요한가

언어 sequence에서는 모든 token이 같은 중요도를 갖지 않는다.

어떤 token은 오래 기억해야 하고, 어떤 token은 금방 잊어도 된다. 어떤 정보는 다음 몇 token에만 필요하고, 어떤 정보는 문단 끝까지 영향을 준다.

고정된 update rule만 쓰면 이런 선택을 표현하기 어렵다.

Mamba는 입력으로부터 state update에 필요한 값들을 만든다.

input_t
  -> update에 쓸 파라미터들
  -> state_{t-1}를 state_t로 갱신

그래서 같은 모델이라도 token마다 state를 갱신하는 방식이 달라진다.

Attention은 현재 token이 과거 token들을 직접 다시 본다.

current query -> past keys/values 전체 참조

Mamba는 과거 전체를 직접 다시 보는 대신 state를 갱신해 간다.

previous state + current input -> next state

이 차이는 모델 구조의 차이이기도 하지만, 추론 시스템의 차이이기도 하다. Attention serving은 KV cache를 관리한다. Mamba serving은 convolution state와 SSM state를 관리한다.