기술 요약: DREAMSTATE

문제 정의

현대적인 순환 신경망(RNN), 특히 RWKV와 같은 아키텍처는 표준 트랜스포머(Transformer)에 비해 추론 효율성과 고정된 크기의 상태 표현 측면에서 상당한 이점을 제공합니다. 그러나 본 논문은 현재의 RNN 설계에서 두 가지 결정적인 한계를 식별합니다:

1. 활용되지 못한 상태 표현 잠재력: RNN의 상태는 단기 모델링에는 강력하지만, 편집 가능한 지식 표현으로서 다루어지는 경우는 드뭅니다. 내부 상태 분포를 직접 생성하거나 조작할 수 있도록 모델링하는 것에 대한 연구가 부족합니다.
2. 정적 재귀의 구조적 노이즈: 트랜스포머에서 모델이 무관한 문맥에 과도하게 주의를 할당하는 "어텐션 노이즈(attention noise)"와 유사하게, RNN은 "구조적 노이즈(structural noise)"를 겪습니다. RWKV의 핵심 재귀 메커니즘은 훈련 후 고정되는 정적 WKV 파라미터($\theta_{WKV}$)에 의존합니다. 이러한 정적 특성은 가중치 편향으로 작용하여 다양한 문맥에 최적화되지 못하며, 모델이 중요한 정보에 집중하기 위해 자신의 "운영 법칙"을 동적으로 조정하는 것을 방해합니다.

저자들은 만약 내부 상태가 문맥의 구조화된 표현이라면, 이를 직접 모델링하고 조작하는 것이 가능할 것이라고 주장합니다. 또한, 구조적 노이즈를 완화하기 위해 재귀 파라미터의 정적 가정을 도전적으로 검토해야 한다고 제안합니다.

방법론

본 논문은 이러한 문제를 해결하기 위해 생성형 모델링(특히 확산 모델, Diffusion Models)에 기반한 두 단계 접근 방식인 DREAMSTATE 프레임워크를 제안합니다.

1. 생성적 상태 합성 (DREAMSTATE 프레임워크)

저자들은 먼저 RWKV 내부 상태를 학습 가능한 매니폴드($\mathcal{M}$)에서 샘플링된 데이터 포인트로 취급합니다.

• 모델: 조건부 디퓨전 트랜스포머(Conditional Diffusion Transformer, DiT)가 텍스트 임베딩 $c$에 조건화된 확률 분포 $p(S|c)$를 모델링하도록 훈련됩니다 (여기서 $S$는 상태).
• 과정: 고차원 상태 텐서(플래튼된 멀티 헤드 상태)는 일련의 패치(patches)로 취급됩니다. DiT는 입력 문맥에 조건화되어 상태에 추가된 노이즈를 예측하는 법을 배웁니다.
• 목표: 상태 매니폴드가 잘 정의되어 있음을 검증하고, 특정 운영 모드(예: 특정 페르소나)를 위한 제어된 상태 생성을 가능하게 합니다.

2. 동적 파라미터 합성 (하이브리드 아키텍처)

구조적 노이즈를 해결하기 위해, 저자들은 정적 재귀 파라미터가 동적으로 생성되는 새로운 하이브리드 아키텍처를 제안합니다.

• 메커니즘: "파라미터 DiT(Parameter DiT)"는 가변 길이 입력 문맥으로부터 글로벌 특징(초기 구현에서는 첫 번째 토큰 $x[:,0,:]$의 임베딩)을 추적하여 문맥 특화 파라미터 벡터 $\theta_{gen}$을 생성합니다.
• 대상 파라미터: DiT는 WKV 파라미터를 구성하는 플래튼된 선형 투영 행렬($W_r, W_k, W_v$)을 생성합니다.
• 융합: 생성된 파라미터는 보간법을 통해 원래의 정적 베이스 파라미터($\theta_{static}$)와 융합됩니다:
  $$\theta_{WKV-final} = \alpha \cdot \theta_{static} + (1 - \alpha) \cdot \theta_{gen}$$
  여기서 $\alpha$는 학습 가능한 하이퍼파라미터입니다. 이를 통해 모델은 안정적인 지식을 유지하면서도 문맥 특화된 변조를 도입할 수 있습니다.
• 훈련: 시스템은 다중 목적 손실 함수를 사용하여 엔드투엔드(end-to-end)로 훈련됩니다:
  $$L_{total} = \lambda_1 L_{LLM}(\theta_{WKV-final}) + \lambda_2 L_{param\_diff}(\psi)$$
  이는 다음 토큰 예측(표준 교차 엔트로피)과 파라미터 생성을 위한 디퓨전 손실 사이의 균형을 맞춥니다.

주요 기여

1. DREAMSTATE 프레임워크: RWKV 모델의 동적인 상태를 생성적으로 모델링하는 방법을 제시합니다. 저자들은 내부 상태가 서로 다른 개념(예: "프로그래머" 대 "시인")이 뚜렷한 클러스터를 형성하는 구조화된 매니폴드임을 입증하여, 상태 조작을 통한 제어된 추론을 가능하게 함을 검증했습니다.
2. 동적 파라미터 아키텍처: 디퓨전 모델이 글로벌 문맥을 기반으로 재귀 파라미터를 생성하는 새로운 하이브리드 설계를 제안합니다. 이는 RNN의 운영 법칙을 정적인 것이 아닌 동적인 것으로 만듦으로써 "구조적 노이즈"에 대응하기 위한 원칙적인 방법으로 제안되었습니다.
3. 다중 목적 훈련 패러다임: 파라미터 생성과 언어 모델링을 공동으로 최적화하는 훈련 전략을 정의하고 성공적으로 입증함으로써, 이러한 하이브리드 접근 방식의 실현 가능성을 증명했습니다.
4. 개념 단위 자기회귀(Concept-Wise Autoregression): 본 연구는 전통적인 토큰 단위 자기회귀와 "개념 단위 자기회귀"(상태/파라미터 생성을 통해)를 결리하는 것이 미래 모델 설계의 유효한 방향임을 보여주는 개념 증명(proof-of-concept)을 제공합니다.

실험 결과

• 상태 매니폴드 시각화: 저자들은 t-SNE를 사용하여 실제 RWKV-7 상태가 특정 전문가 페르소나(예: 리눅스 터미널, 이더리움 개발자, 스토리텔러)에 대응하는 뚜렷한 클러스터를 형성함을 보여주었습니다. 이는 상태가 구조화되고 잘 정의된 객체임을 확인시켜 주었습니다.
• 제어된 추론: DREAMSTATE(또는 보간된 상태)로 생성된 상태로 RWKV-7 모델을 초기화함으로써, 모델은 베이스라인에 비해 훨씬 더 주제적으로 일관되고 창의적인 텍s트를 생성했습니다. 예를 들어, "인내(perseverance)" 프롬프트와 "사막의 기후 변화(desert climate change)" 프롬프트 사이를 보간하면 두 개념을 은유적으로 혼합한 이야기를 만들어냈습니다.
• 훈련 안정성: 하이브리드 아키텍처는 The Pile의 일부 데이터셋을 사용하여 성공적으로 훈련되었습니다. 훈련 손실 곡선은 언어 모델링 손실과 파라미터 디퓨전 손실 모두에서 안정적이고 일관된 감소를 보여주었으며, 이는 공동 최적화 과정의 안정성을 확인해 줍니다.

의의 및 주장

본 논문은 RNN 상태 표현에 대한 새로운 연구 방향을 열었다고 주장하며, 상태를 단순한 일시적 변수가 아닌 조작 가능한 확률적 객체로 바라보게 합니다.

• 제어 가능성: 본 연구는 내부 상태를 직접 생성하고 조작하여 모델을 특정 운영 모드로 준비(prime)시킬 수 있음을 보여주며, 이는 추론에 대한 전례 없는 통제력을 제공합니다.
• 적응성: 재귀 파라미터를 글로벌 문맥 추적기에 의해 생성되는 동적 변수로 취급함으로써, 이 아키텍처는 더 유연하고 표현력이 풍부한 월드 모델(world models)을 구축하는 것을 목표로 합니다.
• 아키텍처적 참조: 본 논문은 디퓨전 기반 파라미터 합성(diffusion-based parameter synthesis)과 재귀적 추론(recurrent inference)의 결합을 미래의 대규모 언어 모델(LLM)을 위한 구체적인 아키텍처 참조점으로 위치시키며, 상태를 모델링하는 "방법"이 LLM의 미래를 결정하는 핵심임을 시사합니다.

저자들은 자신들의 작업이 모든 벤치마크에서 기존 모델보다 즉각적으로 우월하다고 주장하기보다는, 동적 파라미터 합성을 통해 구조적 노이즈를 완화하는 설계의 타당성을 검증하는 개념 증명 역할을 한다고 강조합니다.