The State-Prediction Separation Hypothesis

Transformers use the same forward computation stream to both predict the next token and store useful state for future token predictions. We formulate the \emph{state-prediction separation hypothesis}: disentangling the two roles yields better language modeling performance. We design a Transformer variant that uses two computation streams to separate the two functions, and conduct pretraining experiments across various scales. Our experiments show that state-prediction separation consistently offers better data and compute efficiencies, improving validation loss and outperforming standard Transformers by 2--3 percentage points on average on downstream tasks. We also conduct extensive empirical analysis that rules out potential confounders and demonstrates the fundamental difference in the gradients our design entails.

한 줄 요약

Transformer에서 상태 저장과 예측을 분리한 SPS 설계가 언어 모델 성능을 2~3% 개선한다.

핵심 기여도

**SPS 설계**: 상태와 예측을 분리한 두 개의 계산 스트림을 사용한 Transformer 변형 모델 제안.
**성능 개선**: FineWeb-Edu 데이터셋에서 검증 손실 감소 및 다운스트림 태스크에서 2~3% 포인트 개선.
**자원 효율성**: 동일 데이터/계산량에서 더 높은 성능 달성.
**기본 메커니즘 분석**: 경사도 흐름과 상태 분석을 통해 분리가 학습 효율성에 기여함을 입증.

핵심 아이디어

기존 Transformer는 동일한 계산 스트림을 사용하여 다음 토큰 예측과 상태 저장을 동시에 수행한다. 그러나 이 연구는 **state-prediction separation hypothesis**를 제시하며, 두 기능을 분리하면 언어 모델 성능이 향상된다고 주장한다. 이에 따라, **SPS Transformer**는 두 개의 계산 단계를 도입하여 예측과 상태 저장을 분리한다. 첫 번째 단계는 상태를 저장하는 역할을 하며, 두 번째 단계는 토큰 예측을 수행한다. 이 때, 첫 번째 단계의 key-value(KV) 캐시는 유지되지만, 두 번째 단계의 KV는 버려진다. 이는 **state stream**과 **prediction stream**을 구분하는 핵심 아이디어이다.

기술적 접근법

**SPS Transformer**: 두 개의 계산 단계를 사용하여 상태와 예측을 분리.
**KV 캐시 관리**: 첫 번째 단계의 KV는 캐시에 저장, 두 번째 단계의 KV는 버림.
**실험 규모**: 53M에서 1.678B 파라미터까지 다양한 규모의 모델로 사전 학습 실험 수행.
**하드웨어**: NVIDIA H100 GPU 기반으로 추론 효율성 평가 (토큰/초, 메모리 사용량 측정).
**비교 대상**: 기존 Transformer와 동일한 KV 캐시 크기, 동일한 생성 루프 조건에서 비교.

주요 결과

**FineWeb-Edu**: 검증 손실 감소.
**다운스트림 태스크**: 2~3% 포인트 개선.
**WikiText, C4, Pile-Books3, GovReport**: 평균적으로 out-of-distribution 데이터셋에서 NLL 개선.
**ARC-Easy, HellaSwag, PIQA, SciQ, LAMBADA**: 제로샷 정확도 평균적으로 향상.
**추론 효율성**: 동일한 GPU 메모리 사용량에서 SPS 모델의 처리 속도는 기존 Transformer 대비 약 1.6배 빠름.

의의 및 한계

SPS 설계는 기존 Transformer의 한계를 극복하며, 데이터와 계산 효율성을 동시에 향상시킨다. 특히, 고질 데이터가 한정된 상황에서 **각 토큰당 학습 효과 증대**가 중요하다. 그러나 본 연구는 단일 데이터셋(FineWeb-Edu)에서 사전 학습되었으며, 다른 데이터셋 혼합 효과는 검증되지 않았다. 또한, 최대 1.678B 파라미터 규모까지만 실험되었으며, 더 큰 모델에서의 성능 향상 여부는 추가 연구가 필요하다. SPS는 예측 스트림에 추가 계산 단계를 도입하여 **약 2배의 훈련 계산량 증가**를 유발하며, 이는 효율성 개선에 대한 균형 고려가 필요하다.

실용적 활용

SPS는 대규모 언어 모델의 사전 학습 효율성을 높이는 설계로, **데이터 및 컴퓨팅 자원이 제한된 연구 환경**에서 유용하다. 또한, **고성능 GPU 기반의 추론 시스템**에서 메모리 사용량을 줄이면서도 성능을 유지할 수 있어, 산업적 언어 모델 배포에도 적용 가능하다.