Diversity or Precision? A Deep Dive into Next Token Prediction

TL;DR

This paper investigates how the pre-trained token-output distribution shapes exploration in reinforcement learning for large language models. By reinterpreting cross-entropy as a one-step policy-gradient objective and introducing a generalized reward-shaped pre-training objective with a global entropy regulator $\beta$ and rank-aware negative shaping via $(\tilde{\lambda},\hat{\lambda},k)$, the authors systematically study the diversity-precision trade-off. Across dense and Mixture-of-Experts models and multiple training stages, they find that precision-oriented priors—achieved by lowering global entropy or suppressing tail negatives—consistently enhance downstream RL performance on reasoning and math tasks, challenging the notion that higher entropy always benefits exploration. The results highlight the importance of reward design in pre-training for long-chain-of-thought reasoning, suggesting practical strategies for improving end-to-end RL outcomes without sacrificing perplexity.

Abstract

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.

Figures (8)

• Figure 1: Changes of PPL and entropy during pre-training across 1B and 4B dense models, developed based on different configurations.
• Figure 2: Changes of PPL and entropy during pre-training across 5B-A0.3B and 10B-A0.5B MoE models, developed based on different configurations.
• Figure 3: Changes of performance during pre-training across models with various model parameters, developed based on dense and MoE architectures under different configurations.
• Figure 4: Changes of performance during mid-training across 4B dense and 10B-A0.5B MoE models, developed based on different configurations.
• Figure 5: Changes of performance during RL training across various actor models, developed based on a 4B dense architecture under different configurations.