NL-345, Mamba: Linear-Time Sequence Modeling with Selective State Spaces, CoLM 2024

Abstract

파운데이션 모델은 현재 딥러닝의 흥미로운 응용 대부분을 구동하고 있으며, 거의 예외 없이 Transformer 아키텍처와 그 핵심인 attention 모듈에 기반한다. 긴 시퀀스에서 Transformer의 계산 비효율성을 해결하기 위해 linear attention, gated convolution 및 recurrent 모델, structured state space model(SSM)과 같은 많은 subquadratic-time 아키텍처가 개발되어 왔지만, 언어와 같은 중요한 모달리티에서는 attention만큼 좋은 성능을 보이지 못했다. 우리는 이러한 모델들의 핵심 약점이 content-based reasoning을 수행하지 못하는 데 있음을 확인하고, 이를 개선하기 위한 몇 가지 방법을 제안한다. 첫째, SSM 파라미터가 입력의 함수가 되도록 하는 단순한 변경만으로도 discrete modality에서의 약점을 해결할 수 있으며, 모델이 현재 토큰에 따라 시퀀스 길이 차원을 따라 정보를 선택적으로 전파하거나 잊을 수 있게 한다. 둘째, 이 변경으로 인해 효율적인 convolution을 사용할 수 없게 되지만, 우리는 recurrent mode에서 hardware-aware parallel algorithm을 설계한다. 우리는 이러한 selective SSM을 attention은 물론 MLP block조차 없는 단순화된 end-to-end neural network architecture인 Mamba에 통합한다. Mamba는 빠른 inference, 즉 Transformer보다 5배 높은 throughput을 보이며, 시퀀스 길이에 대해 선형적으로 scaling된다. 또한 실제 데이터에서 million-length sequence까지 성능이 향상된다. 일반적인 sequence model backbone으로서 Mamba는 언어, 오디오, 유전체학 등 여러 모달리티에서 state-of-the-art 성능을 달성한다. Language modeling에서 Mamba-3B 모델은 동일한 크기의 Transformer를 능가하며, pretraining과 downstream evaluation 모두에서 두 배 크기의 Transformer와 맞먹는 성능을 보인다.

1 Introduction

파운데이션 모델(Foundation Models, FMs)은 대규모 데이터로 사전학습(pretraining)된 뒤 다운스트림 작업에 적응되는 대형 모델로서, 현대 머신러닝의 효과적인 패러다임으로 자리 잡았다. 이러한 FM의 백본(backbone)은 대개 시퀀스 모델(sequence model)이며, 언어, 이미지, 음성, 오디오, 시계열, 유전체(genomics) 등 다양한 도메인의 입력 시퀀스를 처리한다 (Brown et al. 2020; Dosovitskiy et al. 2020; Ismail Fawaz et al. 2019; Oord et al. 2016; Poli et al. 2023; Sutskever, Vinyals, and Quoc V Le 2014). 이 개념 자체는 특정 모델 아키텍처에 종속되지 않지만, 현대의 FM은 거의 대부분 Transformer(Vaswani et al. 2017)와 그 핵심 구성요소인 attention layer(Bahdanau, Cho, and Bengio 2015)에 기반하고 있다.

Self-attention의 효과성은 문맥 윈도우(context window) 내부에서 정보를 조밀하게 전달(route)할 수 있다는 능력에서 비롯되며, 이를 통해 복잡한 데이터를 모델링할 수 있다. 그러나 이러한 특성은 근본적인 단점도 함께 가져온다. 첫째, 유한한 윈도우 밖의 정보를 모델링할 수 없으며, 둘째, 윈도우 길이에 대해 계산량이 이차적으로(quadratically) 증가한다. 이러한 문제를 해결하기 위해 attention의 더 효율적인 변형들이 대거 제안되었으나(Tay, Dehghani, Bahri, et al. 2022), 대부분 attention의 핵심적인 장점을 희생하는 대가를 치렀다. 현재까지 이들 변형 중 어느 것도 다양한 도메인에서 대규모로 attention만큼 효과적이라는 것이 실증적으로 입증되지는 못했다.

최근에는 structured state space sequence models(SSMs) (Gu, Goel, and Ré 2022; Gu, Johnson, Goel, et al. 2021)가 시퀀스 모델링을 위한 유망한 아키텍처 계열로 부상하였다. 이러한 모델은 recurrent neural network(RNN)와 convolutional neural network(CNN)의 결합으로 해석될 수 있으며, 고전적인 state space model(Kalman 1960)에서 영감을 받았다. 이 계열의 모델은 recurrence 또는 convolution 형태로 매우 효율적으로 계산될 수 있으며, 시퀀스 길이에 대해 선형 혹은 준선형(linear or near-linear)으로 확장된다. 또한 특정 데이터 모달리티에서는 장거리 의존성(long-range dependencies)을 모델링할 수 있는 원리적인 메커니즘을 제공하며(Gu, Dao, et al. 2020), Long Range Arena(Tay, Dehghani, Abnar, et al. 2021)와 같은 벤치마크를 지배해 왔다. 다양한 형태의 SSM(Gu, Goel, and Ré 2022; Gu, Gupta, et al. 2022; Gupta, Gu, and Berant 2022; Y. Li et al. 2023; Ma et al. 2023; Orvieto et al. 2023; Smith, Warrington, and Linderman 2023)은 오디오와 비전 같은 연속 신호 데이터 도메인에서 성공적이었다(Goel et al. 2022; Nguyen, Goel, et al. 2022; Saon, Gupta, and Cui 2023). 그러나 텍스트처럼 이산적(discrete)이며 정보 밀도가 높은 데이터에서는 상대적으로 성능이 떨어졌다.

SSM이라는게 원래 존재했군?

본 논문에서는 선택적 상태공간 모델(selective state space models)이라는 새로운 계열을 제안하며, 이를 통해 Transformer 수준의 모델링 성능을 달성하면서도 시퀀스 길이에 대해 선형적으로 확장되는 모델을 구현한다.

선택 메커니즘(Selection Mechanism).
우선, 기존 모델들의 핵심적인 한계를 지적한다. 그것은 입력에 의존하는 방식(input-dependent manner)으로 데이터를 효율적으로 선택(select)할 수 없다는 점이다. 즉, 특정 입력에 집중하거나 무시하는 능력이 부족하다. 우리는 selective copy 및 induction heads와 같은 중요한 합성(synthetic) 태스크에서 얻은 직관을 바탕으로, SSM의 파라미터를 입력의 함수로(parameterizing the SSM parameters based on the input) 만드는 간단한 선택 메커니즘을 설계하였다. 이를 통해 모델은 불필요한 정보를 걸러내고(filter out irrelevant information), 중요한 정보는 무기한 기억할 수 있게 된다.

하드웨어 친화적 알고리즘(Hardware-aware Algorithm).
하지만 이러한 단순한 변경은 계산 측면에서 기술적 문제를 야기한다. 사실 기존 SSM들은 계산 효율성을 위해 반드시 시간 및 입력 불변(time- and input-invariant)이어야 했다. 우리는 convolution 대신 recurrence 기반의 scan 연산을 사용하는 하드웨어 친화적 알고리즘을 통해 이 문제를 해결한다. 또한 GPU 메모리 계층 구조 사이에서 발생하는 IO 접근을 줄이기 위해, 확장된 상태(expanded state)를 실제로 메모리에 물질화(materialize)하지 않는다. 그 결과, 제안한 구현은 이론적으로도(모든 convolution 기반 SSM이 pseudo-linear scaling을 가지는 반면, 우리는 sequence length에 대해 truly linear scaling을 달성) 그리고 실제 현대 하드웨어 상에서도(A100 GPU에서 최대 3배 빠름) 기존 방법보다 더 빠르다.

아키텍처(Architecture).
또한 우리는 기존의 깊은 시퀀스 모델 아키텍처를 단순화하였다. 기존 SSM 아키텍처(Dao, Fu, Saab, et al. 2023)의 설계와 Transformer의 MLP block을 하나의 블록으로 결합함으로써, attention은 물론 MLP block조차 포함하지 않는 단순하고 균일한(homogenous) 아키텍처인 Mamba를 제안한다.

Selective SSM, 그리고 이를 기반으로 한 Mamba 아키텍처는 완전한 recurrent model이며, 일반적인 시퀀스 기반 foundation model의 백본으로 적합한 여러 특성을 가진다.

  1. 높은 품질(High quality): 선택 메커니즘(selectivity)은 언어 및 유전체와 같은 정보 밀도가 높은 모달리티에서 강력한 성능을 제공한다.

  2. 빠른 학습 및 추론(Fast training and inference): 학습 시 계산량과 메모리 사용량이 시퀀스 길이에 대해 선형적으로 증가하며, 추론 시 autoregressive unrolling은 이전 요소들의 캐시가 필요 없으므로 step당 상수 시간(constant time)만 소요된다.

  3. 긴 문맥(Long context): 품질과 효율성이 결합되어 실제 데이터에서 최대 백만 길이(sequence length 1M)까지 성능 향상을 이끌어낸다.

우리는 Mamba가 범용 시퀀스 FM 백본으로서 잠재력을 가진다는 점을 여러 모달리티와 설정에서 실험적으로 검증하였다.

  • 합성 태스크(Synthetics). selective copying과 induction heads 같은 중요한 합성 태스크에서, Mamba는 문제를 쉽게 해결할 뿐 아니라 백만 토큰 이상의 길이까지 무한히 일반화(extrapolate)할 수 있다.

  • 오디오와 유전체(Audio and Genomics). Mamba는 오디오 waveform과 DNA sequence 모델링에서 SaShiMi, Hyena, Transformer 같은 기존 최고 성능 모델을 능가하며, pretraining 품질과 downstream metric 모두에서 우수한 결과를 보였다. 특히 긴 문맥을 사용할수록 최대 백만 길이까지 성능이 향상되었다.

  • 언어 모델링(Language Modeling). Mamba는 Transformer 수준의 성능을 진정으로 달성한 최초의 선형 시간(linear-time) 시퀀스 모델이다. 최대 1B 파라미터 규모의 scaling law 실험에서, Mamba는 LLaMa(Touvron et al. 2023) 기반의 강력한 Transformer 학습 레시피들을 포함한 다양한 baseline을 능가하였다. Mamba language model은 유사한 크기의 Transformer보다 5배 높은 생성 처리량(generation throughput)을 가지며, Mamba-3B는 자기 크기의 두 배에 달하는 Transformer와 동등한 품질을 달성하였다(예: Pythia-3B 대비 commonsense reasoning 평균 4점 향상, 심지어 Pythia-7B를 초과).

모델 코드와 사전학습된 체크포인트는 https://github.com/state-spaces/mamba 에 공개되어 있다.

아래는 2 State Space Models 섹션 번역입니다. 수식은 렌더링되도록 $$로 감쌌습니다.

2 State Space Models

구조화된 상태 공간 시퀀스 모델(Structured State Space Sequence Models, S4)은 딥러닝을 위한 비교적 최근의 시퀀스 모델 계열로, 넓게는 RNN, CNN, 그리고 고전적인 상태 공간 모델과 관련이 있다. 이들은 1차원 함수 또는 시퀀스 $$x(t) \in \mathbb{R} \mapsto y(t) \in \mathbb{R}$$ 를 암묵적인 잠재 상태 $$h(t) \in \mathbb{R}^N$$ 를 통해 매핑하는 특정 연속 시스템 (1)에서 영감을 받았다.

구체적으로, S4 모델은 네 개의 파라미터 $$(\Delta, \mathbf{A}, \mathbf{B}, \mathbf{C})$$ 로 정의되며, 이들은 두 단계에 걸쳐 sequence-to-sequence 변환을 정의한다.

Discretization. 첫 번째 단계는 “연속 파라미터” $$(\Delta, \mathbf{A}, \mathbf{B})$$ 를 “이산 파라미터” $$(\overline{\mathbf{A}}, \overline{\mathbf{B}})$$ 로 변환한다. 이 변환은 고정된 공식

$$
\overline{\mathbf{A}} = f_A(\Delta, \mathbf{A})
$$

및

$$
\overline{\mathbf{B}} = f_B(\Delta, \mathbf{A}, \mathbf{B})
$$

를 통해 이루어지며, 이때 쌍 $$(f_A, f_B)$$ 를 이산화 규칙(discretization rule)이라고 부른다. 식 (4)에 정의된 zero-order hold, 즉 ZOH와 같은 다양한 규칙을 사용할 수 있다.

이산화는 연속시간 시스템과 깊은 관련이 있으며, 이를 통해 해상도 불변성(resolution invariance)과 같은 추가적인 성질을 부여하거나, 모델이 적절히 정규화되도록 자동으로 보장할 수 있다. 또한 이는 RNN의 게이팅 메커니즘과도 관련이 있으며, 이 점은 Section 3.5에서 다시 살펴볼 것이다. 그러나 기계적인 관점에서 보면, 이산화는 단순히 SSM의 forward pass 계산 그래프에서 첫 번째 단계로 볼 수 있다. SSM의 다른 변형들은 이산화 단계를 우회하고 대신 $$(\overline{\mathbf{A}}, \overline{\mathbf{B}})$$ 를 직접 파라미터화할 수도 있으며, 이는 더 이해하기 쉬울 수 있다.

Computation. 파라미터가

$$
(\Delta, \mathbf{A}, \mathbf{B}, \mathbf{C})
\mapsto
(\overline{\mathbf{A}}, \overline{\mathbf{B}}, \mathbf{C})
$$

로 변환된 후, 모델은 두 가지 방식으로 계산될 수 있다. 하나는 선형 recurrence (2)로 계산하는 방식이고, 다른 하나는 전역 convolution (3)으로 계산하는 방식이다.

일반적으로 모델은 효율적이고 병렬화 가능한 학습을 위해 convolutional mode (3)를 사용한다. 이 경우 전체 입력 시퀀스를 미리 볼 수 있다. 반면, autoregressive inference에서는 입력이 한 time-step씩 주어지므로, 효율적인 추론을 위해 recurrent mode (2)로 전환한다.

Linear Time Invariance (LTI). 식 (1)부터 (3)까지의 중요한 성질은 모델의 dynamics가 시간에 따라 일정하다는 점이다. 다시 말해 $$(\Delta, \mathbf{A}, \mathbf{B}, \mathbf{C})$$, 그리고 결과적으로 $$(\overline{\mathbf{A}}, \overline{\mathbf{B}})$$ 역시 모든 time-step에서 고정되어 있다. 이 성질을 linear time invariance, 즉 LTI라고 부르며, recurrence 및 convolution과 깊이 연결되어 있다. 비공식적으로, 우리는 LTI SSM을 임의의 선형 recurrence (2a) 또는 convolution (3b)과 동등한 것으로 생각하며, LTI라는 용어를 이러한 모델 계열을 포괄하는 상위 개념으로 사용한다.

지금까지 모든 구조화된 SSM은 LTI였다. 예를 들어, 이들은 convolution으로 계산되었다. 이는 Section 3.3에서 논의할 근본적인 효율성 제약 때문이다. 그러나 이 논문의 핵심 통찰은 LTI 모델이 특정 유형의 데이터를 모델링하는 데 근본적인 한계를 가진다는 점이며, 우리의 기술적 기여는 효율성 병목을 극복하면서 LTI 제약을 제거하는 데 있다.

LTI 정리를 한번 보고가자

LTI가 뭔데? 진짜 쉽게 설명해보자

논문에서 말하는 Linear Time Invariance, LTI는 한마디로 이거야.

“모델이 모든 시점에서 똑같은 규칙으로 정보를 처리한다.”

즉, 1번째 단어를 볼 때도, 10번째 단어를 볼 때도, 1000번째 단어를 볼 때도 처리 규칙이 변하지 않는다는 뜻이야.

1. 상태 공간 모델을 아주 단순하게 보면

SSM은 대충 이런 모델이야.

지금까지 본 정보를 hidden state라는 메모리에 저장해두고,
새 입력이 들어오면 메모리를 업데이트한 뒤,
그 메모리로 출력을 만든다.

예를 들어 문장을 왼쪽에서 오른쪽으로 읽는다고 해보자.

나는 / 오늘 / 학교에 / 갔다

모델은 단어를 하나씩 보면서 내부 메모리 h를 바꾼다.

h_1 = "나는"을 본 뒤의 메모리
h_2 = "나는 오늘"을 본 뒤의 메모리
h_3 = "나는 오늘 학교에"를 본 뒤의 메모리
...

논문 식으로는 이게 이런 형태야.

$$
h_t = \overline{\mathbf{A}}h_{t-1} + \overline{\mathbf{B}}x_t
$$

말로 풀면:

새 메모리 = 예전 메모리를 조금 바꾼 것 + 새 입력을 반영한 것

여기서 중요한 건 두 개야.

$$
\overline{\mathbf{A}}
$$

는 이전 메모리를 얼마나 유지하거나 바꿀지 정하는 규칙이고,

$$
\overline{\mathbf{B}}
$$

는 새 입력을 메모리에 얼마나 넣을지 정하는 규칙이야.

2. LTI는 “규칙이 항상 고정”이라는 뜻

LTI에서는 이 두 규칙이 시간에 따라 바뀌지 않아.

즉,

1번째 단어 처리할 때도 같은 A, B
2번째 단어 처리할 때도 같은 A, B
3번째 단어 처리할 때도 같은 A, B
...

항상 같은 규칙을 쓴다.

비유하면, 컨베이어 벨트 위에 물건이 지나가는데, 모든 물건을 똑같은 기계 설정으로 처리하는 거야.

사과 → 같은 기계 설정
돌멩이 → 같은 기계 설정
금덩이 → 같은 기계 설정
쓰레기 → 같은 기계 설정

기계가 “이건 중요한 물건이니까 잘 보관해야지”, “이건 쓰레기니까 무시해야지”처럼 내용을 보고 규칙을 바꾸지 못한다는 뜻이야.

이게 LTI의 핵심이야.

3. 왜 이름이 Linear Time Invariance일까?

쪼개보면 쉽다.

Linear는 선형이라는 뜻이야. 여기서는 계산이 대략 이런 식으로 생겼다는 뜻이야.

$$
h_t = \overline{\mathbf{A}}h_{t-1} + \overline{\mathbf{B}}x_t
$$

곱하고 더하는 형태지. 복잡한 조건문 같은 게 없어.

Time Invariance는 시간에 대해 변하지 않는다는 뜻이야.

즉, 시점이 바뀌어도

$$
\overline{\mathbf{A}}, \overline{\mathbf{B}}, \mathbf{C}
$$

같은 규칙이 고정되어 있어.

그래서 Linear Time Invariance는 쉽게 말해:

시간에 따라 바뀌지 않는 선형 규칙으로 시퀀스를 처리하는 모델

이라는 뜻이야.

4. 그럼 뭐가 문제야?

문제는 현실 데이터에서는 어떤 입력은 중요하고, 어떤 입력은 중요하지 않다는 거야.

예를 들어 이런 문장이 있다고 해보자.

내 비밀번호는 ... 음 ... 사실 ... 7392야

여기서 모델이 기억해야 할 핵심은 7392야.
음, 사실 같은 건 별로 중요하지 않아.

그런데 LTI 모델은 모든 시점에서 같은 규칙을 쓰니까,

"음"이 들어와도 같은 방식으로 메모리 업데이트
"7392"가 들어와도 같은 방식으로 메모리 업데이트

이렇게 돼.

물론 값 자체는 다르니까 결과가 완전히 같지는 않지만, 입력의 중요도에 따라 메모리 업데이트 방식 자체를 크게 바꾸는 능력은 약하다는 거야.

즉, LTI는 이런 식의 선택을 잘 못해.

이 토큰은 중요하니까 오래 기억하자.
이 토큰은 잡음이니까 무시하자.
여기서부터는 이전 내용을 리셋하자.

이게 Mamba 논문에서 말하는 기존 SSM의 한계야.

5. convolution이랑 왜 연결된다는 거야?

논문에서 LTI SSM이 recurrence나 convolution과 연결된다고 했지.

이것도 쉽게 말하면 이거야.

SSM을 한 칸씩 계산하면 recurrence야.

h_1 계산
h_2 계산
h_3 계산
...

그런데 규칙이 시간마다 똑같으면, 이 전체 과정을 convolution처럼 한 번에 계산할 수 있어.

즉 LTI라서 가능한 일이야.

왜냐하면 모든 위치에서 같은 필터, 같은 규칙을 쓰니까.

CNN에서 이미지에 같은 필터를 쭉 미는 것처럼, LTI SSM도 시퀀스 전체에 같은 커널을 적용하는 convolution으로 바꿔 계산할 수 있어.

그래서 기존 S4 같은 모델은 학습할 때 빠르게 계산하려고 convolution 형태를 많이 썼어.

요약하면:

LTI라서 규칙이 고정됨
→ 규칙이 고정되어 있으니 convolution으로 바꿀 수 있음
→ convolution은 병렬 계산이 쉬움
→ 학습이 빠름

좋은 점은 빠르다는 거야.

나쁜 점은 유연성이 떨어진다는 거고.

6. Mamba는 여기서 뭘 바꾼 거야?

Mamba의 핵심 아이디어는 이거야.

입력을 보고 SSM의 규칙을 바꾸자.

기존 SSM은 대충 이랬어.

모든 토큰에 같은 A, B, C 적용

Mamba는 이렇게 바꿔.

현재 입력 x_t를 보고 B_t, C_t, Δ_t를 다르게 만들자

즉,

중요한 토큰이 들어오면 기억을 강하게 업데이트
별로 안 중요한 토큰이면 거의 무시
문맥이 끊기면 상태를 리셋

이런 식으로 작동할 수 있게 만든 거야.

그래서 Mamba에서는 기존의 “항상 같은 규칙”인 LTI 제약을 깨고, 선택적 selective SSM을 만든다.

7. 한 줄로 정리하면

LTI는:

모든 시간 step에서 같은 규칙으로 메모리를 업데이트하는 방식

장점은:

convolution으로 바꿔서 빠르게 계산할 수 있다.

단점은:

입력 내용에 따라 “기억할지, 무시할지, 리셋할지”를 유연하게 선택하기 어렵다.

Mamba는:

입력에 따라 SSM의 규칙을 바꾸게 만들어서, 중요한 정보는 기억하고 쓸모없는 정보는 버릴 수 있게 만든 모델

이렇게 보면 돼.

진짜 비유로 끝내면:

기존 LTI SSM = 모든 손님에게 같은 말만 하는 점원
Mamba = 손님이 누군지 보고 다르게 응대하는 점원

그래서 Mamba가 “selective”하다고 부르는 거야.

Structure and Dimensions. 마지막으로, 구조화된 SSM이 “structured”라고 불리는 이유는 이들을 효율적으로 계산하려면 $$\mathbf{A}$$ 행렬에 구조를 부여해야 하기 때문이라는 점을 언급한다. 가장 널리 사용되는 구조는 diagonal 구조이며, 본 논문에서도 이를 사용한다.

이 경우 행렬

$$
\mathbf{A} \in \mathbb{R}^{N \times N},
\qquad
\mathbf{B} \in \mathbb{R}^{N \times 1},
\qquad
\mathbf{C} \in \mathbb{R}^{1 \times N}
$$

은 모두 $$N$$ 개의 숫자로 표현될 수 있다. batch size가 $$B$$, 길이가 $$L$$, 채널 수가 $$D$$ 인 입력 시퀀스 $$x$$ 에 대해 동작하기 위해, SSM은 각 채널에 독립적으로 적용된다. 이 경우 전체 hidden state는 각 입력마다 $$DN$$ 차원을 가지며, 이를 시퀀스 길이 전체에 대해 계산하려면 시간과 메모리가

$$
O(BLDN)
$$

만큼 필요하다. 이것이 Section 3.3에서 다루는 근본적인 효율성 병목의 원인이다.

General State Space Models. state space model이라는 용어는 매우 넓은 의미를 가지며, 단순히 잠재 상태를 가진 임의의 recurrent process라는 개념을 나타낸다. 이 용어는 여러 분야에서 서로 다른 개념을 가리키는 데 사용되어 왔다. 예를 들어, 강화학습에서는 Markov decision process, MDP를, 계산신경과학에서는 dynamic causal modeling, DCM을, 제어 분야에서는 Kalman filter를, 머신러닝에서는 hidden Markov model, HMM과 linear dynamical system, LDS를, 그리고 딥러닝에서는 일반적인 recurrent 모델 및 때로는 convolutional 모델을 가리키는 데 사용된다.

이 논문 전체에서 우리는 “SSM”이라는 용어를 오직 structured SSM 또는 S4 모델 계열을 가리키는 데 사용하며, 이 용어들을 서로 바꾸어 사용한다. 편의상, 선형 recurrence 관점 또는 전역 convolution 관점 중 하나에 초점을 맞춘 파생 모델들도 포함할 수 있으며, 필요한 경우에는 세부적인 차이를 명확히 설명한다.

SSM Architectures. SSM은 end-to-end neural network architecture에 통합될 수 있는 독립적인 sequence transformation이다. 우리는 때때로 SSM architecture를 SSNN이라고 부르기도 하는데, 이는 SSM layer와 SSM architecture의 관계가 linear convolution layer와 CNN의 관계와 유사하기 때문이다. 여기서는 가장 잘 알려진 몇 가지 SSM architecture를 논의하며, 이들 중 다수는 본 논문의 주요 baseline으로도 사용된다.

• Linear attention은 recurrence를 포함하는 self-attention의 근사로, 퇴화된 linear SSM으로 볼 수 있다.

• H3는 이 recurrence를 일반화하여 S4를 사용한다. 이는 두 개의 gated connection 사이에 SSM이 끼워진 architecture로 볼 수 있다. H3는 또한 main SSM layer 앞에 표준 local convolution을 삽입하며, 이를 shift-SSM으로 해석한다.

• Hyena는 H3와 동일한 architecture를 사용하지만, S4 layer를 MLP로 파라미터화된 global convolution으로 대체한다.

• RetNet은 architecture에 추가적인 gate를 더하고 더 단순한 SSM을 사용한다. 이를 통해 convolution 대신 multi-head attention, MHA의 변형을 사용하는 또 다른 병렬화 가능한 계산 경로를 허용한다.

• RWKV는 언어 모델링을 위해 설계된 최근의 RNN으로, 또 다른 linear attention 근사인 attention-free Transformer에 기반한다. 핵심 “WKV” 메커니즘은 LTI recurrence를 포함하며, 두 SSM의 비율로 볼 수 있다.

그 밖의 밀접하게 관련된 SSM 및 architecture들은 Appendix B의 확장된 related work에서 더 자세히 논의한다. 특히 S5, QRNN, SRU를 강조하는데, 우리는 이들을 본 논문의 핵심 selective SSM과 가장 밀접하게 관련된 방법으로 본다.
























Reference

댓글

댓글 쓰기