기술 요약: 재귀적 상태 분해를 위한 WriteSAE

문제 제기

상태 공간 및 하이브리드 재귀 언어 모델 (예: Gated DeltaNet, Mamba-2, RWKV-7, Qwen3.5) 은 상태 전파를 위해 행렬 캐시를 활용하며, 여기서 각 토큰은 $d_k \times d_v$ 행렬에 랭크 1 외적 ($k_t v_t^\top$) 을 씁니다. 기존 희소 자동 인코더 (SAE) 는 잔여 스트림에서 작동하여 방출 후의 상태를 읽습니다. 따라서 표준 SAE 는 재귀의 상류에서 발생하는 네이티브 쓰기 메커니즘에 직접 접근하거나 분해할 수 없습니다. 또한 상태가 랭크 1 외적을 통해 업데이트되므로, 표준 벡터 기반 SAE 디코더 원자는 인과적 개입 (캐시 패칭) 에 필요한 구조적 대응 관계를 유지하면서 단일 캐시 쓰기 작업을 대체할 수 없습니다.

방법론

본 논문은 재귀 모델의 행렬 캐시 쓰기를 분해하고 편집하도록 설계된 최초의 희소 자동 인코더인 WriteSAE를 소개합니다.

아키텍처 및 훈련

• 네이티브 형태 매칭: WriteSAE 디코더 원자는 호스트 아키텍처의 네이티브 쓰기 원시 연산 ($k_t v_t^\top$) 과 일치하도록 랭크 1 외적 ($v_i w_i^\top$) 으로 제한됩니다. 이를 통해 단일 원자가 특정 캐시 슬롯 업데이트를 대체할 수 있습니다.
• 이선형 인코더: 행렬 상태 $S_t$를 효율적으로 인코딩하기 위해, 모델은 평탄화된 밀집 인코더 대신 이선형 인코더 ($a_i = v_i^\top S_t w_i$) 를 사용합니다. 이는 원자당 파라미터 수를 $d_k d_v$ ($128 \times 128$인 경우 16,384) 에서 $d_k + d_v$ (256) 로 줄여 64 배 감소시킵니다.
• 훈련 목적 함수: 모델은 TopK 희소성 제약과 죽은 특징을 부활시키기 위한 보조 손실 함수를 사용하여 평균 중심 상태 ($x = \text{vec}(S_t - M)$) 로 훈련됩니다. 훈련 코퍼스는 특정 레이어와 헤드를 통해 Qwen3.5-0.8B(및 기타 모델) 로 처리된 OpenWebText 시퀀스로 구성됩니다.

특징 분류

훈련된 사전은 원자의 쓰기 방향과 네이티브 쓰기 방향 사이의 중앙값 코사인 유사도에 따라 두 가지 클래스로 분류됩니다:

1. 레지스터 (Registers): 네이티브 쓰기와 높은 코사인 정렬을 보이는 원자 (캐시에서 복원 가능).
2. 번들 (Bundles): 캐시 전반에 걸쳐 분산된 쓰기 방향을 가진 원자.

인과적 개입 테스트

본 논문은 세 가지 주요 검증 메커니즘을 제안합니다:

1. 캐시 슬롯 대체: 네이티브 쓰기 ($k_t v_t^\top$) 를 매칭된 프로베니우스 노름 SAE 원자 ($\hat{S}_t = S_t - \Delta_{nat} + \Delta_{atom}$) 로 대체하고 하류 KL 발산을 측정합니다.
2. 폐형 로그잇 시프트: 토큰별 로그잇 시프트를 예측하는 세 가지 인자 표현식을 유도합니다: $\Delta \ell \approx G_{t_0 \to t} \langle w_i, q_t \rangle \langle v_i, W_U[\text{tok}] \rangle$. 여기서 $G$는 게이트 곱, $q_t$는 읽기 쿼리, $W_U$는 언임베딩 행입니다.
3. 직접 캐시 개입: 특정 위치에서 원자를 지우거나 설치하여 토큰 확률과 생성 행동의 변화를 관찰합니다.

주요 기여

1. 아키텍처 매칭 사전: 원자가 랭크 1 외적인 사전으로, 기존 잔여 스트림 SAE 에는 없었던 네이티브 캐시 쓰기를 직접 대체할 수 있는 기능을 제공합니다.
2. 대체 검증: 학습된 원자가 높은 충실도로 네이티브 쓰기를 대체할 수 있음을 입증했습니다. Qwen3.5-0.8B (레이어 9, 헤드 4) 에서 4,851 개의 발화 중 **92.4%**에서 매칭된 노름 제거 실험 (ablation) 보다 우수한 성능을 보였습니다.
3. 폐형 예측: 파라미터가 없는 세 가지 인자 공식으로 발화별 로그잇 시프트를 예측하며, 200 개의 원자별 $\epsilon$ 셀 전반에서 중앙값 $R^2$가 0.98입니다.
4. 교차 아키텍처 전이: 대체 성공률과 특징 분할 (레지스터 대 번들) 이 다양한 행렬 재귀 아키텍처 (Gated DeltaNet, RWKV-7, Mamba-2) 에서 유효함을 검증했으나, 성공률은 기질의 "쓰기 랭크"에 따라 달라집니다.

실험 결과

대체 성능

• Qwen3.5-0.8B (Gated DeltaNet): L9 H4 에서 발화의 92.4% 에서 제거 실험보다 원자가 우수했습니다. 87 개 원자에 대한 인구 테스트는 89.8% 에서 유지되었습니다. 엄격한 순서 $KL_{atom} < KL_{ablate} < KL_{random}$이 발화의 89.5% 에서 성립합니다.
• Mamba-2-370M: Mamba-2 가 전체 랭크 1 외적이 아닌 대각선 상태 업데이트를 사용함에도 불구하고, 2,500 개의 발화에서 **88.1%**의 대체 성공률을 보였습니다.
• RWKV-7: 랭크 2 쓰기 규칙과 일치하는 중간 성능을 보입니다.

로그잇 시프트 예측

폐형 표현식은 측정된 효과를 정확하게 추적합니다. L9 H4 에서 예측된 로그잇 시프트와 측정된 로그잇 시프트 간의 중앙값 $R^2$는 0.98 입니다. 그러나 이 폐형은 Mamba-2 와 Qwen3.5-4B 에서 실패 (음수 $R^2$) 하며, 게이트 인자 $G$가 기질 특정적임을 나타냅니다.

행동 개입

• 삭제: 자연스러운 발화 위치에서 특정 레지스터 원자 (F412) 를 삭제하면 목표 토큰 ("space") 의 로그 확률이 중앙값 -0.116 nats만큼 감소했습니다 ($p=1.07 \times 10^{-6}$).
• 생성 조향: 세 개의 연속된 위치에서 폐형 방향을 지속적으로 설치하면, 탐욕적 디코딩 하에서 중간 순위 토큰의 "연속 내 목표" 비율이 **33.3% 에서 100%**로 증가했습니다.
• 지문 수준 증폭: 4B 모델에서 경계 상관 특징을 증폭하면 5 배 용량에서 줄바꿈 생성이 33% 감소했습니다 (400 토큰당 16.8 개에서 11.2 개로), 문서 구조에 대한 제어를 입증했습니다.

아키텍처 범위

본 논문은 "쓰기 랭크" 효과성의 위계를 확립합니다:

• Gated DeltaNet (랭크 1): 가장 높은 레지스터 코사인 분리 (0.262) 와 대체 충실도.
• RWKV-7 (랭크 2): 중간 분리 (0.180).
• Mamba-2 (대각선/스칼라): 가장 낮은 분리 (0.0575) 이지만, 여전히 약 88% 의 대체 성공률을 달성합니다.

중요성 및 주장

본 논문은 잔여 스트림 분석을 넘어 행렬 재귀 쓰기 사이트에 대한 최초의 인과적 사전을 제공한다고 주장합니다. 이는 다음을 입증합니다:

1. 쓰기 랭크가 결정적입니다: 상태를 해석 가능한 "레지스터"로 분해하는 능력은 네이티브 쓰기 규칙의 랭크 (랭크 1 > 랭크 2 > 대각선) 와 상관관계가 있습니다.
2. 인과적 대체가 가능합니다: 학습된 원자는 하류 KL 발산을 최소화하면서 네이티브 쓰기를 대체할 수 있어, 재귀 모델에서 "레지스터" 개념을 검증합니다.
3. 폐형 해석 가능성: Gated DeltaNet 의 경우, 특징이 로그잇에 미치는 영향을 파라미터 피팅 없이 분석적으로 예측할 수 있어, 사전과 모델의 게이트 및 읽기 메커니즘을 직접 연결합니다.

저자들은 겸손하게 한계를 지적합니다: 폐형 계수는 기질 간에 전이되지 않으며 (예: Mamba-2 에서 실패), 원자별 정체성은 시드 특정적이지만 (클래스 수준 속성은 안정적임). 이 연구는 상태 공간 모델의 기계적 해석 가능성에 대한 새로운 벤치마크를 수립하며, "쓰기" 원시 연산이 사전 학습과 인과적 편집을 위한 실현 가능한 대상임을 보여줍니다.