Diversiteit of Precisie? Een Diepgaande Duik in Voorspelling van de Volgende Token
Diversity or Precision? A Deep Dive into Next Token Prediction
December 28, 2025
Auteurs: Haoyuan Wu, Hai Wang, Jiajia Wu, Jinxiang Ou, Keyao Wang, Weile Chen, Zihao Zheng, Bei Yu
cs.AI
Samenvatting
Recente ontwikkelingen hebben aangetoond dat reinforcement learning (RL) de redeneervaardigheden van grote taalmodellen (LLM's) aanzienlijk kan verbeteren. De effectiviteit van dergelijke RL-training hangt echter cruciaal af van de verkenningsruimte die wordt gedefinieerd door de token-uitvoerdistributie van het vooraf getrainde model. In dit artikel herzien we de standaard cross-entropy loss en interpreteren we deze als een specifiek geval van policy gradient optimalisatie toegepast binnen een eenstaps-episode. Om systematisch te bestuderen hoe de voorgetrainde distributie het verkenningspotentieel voor daaropvolgende RL vormt, stellen we een gegeneraliseerd vooraf-trainingsdoel voor dat on-policy RL-principes toepast op supervised learning. Door next-token voorspelling te formuleren als een stochastisch beslissingsproces, introduceren we een reward-shaping strategie die expliciet diversiteit en precisie in evenwicht brengt. Onze methode gebruikt een positieve schaalfactor voor beloningen om de waarschijnlijkheidsconcentratie op ground-truth tokens te controleren en een rangbewust mechanisme dat hoge- en lage-rang negatieve tokens asymmetrisch behandelt. Hierdoor kunnen we de voorgetrainde token-uitvoerdistributie hervormen en onderzoeken hoe we een gunstigere verkenningsruimte voor RL kunnen bieden, wat uiteindelijk de end-to-end redeneerprestaties verbetert. In tegenstelling tot de intuïtie dat een hogere distributie-entropie effectieve verkenning vergemakkelijkt, constateren we dat het opleggen van een op precisie gerichte prior een superieure verkenningsruimte voor RL oplevert.
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.