RNN에서 Mamba까지: 시퀀스 모델 진화의 단일 원리

2017년 Transformer가 등장한 이후 NLP 아키텍처의 역사는 단순해 보인다 — RNN이 있었고, Transformer가 대체했다. 그런데 2023년 Mamba가 나타나 RNN과 닮은 구조로 Transformer를 능가하는 결과를 보였다. 왜 우리는 한 바퀴를 돌아 다시 순환 구조로 돌아왔는가?

문제의 핵심: Sequential Dependency

모든 이야기는 하나의 수식에서 시작한다.

$$h_t = f(h_{t-1}, x_t)$$

RNN의 forward pass는 이 의존 관계를 $T$번 반복한다. $h_t$를 계산하려면 $h_{t-1}$이 먼저 있어야 하고, $h_{t-1}$은 $h_{t-2}$를 기다린다. 이 chain의 길이가 정확히 $T$다.

GPU는 수천 개의 코어가 동시에 같은 연산을 다른 데이터에 적용하는 SIMT 구조다. RNN의 sequential dependency는 이 구조와 정면으로 충돌한다. $T = 1000$일 때 각 스텝에서 나머지 수천 개 코어는 대기한다. 이론적 GPU utilization은 $O(1/T)$다.

Transformer의 해법과 그 대가

Transformer의 self-attention은 이 문제를 정면 돌파한다.

$$\text{Attention}(Q, K, V) = \text{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}}\right) V$$

$QK^\top$ 행렬의 모든 $(i, j)$ 쌍은 서로 독립적으로 계산 가능하다. Critical path는 matmul 내부의 tree reduction에 해당하는 $O(\log T)$로 줄어든다. Vaswani 2017의 실험이 이를 실증한다 — WMT’14 En→De 학습에서 Google NMT(LSTM) 16일 대비 Transformer-Big은 3.5일, 약 5배 빠르고 BLEU도 높았다.

그러나 이 병렬성에는 명확한 대가가 따른다.

CNN 기반 접근(WaveNet, TCN)도 같은 방향을 시도했다. Dilated convolution을 $1, 2, 4, \ldots, 2^{L-1}$로 쌓으면 receptive field가 $O(2^L)$로 증가하고, 모든 포지션을 동시에 처리할 수 있다. Bai 2018의 TCN이 여러 시퀀스 태스크에서 LSTM을 능가했지만, receptive field가 아키텍처로 고정된다는 한계와 global attention의 부재로 Transformer에 자리를 내줬다.

Linear Attention: 연산 순서의 재발견

Transformer의 $O(T^2)$ 병목을 해결하는 핵심 통찰은 Katharopoulos 2020(“Transformers are RNNs”)에서 나왔다. Softmax를 kernel feature map $\phi$로 근사하면 연산 순서를 바꿀 수 있다.

$$\text{Attn}(Q, K, V) \approx \phi(Q)\,\bigl(\phi(K)^\top V\bigr)$$

괄호의 위치가 전부다. $\phi(K)^\top V$를 먼저 계산하면 $d \times d$ 행렬이 되고 $T \times T$ 행렬은 사라진다. 그리고 이 구조는 정확히 RNN의 update rule과 동형이다.

$$S_t = S_{t-1} + \phi(k_t)\,v_t^\top$$

$S_t$가 recurrent state다. Inference 시 KV-cache($O(T \cdot d)$) 대신 $S \in \mathbb{R}^{d \times d}$ 하나만 유지하면 된다. GPT-3 규모에서 $T = 2048$ KV-cache가 약 19 GB인 데 비해 linear attention state는 0.6 GB이며 $T$에 무관하다.

RWKV(Peng 2023)는 이 원리를 실제 LLM으로 구현했다. Time-mixing block이 지수적 시간 감쇠를 적용한 recurrent state를 유지하면서 Transformer-quality 학습을 달성했다.

SSM: 이론적 토대

State Space Model(SSM)은 제어이론의 고전적 개념을 딥러닝으로 흡수한 결과다.

$$\dot x(t) = A\,x(t) + B\,u(t), \qquad y(t) = C\,x(t)$$

HiPPO(Gu 2020)의 핵심 통찰은 이 ODE의 최적 해가 Legendre polynomial projection과 일치한다는 것이다 — 즉 특정 $A$ 행렬 구조가 “지금까지 본 입력의 최적 압축”을 보장한다.

S4(Gu 2022)는 이를 효율적으로 학습 가능하게 만들었다. $A$를 diagonal + low-rank로 구조화하면 Cauchy kernel을 통해 길이 $T$의 convolution kernel을 $O(N \log T)$에 계산할 수 있다. 학습은 FFT 기반 convolution, inference는 $O(N)$ recurrence — 두 모드를 동시에 지원한다.

Mamba(Gu & Dao 2023)는 S4의 남은 한계를 해결한다. S4의 $A, B, C$는 input-independent — 모든 토큰에 같은 dynamics가 적용된다. Mamba는 이를 selective하게 만든다.

$$A_t,\, B_t,\, C_t = f(x_t)$$

관련 있는 정보는 $A_t \approx I$로 state에 보존하고, 무관한 정보는 $A_t \approx 0$으로 잊는다. 이것이 attention의 selectivity를 $O(T)$ 복잡도로 구현하는 방식이다. 그리고 linear recurrence는 associative binary operation이므로 Blelloch 1990의 parallel scan으로 $O(\log T)$ depth에 병렬 실행된다.

LRA Path-X 벤치마크(16K context, path connectivity 판별)에서 vanilla Transformer와 S4가 모두 random chance(50%)에 머무를 때 Mamba가 처음으로 의미 있는 성능을 기록했다.

트레이드오프

각 아키텍처의 강점과 한계는 명확하다.

Architecture	Time	Memory	Critical Path	In-context Learning
RNN/LSTM	$O(TH^2)$	$O(TH)$	$T$	약함
Transformer	$O(T^2H)$	$O(T^2)$	$O(\log T)$	강함
S4	$O(TN \log T)$	$O(TN)$	$O(\log T)$	약함
Mamba	$O(TN)$	$O(TN)$	$O(\log T)$	중간

Linear attention의 한계도 명확하다. Softmax attention이 임의의 attention pattern을 표현할 수 있는 데 비해, linear attention은 feature map dimension $d_\phi$으로 제한된 low-rank approximation이다. Sharp, sparse attention pattern 표현이 어렵고 이것이 in-context learning의 약점으로 이어진다.

정리

• RNN의 sequential dependency는 GPU utilization을 $O(1/T)$로 만든다. Transformer는 $O(T^2)$ 메모리 비용으로 이 한계를 돌파했다.
• Linear attention은 연산 순서 변경으로 $O(T^2) \to O(T)$를 달성하며, 이 구조가 RNN-like recurrence와 동형임을 보였다.
• SSM(S4, Mamba)은 제어이론적 토대 위에서 selective state evolution과 hardware-aware parallel scan을 결합해 long context의 실질적 해법이 됐다.
• 아키텍처 진화의 원동력은 일관되다 — 어떻게 sequential dependency를 없애면서 selective한 정보 보존을 유지할 것인가.

이 챕터들이 보여주는 것은 아키텍처의 단순한 대체 역사가 아니다. RNN의 비선형 게이트를 제거하면 병렬화가 가능해지고, Transformer의 $T^2$ 항을 kernel trick으로 우회하면 RNN 형태가 돌아온다. 같은 문제의 다른 표현이 서로를 수렴시키고 있다.