NITP: LLM 사전 훈련을 위한 다음 암시적 토큰 예측

초록

표준 다음 토큰 예측(Next-Token Prediction, NTP)은 언어 모델을 출력 로짓 공간의 이산 레이블만으로 지도 학습한다. 본 연구에서는 이러한 희소 원-핫 지도 학습이 잠재 표현 공간을 충분히 제약하지 못하여, 은닉 상태가 일반화를 제한할 수 있는 퇴화되거나 이방성인 구성으로 표류할 수 있다고 주장한다. 이 문제를 해결하기 위해, 우리는 이산 예측을 표현 공간에서의 밀집 연속 지도 학습으로 보완하는 NITP(Next Implicit Token Prediction)를 제안한다. NITP는 동일 모델의 얕은 계층 표현을 안정적인 자기지도 학습 목표로 사용하여, 다음 토큰의 암묵적 의미 내용을 예측하도록 모델을 학습시킨다. 이론적 분석을 통해 NITP가 제약 부족 자유도를 완화하고 컴팩트하고 구조화된 표현 기하를 장려함으로써 최적화 지형을 정규화함을 보여준다. 실험적으로, 0.5B에서 9B 파라미터 범위의 밀집 모델 및 MoE 모델 전반에 걸쳐 NITP는 무시할 수준의 계산 오버헤드로 하위 작업 성능을 일관되게 향상시킨다. 9B MoE 모델에서 NITP는 MMLU-Pro에서 5.7%의 절대적 개선, C3에서 6.4%, CommonsenseQA에서 4.3%의 향상을 달성하며, 추가 학습 FLOPs는 약 2%에 불과하고 추론 비용은 증가하지 않는다. 구현 코드는 https://github.com/aHapBean/NITP에서 확인할 수 있다.

English

Standard next-token prediction (NTP) supervises language models solely through discrete labels in the output logit space. We argue that this sparse one-hot supervision leaves the latent representation space under-constrained, allowing hidden states to drift into degenerate and anisotropic configurations that can limit generalization. To address this issue, we propose Next Implicit Token Prediction (NITP), which augments discrete prediction with dense continuous supervision directly in the representation space. NITP trains the model to predict the implicit semantic content of the next token, using shallow-layer representations from the same model as stable self-supervised targets. We provide theoretical analysis showing that NITP regularizes the optimization landscape by mitigating under-constrained degrees of freedom and encouraging a compact, structured representation geometry. Empirically, across dense and MoE models ranging from 0.5B to 9B parameters, NITP consistently improves downstream performance with negligible computational overhead. On a 9B MoE model, NITP achieves a 5.7% absolute improvement on MMLU-Pro, along with gains of 6.4% on C3 and 4.3% on CommonsenseQA, with approximately 2% additional training FLOPs and no additional inference cost. Our implementation is available at https://github.com/aHapBean/NITP.