상태 공간 모델: 왜 선형 재귀가 긴 시퀀스에서 어텐션에 필적하는가?

어텐션 메커니즘은 긴 컨텍스트를 실용적으로 만들었지만, 비용을 많이 들게 했습니다. 새로운 토큰이 모든 이전 토큰을 되돌아보아야 하므로, 시퀀스의 길이에 따라 비용이 이차적으로 증가합니다. 상태 공간 모델(State Space Models)은 다른 방식을 취합니다. 과거를 고정된 크기의 상태로 압축하고, 이 상태를 선형 재귀(linear recurrence)로 업데이트하며, 그 상태가 미래에 필요한 정확한 정보만을 담도록 합니다. 잘 구현될 경우, 이러한 방식은 거의 선형적인 스케일링, 토큰당 상수 시간 생성, 그리고 매우 긴 시퀀스에서 놀라울 정도로 강력한 결과를 제공합니다.

이 글에서는 상태 공간 모델이 어디서 왔는지, 왜 단순한 재귀가 실패하는지, 그리고 현대 버전인 구조화된 상태 공간(structured state spaces)과 선택적 스캔(selective scans)이 어떻게 이를 해결했는지를 설명합니다. 예시로 사용되는 것은 긴 시퀀스 작업입니다: 1만 단계 후 신호의 다음 값을 예측하는 작업으로, 이 경우 어텐션은 느리고 일반적인 순환 신경망(recurrent network)은 정보를 잊어버립니다.

미분 방정식에서 시퀀스로

연속 시간 선형 상태 공간 모델은 입력 신호에 따라 진화하는 은닉 상태를 설명합니다:

h˙(t)=Ah(t)+Bx(t)
y(t)=Ch(t)+Dx(t)

여기서 h(t)는 N차원의 상태이며, x(t)는 입력이고, y(t)는 출력입니다. A, B, C, D는 행렬들입니다. 이 상태는 메모리 역할을 합니다: A의 동역학을 통해 필터링된 전체 입력 기록을 요약합니다.

실제 데이터는 이산적인 단계로 도착하므로, 모델은 이산화됩니다. 스텝 크기 Δ를 사용하면, 제로-오더 홀드(zero-order hold) 이산화는 다음을 제공합니다:

\bar{A}=exp(\Delta A)\n\bar{B}=(\Delta A)^{-1}(\exp(\Delta A)-I)\Delta B

그리고 이산 재귀:

h_t=\bar{A}h_{t-1}+\bar{B}x_t\ny_t=Ch_t

이 재귀가 핵심입니다. 새로운 토큰을 처리하는 비용은 시퀀스 길이에 독립적인 고정된 행렬-벡터 곱셈에 불과합니다. 생성 과정에서는 증가하는 키(key)와 값(value)의 캐시가 아니라 현재 상태만 필요합니다. 한편, 학습은 병렬화될 수 있습니다: 이 재귀는 선형적이기 때문에 전체 시퀀스에 대한 컨볼루션(convolution)으로 다시 작성될 수 있으며, 이는 최신 하드웨어에서 효율적으로 실행됩니다.

그렇다면 왜 이 오래된 제어 이론(control theory) 아이디어가 더 일찍 주류가 되지 못했을까요? 그 이유는 일반적인 행렬 A는 일반적인 메모리를 제공하며, 이러한 일반적인 메모리는 긴 시간 범위에 걸쳐 올바른 것을 기억하는 데 취약하기 때문입니다.

메모리 문제: 상태는 무엇을 기억해야 하는가?

입력이 긴 스트림이고 작업이 천 단계 전의 사실을 필요로 한다고 가정해 봅시다. 선형 재귀(linear recurrence)는 매 단계마다 상태를 $\bar{A}$로 곱합니다. 만약 $\bar{A}$의 고윳값(eigenvalues)이 1보다 훨씬 작다면, 오래된 정보는 지수적으로 감쇠하여 사라집니다. 만약 그 값이 1보다 크면, 상태가 폭발합니다. 오직 신중하게 구조화된 A만이 여러 시간 규모에서 정보를 동시에 살아있게 유지할 수 있습니다.

S4 모델에서 소개된 구조화된 상태 공간 접근 방식(structured state space approach)은 A를 초기화하여 HiPPO 연산자(HiPPO operator)를 근사함으로써 이 문제를 해결합니다. HiPPO는 정확한 질문에 답합니다: 시간 $t$까지의 신호 기록이 주어졌을 때, 최근 시간을 더 무겁게 가중하는 측정 하에서 해당 기록에 대한 최적 다항식 근사(polynomial approximation)의 계수는 무엇인가? 상태는 그저 희미해지는 메아리가 아니라, 과거에 대한 압축되고 다중 스케일적인 요약본을 문자 그대로 저장합니다.

실제로 S4는 A를 대각선(diagonal)에 저랭크 행렬(low-rank matrix)을 더한 형태로 매개변수화합니다. 이 구조는 HiPPO와 유사한 메모리 속성을 유지하는 동시에 컨볼루션 커널(convolution kernel)을 준선형 시간(near-linear time)에 계산 가능하게 만듭니다. 그 결과, 어텐션 기반 모델이 너무 느리거나 메모리를 많이 사용하는 영역이었던 오디오 파형, 장문 문서, 유전체 서열 등 수만 단계 길이의 원시 시퀀스를 처리하는 최초의 모델을 얻게 되었습니다.

하지만 S4에는 언어에 중요한 한계가 있었습니다: 그 동역학(dynamics)이 시간 불변적(time-invariant)이었다는 것입니다. 입력과 상관없이 모든 단계에서 동일한 A, B, C가 적용되었습니다. 이 모델은 내용에 기반하여 특정 토큰을 기억하고 다른 토큰을 무시할지 결정할 수 없었습니다. 어텐션은 정확히 그것을 할 수 있습니다: 선택합니다(selects). 고정된 선형 필터는 그럴 수 없습니다.

선택성(Selectivity): 입력이 무엇을 기억할지 선택하게 하기

Mamba 아키텍처는 시간 불변성 제약(time-invariance constraint)을 제거합니다. 이 모델의 파라미터들은 각 단계에서 입력에 대한 함수가 됩니다:

h_t = ar{A}t h{t-1} + ar{B}_t x_t

y_t = C_t h_t

여기서 $ar{A}_t$, $ar{B}_t$ 및 $C_t$는 일반적으로 작은 학습된 투영(learned projections)을 통해 $x_t$ 자체로부터 계산됩니다. 스텝 크기 $\Delta t$ 또한 입력 의존적이 됩니다: 큰 스텝은 과거를 빠르게 잊게 만들고, 작은 스텝은 과거를 유지하게 합니다.

직관적으로 이해하면 쉽습니다. 상태(state)를 제한된 페이지가 있는 공책이라고 생각해 보세요. 시간 불변 모델은 노트의 내용이 중요하든 아니든, 모든 들어오는 메모를 동일한 필체로, 동일한 일정에 따라 기록합니다. 반면, 선택적 모델은 각 메모마다 그것을 크게 쓸지, 작게 쓸지, 건너뛸지를 결정하고, 공간 확보를 위해 오래된 페이지를 지울지 여부를 결정합니다. 이러한 콘텐츠 기반 게이팅(content-based gating) 덕분에 선택적 상태 공간 모델은 이전의 순환 모델들이 해결하지 못했던 선택적 복사 및 유도(induction)-유사 작업을 수행할 수 있게 되었습니다.

단점은 컨볼루션 트릭(convolution trick)을 더 이상 적용할 수 없다는 것입니다. 시간 변화 동역학(time-varying dynamics)이 존재하기 때문에, 순환 구조를 고정된 컨볼루션으로 다시 쓸 수 없습니다. Mamba는 하드웨어 인지형 선택적 스캔(hardware-aware selective scan)을 통해 효율성을 회복합니다. 이는 상태를 빠르고 온칩 메모리(fast on-chip memory)에 유지하면서, 전체 중간 시퀀스를 절대 물질화하지 않고 순서 전체에 걸쳐 순환 구조를 병렬로 계산하는 연관형 스캔(associative scan)입니다. 이 알고리즘은 단순히 수학적 공식뿐만 아니라 모델 자체의 일부가 됩니다.

실용적인 예시: 긴 컨텍스트 복사 (long-context copying)

합성적이지만 매우 보여주는 작업 하나를 고려해 봅시다: 일련의 키-값 쌍(key-value pairs) 뒤에 쿼리 키(query key)가 오는 시퀀스이며, 모델은 일치하는 값을 출력해야 합니다. 어텐션(Attention)은 이 작업을 일치하는 쌍에 직접 주의를 기울임으로써 해결합니다. 반면, 시간 불변 상태 공간 모델은 모든 쌍을 동일한 고정 요약으로 압축하기 때문에 단일 쌍을 분리해낼 수 없어 어려움을 겪습니다.

선택적 모델(selective model)은 다르게 동작합니다. 쌍(pair)을 볼 때, 게이트를 열어 상태에 기록할 수 있고, 필러 토큰(filler tokens)을 볼 때는 게이트를 닫아 상태를 거의 건드리지 않고 유지할 수 있습니다. 질의 시간(query time)에 출력 투영(output projection)은 저장된 값을 읽어낼 수 있습니다. 코드로 보면, 단일 채널의 핵심 재귀(core recurrence)는 다음과 같습니다:

# h: 상태 벡터 (state vector), x_t: t 시점의 입력
# A, B, C, delta: t 시점의 입력 의존적 파라미터 (input-dependent parameters)
h = A_bar * h + B_bar * x_t
...

전체 모델은 이와 같은 많은 채널들을 쌓고 투영(projection) 및 게이팅 비선형성(gating nonlinearity)과 혼합하여 병렬로 스캔합니다. 실무자들에게 주는 교훈은 위의 두 줄 자체가 아니라 설계 패턴(design pattern)입니다. 즉, 고정 크기 상태에 콘텐츠 의존적 업데이트를 결합하면 2차원 어텐션 행렬 없이도 어텐션 유사 선택(attention-like selection) 형태를 구현할 수 있다는 것입니다.

상태 공간 모델이 강점을 갖는 영역과 그렇지 않은 영역

상태 공간 모델은 세 가지 축에서 강점을 가집니다:

• 길이 (Length). 생성 중 토큰당 비용이 일정하며, 메모리는 시퀀스 길이가 아닌 상태 크기(state size)에 의해 결정됩니다. 10만 단계 이상의 시퀀스가 실현 가능해집니다.
• 처리량 (Throughput). 스캔 공식화(scan formulation)는 긴 시퀀스를 높은 하드웨어 활용률로 처리하며, 추론 과정에서 장기 컨텍스트 어텐션(long-context attention)을 느리게 만드는 증가하는 키-값 캐시(key-value cache)를 피합니다.
• 연속적이고 불규칙한 데이터 (Continuous and irregular data). 모델이 미분 방정식(differential equation)에서 시작하기 때문에, 토큰 기반 어텐션에 추가적인 장치가 필요한 경우에도 불규칙하게 샘플링된 시계열(time series), 오디오 및 센서 스트림을 자연스럽게 처리합니다.

이들이 모든 면에서 이기는 것은 아닙니다. 작업이 긴 컨텍스트 내의 여러 임의적인 조회(arbitrary lookups)를 요구하는 경우, 예를 들어 흩어져 있는 몇 가지 사실을 검색하고 이를 결합해야 하는 경우에는 여전히 Attention이 우위를 점합니다. 그 이유는 압축된 상태(compressed state)가 설계상 병목 현상(bottleneck)이기 때문입니다. 차원 N의 상태는 길이 N보다 훨씬 큰 임의적인 시퀀스를 손실 없이 저장할 수 없으며, 유용한 요약 정보만을 저장할 수 있습니다. 따라서 상태 공간 레이어(state space layers)와 간헐적인 Attention 레이어를 혼합하는 하이브리드 아키텍처가 합리적인 기본값입니다. 즉, 시퀀스의 대부분에는 순환 구조(recurrence)를 사용하고 정밀한 검색이 중요한 부분에만 Attention을 사용하는 것입니다.

학습 안정성(Training stability) 또한 또 다른 실질적인 문제입니다. 이산화된 행렬 $\bar{A}$는 안정적으로 유지되어야 하며, 초기화(initialization)는 Transformer에서보다 더 중요합니다. 그리고 속도 주장을 충족시키기 위해서는 선택적 스캔(selective scans)의 구현에 세심한 주의가 필요합니다. Python으로 시간 단계(time steps)를 순진하게 반복하는 것은 느릴 것이며, 발표된 속도 향상은 순환 구조 자체에서 나오는 것이 아니라 융합 커널(fused kernels)에서 나옵니다.

한계점 및 일반적인 함정 (Limitations and common pitfalls)

• 상태는 데이터베이스가 아닌 병목 지점입니다. 고정 크기의 상태가 긴 문서를 그대로 암기할 것이라고 기대해서는 안 됩니다. 상태 공간 모델은 임의의 세부 사항에 대한 정확한 검색이 아니라, 요약되거나 필터링되었거나 예측된 정보를 기반으로 평가해야 합니다.
• 선택성은 보장되는 것이 아니라 학습되는 것입니다. Mamba가 작동하게 만드는 게이팅(gating) 메커니즘은 데이터로부터 학습되어야 합니다. 선택에 대한 훈련 신호가 적은 작업에서는, 선택적 모델이 고정된 모델과 매우 유사하게 행동할 수 있습니다.
• 길이가 무한하다는 의미는 아닙니다. 시퀀스가 훈련 길이를 훨씬 초과하여 증가함에 따라 성능은 여전히 저하되지만, 비용이 폭발하는 모델보다는 더 우아하게(gracefully) 저하됩니다. 항상 배포 길이(deployment length)에서 테스트해야 하며, 오직 훈련 길이에서만 테스트해서는 안 됩니다.
• 벤치마크가 오해를 줄 수 있습니다. 장거리 합성 벤치마크에서 이긴 모델이 여러분의 데이터에서는 우위를 점하지 못할 수 있습니다. 그 이유는 관련 범위(relevant horizon)가 짧을 수 있고, Attention의 정밀한 조회 기능이 더 중요할 수 있기 때문입니다. 실제 작업과 실제 길이에서 측정해야 합니다.

핵심 요약 (The takeaway)

상태 공간 모델(State Space Models)은 시퀀스 모델링을 필터링 문제로 재정의합니다. 즉, 하나의 상태를 유지하고, 이를 선형적으로 업데이트하며, 입력이 무엇이 들어오고 무엇이 사라질지 결정하게 합니다. 구조화된 초기화(Structured initialisation)는 상태에 좋은 다중 스케일 메모리(multi-scale memory)를 부여하고, 입력 의존적 선택성(input-dependent selectivity)은 선택할 수 있는 능력을 제공합니다. 그 결과는 어텐션(attention)의 보편적인 대체재가 아니라, 근본적으로 다른 트레이드오프입니다. 이는 선형적 스케일링(linear scaling)과 상수 시간 생성(constant-time generation)을 얻는 대신, 과거에 대한 완전한 기억이 아닌 압축된 메모리로 비용을 지불하는 것입니다. 긴 시퀀스에서는 그 트레이드오프가 종종 올바른 선택입니다.

LinkedIn에서 저와 연결하세요: linkedin.com/in/umeshtharukaofficial