다양성 vs 정확성: 넥스트 토큰 예측의 심층 분석

Diversity or Precision? A Deep Dive into Next Token Prediction

December 28, 2025

저자: Haoyuan Wu, Hai Wang, Jiajia Wu, Jinxiang Ou, Keyao Wang, Weile Chen, Zihao Zheng, Bei Yu

cs.AI

초록

최근 연구에 따르면 강화 학습(RL)은 대규모 언어 모델(LLM)의 추론 능력을 크게 향상시킬 수 있음이 입증되었습니다. 그러나 이러한 RL 훈련의 효과는 사전 훈련된 모델의 토큰 출력 분포로 정의된 탐색 공간에 크게 의존합니다. 본 논문에서는 표준 교차 엔트로피 손실을 단일 단계 에피소드 내에서 적용된 정책 경사 최적화의 특수한 사례로 해석하여 재조명합니다. 사전 훈련된 분포가 이후 RL을 위한 탐색 가능성을 어떻게 형성하는지 체계적으로 연구하기 위해, 온-정책 RL 원리를 지도 학습에 적용하는 일반화된 사전 훈련 목표를 제안합니다. 다음 토큰 예측을 확률적 결정 과정으로 구성하고, 다양성과 정밀도를 명시적으로 균형 잡는 보상 형성 전략을 도입합니다. 우리의 방법은 실제 정답 토큰에 대한 확률 집중을 제어하는 양의 보상 스케일링 인자와 상위 순위 및 하위 순위 부정적 토큰을 비대칭적으로 처리하는 순위 인식 메커니즘을 사용합니다. 이를 통해 사전 훈련된 토큰 출력 분포를 재구성하고, 어떻게 하면 RL에 더 유리한 탐색 공간을 제공하여 궁극적으로 종단간 추론 성능을 향상시킬 수 있는지 조사합니다. 높은 분포 엔트로피가 효과적인 탐색을 용이하게 한다는 직관과는 반대로, 정밀도 지향 사전 분포를 부과하는 것이 RL에 더 우수한 탐색 공간을 제공한다는 사실을 발견했습니다.

English

Recent advancements have shown that reinforcement learning (RL) can substantially improve the reasoning abilities of large language models (LLMs). The effectiveness of such RL training, however, depends critically on the exploration space defined by the pre-trained model's token-output distribution. In this paper, we revisit the standard cross-entropy loss, interpreting it as a specific instance of policy gradient optimization applied within a single-step episode. To systematically study how the pre-trained distribution shapes the exploration potential for subsequent RL, we propose a generalized pre-training objective that adapts on-policy RL principles to supervised learning. By framing next-token prediction as a stochastic decision process, we introduce a reward-shaping strategy that explicitly balances diversity and precision. Our method employs a positive reward scaling factor to control probability concentration on ground-truth tokens and a rank-aware mechanism that treats high-ranking and low-ranking negative tokens asymmetrically. This allows us to reshape the pre-trained token-output distribution and investigate how to provide a more favorable exploration space for RL, ultimately enhancing end-to-end reasoning performance. Contrary to the intuition that higher distribution entropy facilitates effective exploration, we find that imposing a precision-oriented prior yields a superior exploration space for RL.