多様性か精度か？次トークン予測の深層分析

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.