기술 요약: Raven - 희소 메모리 라우팅을 통한 고재현율 시퀀스 모델링

1. 문제 정의

Mamba-2, GLA, GDN과 같은 상태 공간 모델(SSM) 및 선형 트랜스포머를 포함한 선형 시간 시퀀스 모델들은 긴 컨텍스트 재현(long-context recall)을 위해 메모리를 관리하는 방식에서 근본적인 트레이드오프에 직면해 있습니다. 기존 아키텍처들은 일반적으로 다음과 같은 두 가지 극단에 치우쳐 있으며, 둘 다 특정 과거 토큰을 검색하는 데 있어 구체적인 한계를 가집니다.

• 밀집 상태 업데이트 (SSM): Mamba-2, GLA, GDN과 같은 모델은 매 단계마다 밀집하게(densely) 업데이트되는 고정 크기의 상태를 유지합니다. 모든 메모리 슬롯은 새로운 토큰이 들어올 때마다 업데이트되며, 학습된 감쇠(decay) 메커니즘에 의존하여 정보를 망각합니다. 이는 원칙적으로는 무한한 지속성을 허용하지만, 밀집한 업데이트로 인해 간섭(interference)이 발생하며, 이로 인해 특정 개별 토큰이 후속 정보에 의해 덮어씌워지거나 희석되지 않고 보존되기가 어려워집니다.
• 희소 위치 기반 쓰기 (슬라이딩 윈도우 어텐션 - SWA): SWA는 명시적인 토큰 표현을 저장함으로써 희소하게(sparsely) 쓰기를 수행하지만, 이를 고정된 최근 윈도우 내로 제한합니다. 이는 FIFO(First-In-First-Out) 전략을 사용하여 위치에 따라 오래된 토큰을 강제로 퇴출(eviction)합니다. 이 방식은 윈도우 내에서는 신뢰할 수 있는 재현을 가능하게 하지만, 관련 토큰이 퇴출되면 해당 토큰의 중요도와 상관없이 재현율이 거의 0으로 떨어집니다.

핵심 문제는 기존 모델들이 정보가 어디에(where) 기록되는지와 그것이 얼마나 오래(how long) 지속되는지를 혼동한다는 점입니다. SSM은 모든 곳에 쓰고 감쇠에 의존하며, SWA는 위치에 따라 망각합니다. 두 방식 모두 콘텐츠 의존적인 희소 쓰기와 명시적이고 제어 가능한 감쇠를 효과적으로 결합하지 못합니다.

2. 방법론: 라우팅 슬롯 메모리 (RSMs)

저자들은 정보의 기록 위치와 지속 기간을 분리하는 통합 프레임워크인 **라우팅 슬롯 메모리(Routing Slot Memories, RSMs)**를 제안합니다. RSM은 고정된 $M$개의 메모리 슬롯을 유지하며, 각 타임스텝 $t$에서 **라우터(router)**를 사용하여 어떤 하위 집합의 슬롯이 새로운 정보를 수신하고 감쇠를 적용할지 결정합니다.

2.1 RSM 프레임워크

RSM의 일반적인 재귀식은 다음과 같이 정의됩니다:
$$S_t = (1 - r_t) \odot S_{t-1} + r_t \odot (D_t S_{t-1} A_t + U_t)$$
여기서:

• $S_t \in \mathbb{R}^{M \times d}$는 $M$개의 독립적인 슬롯을 가진 메모리 상태입니다.
• $r_t \in \mathbb{R}^M$은 각 슬롯의 쓰기 강도를 결정하는 라우팅 벡터입니다.
• $D_t$와 $A_t$는 감쇠 및 특징 혼합(feature mixing)을 제어합니다.
• $U_t$는 입력 $x_t$로부터 유도된 쓰기 내용입니다.

이 프레임워크는 기존 모델들을 하나의 스펙트럼 상에 배치합니다:

• SSM: 모든 슬롯을 점진적 감쇠와 함께 업데이트하는 밀집 라우터($r_t = \mathbf{1}_M$)를 사용합니다.
• SWA: 시간 의존적인 원-핫(one-hot) 라우터($r_t = e_t$)를 사용하여 정확히 하나의 슬롯을 업데이트하며, 강제 퇴출(감쇠 없음)을 수행합니다.
• GSA/ABC: 입력 의존적 라우터를 사용하지만 일반적으로 모든 슬롯을 업데이트(밀집 쓰기)하므로, 메모리를 격리하는 능력이 제한됩니다.

2.2 Raven 아키텍처

Raven은 희소한 입력 의존적 라우팅과 명시적 감쇠를 결합하여 "중간 지점"을 차지하도록 설계된 RSM의 구체적인 구현체입니다.

• 이중 상태 구조: SWA와 마찬가지로 Raven은 별도의 Key($S^k_t$) 및 Value($S^v_t$) 상태를 유지합니다.
• 희소 입력 의전적 라우팅: 고정된 윈도우나 밀집 업데이트 대신, Raven은 입력 $x_t$로부터 유도된 라우터 $g_t$를 사용합니다. 이는 (DeepSeek MoE 라우팅에서 변형된) Top-K 선택 전략을 채택하여 업데이트할 소수의 $K$개 슬롯을 선택합니다.
  • 점수는 시그모이드 함수를 통해 계산됩니다: $m_t = \sigma(W x_t)$.
  • 상위 $K$개의 항목만 유지되며, 나머지는 0으로 설정됩니다.
  • 라우터는 선택된 슬롯이 적절한 쓰기 강도를 받을 수 있도록 정규화됩니다.
• 선택적 감쇠: 감쇠는 선택된 슬롯에만 적용됩니다. 재귀식은 다음과 같습니다:
  $$S^k_t = \exp(a_t r_t) \odot S^k_{t-1} + (1 - \exp(a_t r_t)) \odot k^\top_t$$
  여기서 $a_t$는 입력 의존적 스칼라 감쇠(Mamba-2와 유사)입니다. 만약 특정 슬롯이 선택되지 않으면($r_t[i] \approx 0$), $\exp(a_t r_t[i]) \approx 1$이 되어, 해당 슬롯은 **동결(frozen)**되어 다시 선택될 때까지 그 내용을 완벽하게 유지합니다.
• 단순화된 블록 설계: 많은 선형 모델(예: Mamba, GLA)과 달리, Raven은 단거리 컨볼루션을 생략합니다. 로컬 및 장거리 컨텍스트 통합을 위해 오직 라우팅과 감쇠 메커니즘에만 의존하며, 결과적으로 더 단순한 블록 구조를 가집니다.

3. 주요 기여

1. 라우팅 슬롯 메모리 (RSMs): SSM과 SWA를 통합하고, "쓰기 위치"(라우팅)와 "망각 메커니즘"(감쇠)을 명시적으로 분리하는 공식적인 프레임워크를 제시했습니다.
2. Raven 모델: 입력 의존적 희소 라우팅을 사용하여 일부 메모리 슬롯을 업데이트하면서 동시에 업데이트된 슬롯에만 감쇠를 적용하는 RSM의 선형 시간 시퀀스 모델을 구현했습니다.
   casting 과정에서 특정 토큰을 전용 슬롯으로 라우팅하여 동결함으로써, 밀집 SSM에서 흔히 발생하는 파괴적 간섭과 SWA의 강제 퇴출 문제를 방지합니다.
3. 장거리 재현 메커니즘: 검색에 중요한 토큰(예: 패스키)을 전용 슬롯으로 라우팅하고 이를 동결함으로써, Raven은 밀집 SSM의 파괴적 간섭이나 SWA의 하드 퇴출을 방지합니다.
4. 길이 일반화 (Length Generalization): 명시적인 길이 커리큘럼이나 위치 보간 기술(예: RoPE) 없이도, 슬롯 간의 가변적인 유효 시퀀스 길이를 통해 다양한 시간적 역학에 노출됨으로써 훈련 길이의 최대 16배에 달하는 컨텍스트에서도 강력한 성능을 달성합니다.

4. 실험 결과

저자들은 Raven을 Mamba-2, GLA, GDN, SWA, Transformer 변형(FoX)을 포함한 강력한 베이스라인들과 비교하여 검색 집약적 벤치마크 및 일반 언어 모델링 작업에서 평가했습니다.

• Needle-in-a-Haystack (NIAH):
  • 400M 규모: Raven은 16K 토큰까지 거의 완벽한 정확도($\ge 99\%$)를 유지했으며, 32K 토큰(훈련 길이의 16배)에서도 강력한 성능($>91\%$)을 보였습니다. 반면 Mamba-2와 GDN은 8K 이후 성능이 크게 저하되었습니다.
  • 800M 규모: Raven은 64K 토큰에서 NIAH-1에 대해 91.0%의 정확도를 달성하여, 차순위 베이스라인(GDN 45.2%)을 크게 앞질렀습니다.
• 실제 세계 재현 작업:
  • SWDE(웹 추출), FDA(문서 추출), SQuAD에서 Raven은 선형 모델들을 일관되적으로 능가했습니다. 예를 들어, FDA에서 Raven은 선형 모델 중 유일하게 22% 이상의 정확도를 넘었으며, 최상의 SSM 베이스라인보다 약 8포인트 향상된 결과를 보여주었습니다.
• 언어 모델링:
  • Raven은 표준 제로샷 벤치마크(Lambada, PIQA, HellaSwag 등)에서 밀집 SSM(Mamba-2, GLA)과 대등하거나 더 나은 성능을 보여주어, 희소 라우팅이 일반적인 언어 모델링 능력을 저해하지 않음을 입증했습니다.
• 하이브리드 아키텍처:
  • 풀 어텐션 레이어와 결격된 하이브리드 모델로 통합되었을 때, Raven은 특히 16K~64K 범위의 긴 컨텍스트 재현 작업에서 Mamba-2, GDN, 또는 SWA-RoPE를 사용하는 하이브리드 모델보다 훨씬 뛰어난 성능을 보였습니다.

5. 의의 및 주장

본 논문은 Raven이 선형 시퀀스 모델링의 근본적인 아키텍처적 간극을 해결한다고 주장합니다. 정보가 어디에 기록되는지와 그것이 얼마나 오래 지속되는지를 명시적으로 분리함으로써, Raven은 이차(quadratic) 어텐션이나 복잡한 하이브리드 설정 없이도 선형 모델이 이전에 도달할 수 없었던 수준의 장거리 재현을 달enc합니다.

주요 주장은 다음과 같습니다:

• SWA 대비 우월성: Raven은 위치 기반 퇴출을 콘텐츠 의전적 라우팅과 선택적 감쇠로 대체하여, SWA가 완전히 실패하는 작업에서도 완벽한 재현을 가능하게 합니다.
• 밀집 SSM 대비 우월성: Raven은 선택적 쓰기를 통해 밀집한 상태 업데이트의 간섭를 피함으로써, 특정 토큰이 전용 슬롯에서 영구적으로 보존될 수 있도록 합니다.
• 단순성: Raven은 단거리 컨볼션 없이도 이러한 이점을 달성하여, 기존 선형 모델들에 비해 더 단순한 구조를 가집니다.
• 창발적 전문화 (Emergent Specialization): 분석 결과, Raven은 명시적인 감독 없이도 검색에 중요한 토큰(예: 비밀번호)을 위한 전용 슬롯을 할당하고 다른 슬롯을 일반적인 콘텐츠에 사용하는 방식으로 메모리 슬롯을 자연스럽게 할당하며, 구조화되고 콘텐츠 인지적인 메모리 할당을 학습함을 보여줍니다.

저자들은 RSM 프레임워크가 라우팅과 망각은 감쇠 메커니즘의 부산물이 아니라, 독립적인 설계 축으로서 다루어져야 한다는 점을 시사하며 결론을 맺습니다.