상태-예측 분리 가설과 State-Prediction Separation Transformer

Transformer 계열 모델은 각 시점의 활성화가 즉시 다음 토큰 예측과 이후 토큰을 위한 상태 보관이라는 두 역할을 동시에 수행하면서 역할이 충돌할 가능성이 존재한다. 이 논문은 예측과 상태 준비를 구조적으로 분리하면 동일한 파라미터 수에서 더 낮은 검증 NLL과 더 나은 제로샷 성능을 얻는다는 점을 실험으로 입증했다. 데이터가 희소해지는 상황에서 토큰당 학습 효율을 올리는 방식은 사전학습 비용과 데이터 예산을 보다 효율적으로 활용하게 한다.

Transformer의 각 위치(hidden state)는 두 가지 목적을 동시에 수행한다. 하나는 즉시 다음 토큰을 예측하기 위한 표현을 만들고 다른 하나는 이후 위치들이 참조할 수 있는 키-값 항목을 생성하여 장기 상태를 형성하는 것이다. 이 두 목적이 동일한 활성화와 동일한 파라미터 경로로 최적화될 때 서로 경쟁하여 표현이 타협 상태에 놓일 수 있다.

해결 원리는 예측과 상태 준비를 서로 다른 시간 스텝과 스트림으로 분리하는 것이다. 구체적으로 각 입력 토큰 x_i 뒤에 토큰 ρ_i를 삽입하여 시퀀스를 확장하고 입력 스트림은 영구적인 KV 항목만을 축적하도록 유지하는 반면 예측 스트림의 KV 항목은 슬라이딩 윈도우 w 내에서만 참조되도록 한다. 이렇게 하면 모델의 gradient가 예측과 미래 손실로부터 오는 신호를 서로 다른 위치로 라우팅하게 되고 입력 스트림은 장기 상태 준비에, 예측 스트림은 즉시 예측에 특화될 수 있다.

이 분리는 실제 학습과 추론에서 실질적 변화를 만든다. SPS는 동일한 파라미터 수에서 검증 손실을 더 빠르게 낮추며 같은 수준의 성능을 달성하기 위해 필요한 토큰 수가 표준 Transformer보다 적게 나타났다. 또한 입력 스트림에 쌓이는 persistent keys가 추론 시 더 큰 기여를 함이 제한 실험에서 확인되어 분리된 스트림이 단순한 계산 증가가 아닌 구조적 이득을 제공함이 확증됐다.

전체 접근 방식은 각 입력 토큰 뒤에 예측 전용 슬롯을 삽입하여 두 개의 상호작용하는 표현 스트림을 만드는 것이다. 원문에서는 확장된 시퀀스를 x_1, ρ_1, x_2, ρ_2, … 형태로 구성하고 입력 스트림은 모든 인과적 입력을 참조하도록 하며 예측 스트림은 최근 w개의 항목만 참조하도록 attention 마스크를 설계했다. 이렇게 정의된 attention 패턴은 persistent KV 캐시에는 오직 입력 항목만 남기고 예측 항목은 윈도우를 벗어나면 폐기되도록 만든다.

핵심 메커니즘은 attention 마스크와 슬라이딩 윈도우를 통해 gradient의 소스와 도착지를 분리하는 것이다. 각 쿼리 q가 있을 때 SPS의 attention 집합은 모든 인과적 입력과 최근의 항목의 합집합으로 정의되어 예측 위치에서는 해당 입력 위치를 포함한 추가 항목을 볼 수 있게 한다. 이로 인해 입력 위치의 표현은 이후 손실들로부터 누적되는 미래 손실 그라디언트를 주로 받는 반면 예측 위치의 표현은 현재 손실에 의해 주로 최적화된다.

학습 및 구현 세부에서는 FineWeb-Edu 데이터를 사용하여 다양한 스케일(53M~1.678B)에서 20B 토큰 기본 예산으로 사전학습을 수행했다. 학습은 nanoGPT 기반 하이퍼파라미터와 AdamW 옵티마이저를 사용했고 FlashAttention 트리톤 커널을 해당 attention 패턴에 맞춰 수정하여 슬라이딩 윈도우를 효율적으로 시뮬레이션했다. 윈도우 크기 w는 경험적으로 64를 모든 스케일에서 채택하였고 일부 실험에서 Reverse SPS와 Delayed State 등 대비군을 동일한 설정으로 비교하였다.

SPS는 모든 스케일에서 FineWeb-Edu 검증 NLL을 표준 Transformer보다 낮게 기록했다. 개선 폭은 스케일이 커질수록 확대되는 경향을 보였고 예컨대 XS에서 약 −0.042, XL에서 약 −0.068의 NLL 개선이 관찰되었다. 동일한 GPU시간 대비 학습 곡선에서도 SPS가 표준을 앞서는 경향을 보였고 표준 모델이 토큰을 두 배로 보았을 때에도 SPS가 더 낮은 손실을 유지했다.

일반화 실험에서 SPS는 네 개의 분포 외 코퍼스에 대한 평균 NLL을 약 0.09~0.11만큼 낮추었고 다섯 개 제로샷 벤치마크의 평균 정확도는 약 2.3%에서 3.1%까지 향상되었다. 이러한 품질 향상은 데이터 효율로 직결되어 같은 성능을 달성하는 데 필요한 토큰 수가 감소했다는 점에서 실용적 의미를 가졌다. 윈도우 크기 민감도 실험에서는 w=64가 소폭 우수했고 Reverse SPS는 작은 w에서 성능 저하가 급격히 발생했으나 SPS는 더 견고했다.

효율성 측면에서 SPS의 persistent KV 캐시 크기는 표준과 거의 동일하여 추론 시 peak memory 비율이 약 1.01 수준에 머물렀다. end-to-end throughput은 표준 대비 대략 6%에서 10% 이내의 저하로 비교적 작은 오버헤드만 존재했다. 학습 시에는 각 입력에 대해 추가 전진 단계가 들어가므로 per-step 연산이 증가하지만 이 비용 대비 데이터 효율과 최종 성능 개선이 전체적인 이득으로 평가되었다.

전체 아키텍처는 기존의 pre-normalized Transformer 블록을 유지하되 각 입력 토큰 뒤에 토큰을 삽입하여 시퀀스 길이를 두 배로 확장한다. 두 스트림은 동일한 파라미터를 공유하는 표현 공간을 사용하나 attention 마스크로 가시성 규칙을 달리하여 입력 스트림은 모든 인과적 입력을 참조하고 예측 스트림은 슬라이딩 윈도우 w 내의 항목만 참조하도록 구현되었다. persistent KV 캐시는 입력 항목만 보존하며 예측 항목은 윈도우를 벗어나면 폐기된다.

수학적으로 SPS는 각 위치 i 의 gradient를 예측 관련 항목과 상태 준비 관련 항목으로 분해한 기존 식을 확장하여 처리한다. 논문은 각 스트림 p 에 대해 r(p,k)=||∇{θ_p} ℓ{i+k}||2 / ||∇{θ_p} ℓ_i||_2 를 계산하여 미래 손실 기여도의 상대적 크기를 측정했다. 이 지표는 SPS의 입력 스트림이 높은 r(p,k) 값을 유지함으로써 미래 손실 신호를 더 많이 흡수한다는 사실을 정량적으로 보여준다.

구현 및 학습 세부에서는 GPT-2 토크나이저와 FineWeb-Edu 코퍼스를 사용하였고 학습 파라미터는 AdamW, β_1=0.9, β_2=0.95, peak LR=6e-4, 글로벌 배치 크기 96 시퀀스 등을 채택했다. attention 연산은 Triton 기반 FlashAttention을 슬라이딩 윈도우 패턴과 마스크에 맞추어 수정하여 처리량을 확보했다. 모든 실험은 동일한 시드와 데이터 순서를 사용하였고 몇몇 핵심 비교에서는 다중 시드 검증을 수행하여 통계적 유의성을 확인했다.

논문에서 제시된 실험은 단일 코퍼스인 FineWeb-Edu로 사전학습을 수행했기 때문에 다른 데이터 혼합이나 완전히 다른 도메인에서 동일한 이득이 보장되지는 않는다. 최대 규모 실험은 1.678B 파라미터로 제한되어 있으며 논문은 SPS-표준 간격이 더 큰 모델 스케일에서도 계속 확장될 가능성을 제시하지만 이를 검증하려면 더 큰 스케일 실험이 필요하다고 명시했다. 학습 시에는 각 입력에 대해 추가 전진 단계가 들어가므로 per-step 연산과 전체 학습 비용이 증가하며 낮은 오버헤드로 동일한 분리를 구현할 수 있는지 여부는 미해결 문제로 남아 있다.