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

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

왜 중요한가

핵심 기여

State-Prediction Separation 구조 제안

입력 토큰과 예측 토큰을 교차하는 두 개의 스트림으로 모델을 재구성하여 예측 전용의 <predict> 슬롯과 영구 상태용 입력 슬롯을 분리했다. <predict> 슬롯은 제한된 슬라이딩 윈도우 내에서만 KV 항목을 유지하고 그 밖의 항목은 폐기하여 예측 역할이 장기 상태로 전파되는 것을 막았다. 이러한 구조적 변경은 학습 중 gradient가 예측과 상태 준비로 분산되는 방식을 재배치하여 두 역할 간의 경쟁을 완화시켰다.

대규모 사전학습 실험으로 일관된 성능 개선 확인

53M에서 1.678B까지 다섯 가지 스케일에서 SPS를 표준 Transformer, 2x Memory, Delayed State와 비교하여 검증 NLL과 제로샷 정확도에서 일관된 개선을 보였다. 모든 스케일에서 SPS가 동일한 영구 KV 캐시 크기를 유지하면서 검증 NLL을 낮췄고 제로샷 정확도는 대략 2.3%에서 3.1%까지 향상되었다. 특히 1.6B급 스케일에서는 표준 모델이 필요로 하는 토큰의 약 절반 수준으로 동일한 품질을 달성하는 데이터 효율을 보였다.

gradient 흐름과 persistent state 의 중요성 실험적 분석

각 위치에서 현재 손실과 미래 손실이 파라미터에 미치는 기여도를 수치화한 후 스트림별로 비교하여 SPS가 미래 손실 그라디언트를 입력 스트림으로 더 많이 라우팅함을 확인했다. 슬라이딩 윈도우로 영구 상태를 제한했을 때 SPS의 성능 손실이 Delayed State보다 훨씬 크게 나타나 SPS의 입력 스트림이 실제로 추론 시 중요한 persistent state를 제공함이 입증되었다. 이 분석은 단순한 계산 증가나 메모리 확장만으로는 얻을 수 없는 구조적 이득이 존재함을 보였다.

효율성 및 견고성 검증

학습 시 오버헤드는 존재하나 추론 시에는 persistent KV 캐시 크기가 표준과 거의 동일하여 peak memory 비율이 1.01 수준에 머물렀다. end-to-end throughput은 표준 대비 약 6%에서 10% 내외의 차이에 그쳐 실서비스 관점에서 수용 가능한 오버헤드로 평가되었다. 윈도우 크기 w 에 대한 민감도 실험에서 w=64가 경험적으로 우수했으며 Reverse SPS 대비 작은 윈도우에서 더 견고한 성능을 보였다.

핵심 아이디어 이해하기

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 연산과 전체 학습 비용이 증가하며 낮은 오버헤드로 동일한 분리를 구현할 수 있는지 여부는 미해결 문제로 남아 있다.

실무 활용

SPS는 학습 단계에서 예측과 상태 준비를 구조적으로 분리함으로써 토큰당 학습 효율을 높여 사전학습 예산을 절감할 수 있다. 추론 단계에서는 persistent KV 캐시 크기를 표준 수준으로 유지하므로 메모리 비용 증가가 거의 없고 실서비스 적용 부담이 작다. 공개된 구현을 통해 연구자와 엔지니어가 제시된 attention 패턴을 기존 Transformer 코드베이스에 통합할 수 있다.

대규모 사전학습 예산이 제한된 연구실이나 기업에서 동일한 품질을 더 적은 토큰으로 달성하려는 목적에 적용 가능하다. 실제로 논문에서는 1.6B급 모델에서 표준이 필요로 하는 토큰의 약 절반 수준으로 유사한 검증 손실을 달성한 사례가 보고되었다. 이로 인해 고품질 데이터가 부족한 상황에서 학습 러닝레이트를 더 효율적으로 분배할 수 있다.
추론 메모리 예산을 거의 유지하면서 모델의 데이터 효율을 개선해야 하는 실무 환경에 적합하다. SPS의 persistent KV 캐시는 표준 Transformer와 유사하므로 기존 추론 인프라와의 호환성이 높다. 이를 통해 레거시 배포 환경에서 큰 하드웨어 변경 없이 성능 이득을 실현할 수 있다.
구조적 연구 방향으로서 예측 대상과 상태 표현을 분리하는 다른 손실 함수나 파라미터 분리 기법과 병행하여 탐색하기에 적합하다. 본 논문은 스트림 분리만으로도 성능을 개선했으며 추가로 파라미터 분리나 보조 손실을 도입하면 더 큰 이득이 가능할 여지가 있다. 해당 조합 연구는 후속 연구로 실용적 가치가 있다.

코드 공개 여부: 공개

코드 저장소 보기

키워드

SPS(State-Prediction Separation)(상태-예측 분리)Transformer(Transformer)KV cache(키-값 캐시)sliding window(슬라이딩 윈도우)validation NLL(검증 NLL)

용어 해설

키-값 캐시(KV cache): — Transformer가 과거 토큰의 key와 value를 저장하여 이후 토큰의 attention 연산에서 재사용하는 메커니즘이다. 이 캐시는 생성 과정에서 과거의 표현을 빠르게 참조하도록 하여 문맥 유지와 속도 향상에 기여한다. 본문에서는 어떤 스트림의 KV 항목을 지속적으로 보존하느냐가 성능과 데이터 효율성에 결정적 영향을 미치는 요소로 다뤄진다.
슬라이딩 윈도우(Sliding window): — 모델이 최근 몇 개의 예측용 토큰만 일시적으로 참조하도록 키-값 저장을 제한하는 기법이다. 본문에서는 <predict> 항목을 윈도우 크기 w 이내에서만 참조하도록 하여 예측 스트림을 비지속적(ephemeral)으로 유지한다. 이 방식은 persistent state 크기를 일정 수준으로 유지하면서 지역적 일관성을 보장하는 수단이다.
<predict> 토큰(<predict> token): — 입력 토큰마다 삽입되는 학습 가능한 더미 토큰으로서 바로 다음 토큰의 예측만 담당하도록 설계된 슬롯이다. 이 토큰의 활성화는 짧은 윈도우 동안만 KV 캐시에 남고 이후 폐기되어 예측 역할과 상태 저장 역할을 분리하는 핵심 수단이 된다. 논문에서는 이 토큰을 통해 예측 스트림과 상태 스트림을 구분하여 gradient 흐름을 분리한다.
미래 손실 그라디언트(Future-loss gradient): — 어떤 위치의 파라미터가 이후 시점 j>i 의 손실 ℓ_j 에 기여할 때 생기는 기울기 성분을 지칭한다. 본문은 이 성분이 어느 스트림으로 얼마나 배분되는지가 모델의 상태 표현과 예측 성능에 직접적인 영향을 준다고 분석한다. SPS는 이 미래 손실 그라디언트를 입력 스트림에 더 많이 모으는 것으로 확인됐다.

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

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

한계점

실무 활용

키워드

용어 해설

관련 토론

댓글

관련 기사

SMELT: 반복형 MoE Transformer의 계산량 대비 확장 법칙

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

한계점

실무 활용

키워드

관련 토론

댓글

관련 기사

SMELT: 반복형 MoE Transformer의 계산량 대비 확장 법칙