[Paper Review] Mamba - Linear Time Sequence Modeling with Selective State Spaces 2

3. Selective State Space Models

3.1 Selection as a Means of Compression

figure 2 왼쪽의 Copying task의 기본 버전은 입력과 출력 사이에 일정 거리를 두고 시불변한 모델로 global convolution이나 선형 회귀를 통해서 쉽게 해결할 수 있음. Selective Copying은 입력과 출력 사이 거리가 무작위로 선정되고, 시간에 따라 변화하며 선택적으로 입력을 기억할지 무시할지를 contents에 따라 결정할 수 있음. Induction head 는 LLM의 핵심 능력인 맥락에 따른 추론에 관련된 것

병합 작업에 관한 두가지 실행 예시

• Selective Copying : 기억할 토큰의 위치를 바꿔 Copying Task를 수정한다. 관련있는 토큰을 기억하고 관련없는 토큰을 걸러내려면 내용 인식 추론(content-aware resoning)이 필요하다.
• Induction Heads : 적절한 컨텍스트에서 출력을 낼 시기를 알기 위해서는 내용 인식 추론이 필요하다. LLM의 동작 과정 설명을 위해 가장 많이 쓰이는 매커니즘.

이 작업은 LTI 모델의 실패한 모드를 보여준다. 회귀적 관점에서 constant dynamics(여기에서는 $\bar{A}, \bar{B}$)는 context에서 올바른 정보를 선택하거나, 입력에 따른 방식으로 시퀀스를 따라 hidden state에 영향을 줄 수 없다.
합성곱적 관점에서 global convolutions는 일반 copying task를 해결할 수 있다고 알려져 있는데, 왜냐하면 그게 시간 인식 기반이기 때문이다. 그러나 내용 인식 부족의 문제가 있어서 Selective Copying task에는 어려움이 있다.

시퀀스 모델의 효율성과 효과성(얼마나 효율적으로 목표를 달성했는가 vs 목표를 얼마나 성공적으로 달성했는가)의 트레이드 오프는 얼마나 자신의 state를 잘 압축했는가로 표현된다. 효율적인 모델은 반드시 작은 state를 가져야하고, 반면에 효과적인 모델은 context에서 모든 필요한 정보를 state에 포함해야한다.

이 논문에서 시퀀스 모델 수립에 필요한 기본 원칙을 집중을 위한 내용인식 혹은 입력을 sequential state로 필터링할지에 대한 선택성 Selectivity로 제안한다. 부분적으로 selection mechanism은 정보 전파 혹은 시퀀스 차원에 대한 상호작용을 제어한다.

3.2 Improving SSMs with Selection

Selection 매커니즘을 모델에 적용하는 방법은 시퀀스를 따라 상호작용에 영향을 미치는 파라미터들을 입력에 따라 움직이게 하는 것이다.

가장 다른 점은 SSM과 달리 $\Delta, B, C$ 파라미터를 input에 대해 함수로 정의하고, tensor의 shape을 바꾼 것이다. 부분적으로는 이 파라미터들이 이제 길이 차원 L을 가진다는 것인데, 이는 모델이 시불변에서 시간 변화를 반영하는 쪽으로 바뀌었다는 뜻이다.

$$\begin{matrix}
s_{B}(x) &=& Linear_{N}(x) \
s_{C}(x) &=& Linear_{N}(x) \
s_{\Delta}(x) &=& Broadcast_{D}(Linear_{1}(x) \
\tau_{\Delta} &=& softplus
\end{matrix}$$

$Linear_{d}$ 는 파라미터화 된 차원 $d$ 의 projection이다. $s_{\Delta}$ 와 $\tau_{\Delta}$ 의 선택은 RNN gating 매커니즘과의 연결 때문이다.

3.3 Efficient Implementation of Selective SSMs

선택 매커니즘은 꽤 자연스럽고, 이전 연구에서는 recurrent SSMs에서 $\Delta$를 시간이 지날 수록 변화시키는 것과 같은 선택의 특별한 케이스를 통합하려고 노력했다. 그러나 SSMs 사용의 주요 한계는 연잔적 효율성에 있다. 이는 S4나 다른 파생 알고리즘이 LTI(non-selective) 모델을, 그 중에서도 global convolution 형태의 모델을 사용하는 이유이다.

3.3.1 Motivation of Prior Models

• SSMs와 같은 recurrent model은 표현력과 속도 사이의 균형을 잡아야한다.
• recurrent 모드에서 파생된 것이 convolution 모드라서 recurrent가 좀 더 유연하다. 그러나 이 과정에서 (B, L, D, C) 차원의 latent state h를 계산하고 구체화해야하는데, 이게 입력 x와 출력 y의 차원보다 훨씬 크다. 그래서 state 연산을 우회하고 convolution kernel(B, L, D)를 구체화하는 convolution mode가 더 효휼적이라고 소개됐다.
• 이전 LTI state space model 은 이중 순환-합성곱 형태를 사용해서 효율성 저하 없이 기존 RNN보다 10-100배에 달하는 state 차원을 늘린다.