Attention은 모든 새로운 토큰이 이전의 모든 토큰을 다시 돌아볼 수 있게 해주기 때문에 회상(recall)에 매우 탁월합니다. 하지만 그것은 또한 저주이기도 합니다. 매 단계마다 전체 이력을 다시 읽기 때문에, 실행 비용은 길이의 제곱처럼 상승하고 메모리(KV cache)도 그에 따라 증가합니다. 매우 긴 입력값에 대해서는 한계(wall)에 부딪히게 됩니다. State-space models (SSM, 상태 공간 모델)은 다른 길을 택합니다. 즉, 토큰별로 업데이트되는 고정된 크기의 실행 상태(running state)를 사용하는 것입니다. 그리고 Mamba는 마침내 이 모델들을 경쟁력 있게 만든 버전입니다. 저는 두 모델을 실시간으로 실행하고 스트림을 진행함에 따라 두 비용 카운터를 계산하며, 직접 켜고 끌 수 있는 선택성 게이트(selectivity gate)가 포함된 실제 SSM 재귀(recurrence)를 계산하는 데모를 만들었습니다. 그 아이디어는 다음과 같습니다.

이차적 장벽(Quadratic wall) vs 고정된 상태(Fixed state)

Attention의 토큰당 비용은 위치 번호 t입니다 (이전의 t개 토큰을 모두 다시 읽기 때문). 따라서 n개 토큰 이후의 총 비용은 n(n+1)/2 ≈ n²/2가 되며, 모든 n개의 키(keys)와 값(values)을 캐싱해야 합니다. 반면 SSM은 크기가 d인 고정된 상태만을 건드리므로, 토큰당 비용은 일정(constant)하며, 총 비용은 d·n(선형, linear)이고, 메모리는 영원히 d로 유지됩니다:

// ATTENTION: 매 토큰마다 전체 이력을 다시 읽음
const attnTotal  = n => n * (n + 1) / 2;   // ~ n^2 / 2  -> 장벽(the wall)
const attnMemory = n => n;                  // KV cache가 길이에 따라 증가
...

데모를 통해 토큰 스트림을 단계별로 진행하거나 재생하면서, Attention의 빨간색 곡선이 SSM의 직선인 초록색 선으로부터 멀어지는 것을 관찰할 수 있습니다. 길이를 늘리면 길이가 두 배가 될 때마다 Attention의 비용은 대략 _4배_로 증가하지만, SSM의 비용은 단지 _2배_만 증가합니다. 그 격차는 거대한 심연이 되며, 메모리 출력값을 보면 Attention의 KV cache는 n개의 항목으로 성장하는 반면 SSM은 d로 일정하게 유지되는 것을 볼 수 있습니다.

재귀(Recurrence)는 단 한 줄입니다

SSM의 핵심은 지속적으로 감쇠하는 요약(decaying summary)입니다. 이전 상태의 일부 a를 유지하고, 새로운 입력 x의 b배를 더하며, c를 통해 출력을 읽어냅니다. 이력은 저장되지 않습니다. 모델이 "아는" 모든 것은 하나의 숫자 h로 압축됩니다:

function ssmStep(h, x, a, b){ return a * h + b * x; }  // h_t = a*h + b*x
// 'a'가 1에 가까우면 = 긴 기억력 (long memory), 0에 가까우면 = 빠르게 망각 (forget fast)
// 'b'        = 새로운 입력이 상태 (state)에 얼마나 강하게 들어오는지

스트리밍 (Streaming)은 단순히 루프 (loop)일 뿐입니다: 각 토큰에 대해, 단일 상태 (single state)를 업데이트하고 출력을 내보냅니다. 과거 토큰들이 늘어나는 배열은 존재하지 않으며 — 오직 h만이 단계 사이에 살아남습니다 — 이것이 바로 SSM이 일정한 메모리 (constant memory)로 백만 개의 토큰 스트림을 처리할 수 있는 핵심 이유입니다.

왜 일반적인 SSM만으로는 충분하지 않았는가 — 그리고 Mamba의 해결책

고정된 SSM은 모든 토큰에 대해 동일한 a, b를 사용하므로, 상태 (state)가 하나의 고정된 비율로 누출됩니다. 이 모델은 "이 토큰은 중요하니 유지해라" 또는 "새로운 주제이니 이전 내용은 버려라"라고 말할 방법이 없습니다. 노이즈 (noise)의 급증이 중요한 사실과 정확히 동일한 속도로 감쇠합니다. 이러한 입력에 대한 무지 (input-blindness) 때문에, 고전적인 선형 SSM (linear SSMs)과 RNN은 회상 (recall) 중심의 작업에서 어텐션 (attention)에 뒤처졌습니다.

Mamba의 핵심적인 변화는 **선택성 (selectivity)**입니다: 파라미터 (parameters)를 현재 토큰의 함수로 만드는 것입니다. 이제 모델은 스스로의 메모리를 게이팅 (gating)합니다 — "새로운 주제" 토큰에서는 a→0으로 유도하여 상태를 비우고 (flush); "중요한" 토큰에서는 a→0, b→1로 유도하여 값을 깔끔하게 기록하며; "무관한" 토큰에서는 a→1, b→0으로 유도하여 기존 내용을 유지하고 입력을 무시합니다:

function selectiveA(tok, baseA){
  if (tok.type === "reset") return 0.04;  // "과거를 잊어라"  -> 비우기 (flush)
  if (tok.type === "keep")  return 0.05;  // "이 값을 기록하라" -> 이전 내용 삭제 (clear old)
...

데모에서는 리셋 (reset), 유지 (keep), 노이즈 (noise)와 같은 특수 토큰이 포함된 토큰 스트림에 대해 실제 재귀 (recurrence)를 실행하고, 선택적 상태 (selective state)와 비선택적 상태 (non-selective state)를 두 개의 선으로 그래프로 나타냅니다. 제어 토큰 (control tokens)에서 어떤 일이 일어나는지 살펴보세요: 선택적 선은 급격히 변합니다 (snaps) — 리셋 시에는 비워지고 유지 시에는 고정됩니다 — 반면 비선택적 선은 중요한 토큰과 노이즈를 구분하지 못한 채 하나의 고정된 비율로 계속 누출될 뿐입니다. 이러한 입력 의존적 게이팅 (input-dependent gating)이 Mamba의 핵심 아이디어이며, 이것이 고정된 크기의 상태가 마침내 어텐션 (attention)과 경쟁할 수 있게 만든 원동력입니다.

한 가지 걸림돌은 입력 의존적 파라미터(input-dependent parameters)가 선형 SSM (linear SSMs)의 학습을 빠르게 만들었던 단순한 컨볼루션 트릭 (convolution trick)을 깨뜨린다는 점입니다. 따라서 Mamba는 학습 중에 GPU에서 모든 상태를 한 번에 계산하기 위해 하드웨어 인지 병렬 스캔 (hardware-aware parallel scan)을 사용하며, 추론 시에는 단순한 $O(1)$ 메모리 재귀 (recurrence) 방식으로 동작합니다.

지속되는 트레이드오프 (trade-off)

어텐션 (Attention)은 모든 것을 다시 읽음으로써 $n^2$의 비용과 증가하는 메모리를 대가로 정밀한 회상 (pinpoint recall) 능력을 얻습니다. 반면 SSM은 이를 고정된 크기의 상태 (fixed-size state)와 맞바꿉니다. 즉, 선형 시간 (linear time), 일정 메모리 (constant memory), 그리고 무한히 스트리밍할 수 있는 능력을 얻는 것입니다. 회상 능력은 어텐션의 강점이며, 효율성은 SSM의 강점입니다. 이것이 바로 롱 컨텍스트 (long-context), 스트리밍 (streaming), 온디바이스 (on-device) 시스템이 SSM에 의존하는 이유이며, 최첨단 기술들이 회상을 위한 약간의 어텐션과 저렴한 SSM 레이어를 함께 유지하는 하이브리드 (hybrids) 모델을 점점 더 추구하는 이유입니다.

스트림을 진행시키며 어텐션의 비용이 폭발하는 동안 SSM은 평탄하게 유지되는 것을 확인해 보세요:
https://dev48v.infy.uk/ai/days/day43-state-space-models.html