Search
Duplicate

AI/ Paper/ Mamba/ Appendix

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...
https://arxiv.org/abs/2312.00752

A Discussion: Selection Mechanism

우리의 selection 메커니즘은 gating, hypernetwork, data-dependence와 같은 연관된 개념에 의해 영감 받았다. 이것은 또한 linear attention(Schlag, Irie, and Schmidhuber 2021)의 매커니즘을 사용하는 classical RNN과 연결되는 ‘fast weight’(J. Ba et al. 2016)와 연관된 것으로 볼 수도 있다. 그러나 우리는 이것이 구분할 가치가 있는 별개의 개념이라고 믿는다.
Gating.
gating은 원래 LSTM이나 GRU 같은 RNN의 gating 메커니즘 또는 Theorem 1의 gated equation (5)을 의미했다. 이것은 입력을 RNN의 hidden state로 허용할지 여부를 제어하는 특별한 메커니즘으로 해석되었다. 특히 이것은 시간을 통한 신호의 전파에 영향을 미치고 입력이 시퀀스 길이 차원을 따라 상호작용하는 원인이다.
그러나 gating의 개념은 이후 일반적인 사용에서 단순히 multiplicative interaction(종종 활성화 함수와 함께 사용)을 의미하는 것으로 완화되었다. 예컨대 신경망 아키텍쳐의 elementwise mulitplicative 컴포넌트(시퀀스 길이를 따라 상호작용하지 않음)는 이제 원래 RNN의 의미보다 매우 다른 의미로 일반적으로 gated architectures(Hua et al. 2022; Mehta et al. 2023)라고 불린다. 따라서 우리는 원래 RNN gating 개념과 일반적인 multiplicative gating의 사용이 실제로 매우 다른 의미를 갖는다고 생각한다.
Hypernetworks.
Hypernetworks는 파라미터가 더 작은 신경망에 의해 생성되는 신경망을 의미한다. original 아이디어(Ha, Dai, and Quoc V. Le 2017)은 recurrent 파라미터가 더 작은 RNN에 의해 생성되는 대형 RNN을 정의하는 좁은 의미로 사용되었다.
Data-Dependence.
hypernetwork와 유사하게 data-dependence은 모델의 일부 파라미터는 데이터에 의존하는(Poli et al. 2023) 개념을 의미할 수 있다.
Example: GLU Acitvation.
이러한 컨셉의 문제점을 설명하기 위해 간단한 대각 선형 레이어 y=Dx\bold{y} = \bold{Dx}를 고려한다. 여기서 D\bold{D}는 대각 가중치 파라미터이다. 이제 D\bold{D}가 x\bold{x}의 선형 변환에 의해 생성되고 선택적으로 비선형이라고 가정하자. D=σ(Wx)\bold{D} = \sigma(\bold{Wx}). 대각 행렬이므로 multiplication은 element-wise product가 된다. y=σ(Wx)∘x\bold{y} = \sigma(\bold{Wx})\circ\bold{x}
이것은 단순한 변환이지만, 기술적으로 gating(multiplicatvie ’분기’를 가지므로), hypernetwork(파라미터 D\bold{D}가 다른 레이어에 의해 생성되므로), data-dependent(D\bold{D}가 데이터 x\bold{x}에 의존하므로)의 일반적인 의미를 만족한다. 그러나 사실 이것은 GLU 함수를 정의할 뿐이며 이것은 너무 단순해서 의미 있는 레이어 대신 단순한 활성화 함수(Dauphin et al. 2017; Shazeer 2020)로 간주되는 경우가 많다.
Selection. 따라서 selection 메커니즘은 아키텍쳐 gating, hypernetwork, data-dependence와 같은 아이디어의 특별한 경우로 간주될 수 있지만, 다른 광범위한 구성 —표준 attention 메커니즘(Bahdanau, Cho, and Bengio 2015; Vaswani et al. 2017)을 포함하여 근본적으로 multiplication을 사용하는 모든 것— 도 마찬가지이며, 그렇게 생각하는 것이 유용하지 않다고 생각한다.
대신 우리는 이것을 전통적인 RNN의 gating 메커니즘과 가장 밀접하게 연관된 것으로 보고 있으며, 이는 특별한 경우(Theorem 1)이며 Δ\Delta의 가변 (입력-의존적) 이산화를 통해 SSM과 더 깊은 연결 역사를 가지고 있다. 또한 우리는 이전 용어의 과도한 사용을 명확히 하기 위해 ‘gating’이라는 용어대신 selection을 사용한다. 더 좁게는 모델이 입력을 선택하거나 무시하고 시퀀스 길이에 따라 데이터 상호작용을 촉진하는 메커니즘적 작용을 selection이라고 한다(섹션 3.1). selective SSM과 gated RNN 외에도 입력 의존적 convolution과 attention도 예가 될 수 있다.

C Mechanics of Selective SSMs

Theorem 1의 증명.
N=1,A=−1,B=1,sΔ=Linear(x),τΔ=softplusN=1, A=-1, B=1, s_\Delta = \text{Linear}(x), \tau_\Delta = \text{softplus}를 갖는 selective SSM(알고리즘 2)를 고려하자. 해당하는 연속시간 SSM(1)은 다음과 같다.
h(t)=−h(t)+x(t)h(t) = -h(t) + x(t)
이것은 또한 leaky integrator라고 불린다.
이산화 단계는 다음과 같다.
Δt=τΔ(Parameter+sΔ(xt))=softplus(Parameter+Linear(xt))=softplus(Lienar(xt))\begin{aligned} \Delta_t &= \tau_\Delta(\text{Parameter} + s_\Delta(x_t)) \\ &= \text{softplus}(\text{Parameter} + \text{Linear}(x_t)) \\ &= \text{softplus}(\text{Lienar}(x_t)) \end{aligned}
여기서 Parameter가 학습가능한 bias이고 linear projection으로 접힐 수 있다는 것을 관찰한다.
이제 zero-order hold(ZOH) 이산화 공식을 적용한다.
A‾t=exp⁡(ΔA)=11+exp⁡(Linear(xt))=σ(−Linear(xt))=1−σ(Linear(xt))B‾t=(ΔA)−1(exp⁡(ΔA)−I)⋅ΔB=−(exp⁡(ΔA)−I)=1−A‾=σ(Linear(xt))\begin{aligned} \overline{A}_t &= \exp(\Delta A) = {1\over 1 + \exp(\text{Linear}(x_t))} = \sigma(-\text{Linear}(x_t)) \\ &= 1-\sigma(\text{Linear}(x_t)) \\ \overline{B}_t &= (\Delta A)^{-1}(\exp(\Delta A)-I)\cdot \Delta B = -(\exp(\Delta A)-I) = 1-\overline{A} \\ &= \sigma(\text{Linear}(x_t)) \end{aligned}
따라서 마지막 discrete recurrence (2a)는
gt=σ(Linear(xt))ht=(1−gt)ht−1+gtxt\begin{aligned} g_t &= \sigma(\text{Linear}(x_t)) \\ h_t &= (1-g_t)h_{t-1} + g_tx_t \end{aligned}
바라던 대로.
위의 A‾t=exp⁡(ΔA)\overline{A}_t = \exp(\Delta A)의 유도는 A=−1A = -1인 경우에 해당한다. Δt=softplus(Linear(xt))\Delta_t = \text{softplus}(\text{Linear}(x_t))이고, softplus(x)=log⁡(1+exp⁡(x))\text{softplus}(x) = \log(1+\exp(x))에 이므로
A‾t=exp⁡(−Δ)=exp⁡(−softplus(Linear(xt)))=exp⁡(−log⁡(1+exp⁡(Linear(xt)))=1exp⁡(log⁡(1+exp⁡(Linear(xt)))=11+exp⁡(Linear(xt))=σ(−Linear(xt))=1−σ(Linear(xt))\begin{aligned} \overline{A}_t = \exp(-\Delta) &= \exp(-\text{softplus}(\text{Linear}(x_t))) \\ &= \exp(-\log(1 + \exp(\text{Linear}(x_t))) \\ &= {1\over \exp(\log(1 + \exp(\text{Linear}(x_t)))} \\ &= {1 \over 1 + \exp(\text{Linear}(x_t))} \\ &= \sigma(-\text{Linear}(x_t)) \\ &= 1 - \sigma(\text{Linear}(x_t)) \end{aligned}
마지막은 σ(z)=11+exp⁡(−z)\sigma(z) = {1\over 1 + \exp(-z)}이므로 σ(−z)=11+exp⁡(z)=1−σ(z)\sigma(-z) = {1\over 1+ \exp(z)} = 1 - \sigma(z)라는 것을 이용한다.
또한 B‾t=(ΔA)−1(exp⁡(ΔA)−I)⋅ΔB\overline{B}_t = (\Delta A)^{-1}(\exp(\Delta A)-I)\cdot \Delta B의 유도는 A=−1,B=1A = -1, B=1인 경우에 해당한다.
B‾t=(−Δ)−1(exp⁡(ΔA)−I)⋅Δ=−1Δ(A‾−I)⋅Δ=1−A‾=1−(1−σ(Linear(xt)))=σ(Linear(xt))\begin{aligned} \overline{B}_t &= (- \Delta)^{-1}(\exp(\Delta A)-I)\cdot \Delta \\ &= -{1 \over \cancel{\Delta}}(\overline{A}-I)\cdot \cancel{\Delta} \\ &= 1 - \overline{A} \\ &= 1 - (1-\sigma(\text{Linear}(x_t))) = \sigma(\text{Linear}(x_t))\end{aligned}
\여기서 AA를 스칼라로 사용했기 때문에 II를 1로 대체했다.
A=−1,B=1A = -1, B=1에 대해 A‾=1−σ(Linear(xt))\overline{A} = 1 - \sigma(\text{Linear}(x_t))이고 B‾=σ(Linear(xt))\overline{B} = \sigma(\text{Linear}(x_t))가 되므로, gt=σ(Linear(xt))g_t = \sigma(\text{Linear}(x_t))라 정의하여 다음이 성립한다.
ht=A‾ht−1+B‾xt=(1−gt)ht−1+gtxt\begin{aligned} h_t &= \overline{A}h_{t-1} + \overline{B}x_t \\ &= (1-g_t)h_{t-1} + g_t x_t\end{aligned}
Made with