Diversity or Precision? A Deep Dive into Next Token Prediction

Abstract: Recent advancements have shown that reinforcement learning (RL) can substantially improve the reasoning abilities of 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.

• The paper introduces a reward-shaping method that modulates the precision-diversity trade-off in next-token prediction.
• It recasts next-token prediction as a policy optimization problem, showing that a precision-biased pre-training phase improves RL initialization.
• Empirical results across architectures reveal that low-entropy pre-training leads to enhanced reasoning and stable RL performance.

Diversity or Precision? A Deep Dive into Next Token Prediction

Introduction

This paper interrogates the core assumption in LLM training that diversity, as instantiated by high entropy in next-token distributions, necessarily enables superior downstream reasoning via reinforcement learning (RL). The authors formalize the connection between cross-entropy-based next-token prediction and single-step on-policy RL, and introduce a reward-shaping approach that allows explicit manipulation of the precision-versus-diversity trade-off during pre-training. Their empirical results across architectures and scales indicate that imposing a precision-oriented prior in pre-training yields a more effective RL initialization and improved performance on complex reasoning tasks, challenging conventionally held beliefs in LLM training.

Theoretical Framework

The authors recast next-token prediction as a stochastic policy optimization problem, making explicit the equivalence between standard cross-entropy loss and a single-step policy gradient update with a specific reward structure. Standard cross-entropy is viewed as assigning a maximal reward to the ground-truth token and penalizing all negatives only implicitly via the softmax constraint. This perspective motivates a systematic intervention in reward assignment to modulate the model’s pre-trained output distribution beyond entropic regularization.

Their generalized reward structure consists of a scaled positive reward (parameterized by $\beta$) for the true token and differentiated negative token shaping ($\tilde{\lambda}$ and $\hat{\lambda}$ for high- and low-probability negatives, respectively). This enables both global (distributional entropy) and local (probability mass allocation within the head/tail of the output distribution) entropy control.

Experimental Design

The proposed approach is instantiated in both dense and Mixture-of-Experts (MoE) architectures (with 1B–10B parameter scales), evaluated at three stages: base pre-training, mid-training with reasoning-enhanced data, and RL fine-tuning focused on mathematical tasks. A diverse multi-domain evaluation suite (MMLU, ARC, GSM8K, HumanEval+ etc.) is used to decouple knowledge and reasoning capacity. Reward parameters are swept (notably $\beta = -0.25$ vs. $\beta = 0.5$) to compare low-precision/high-precision regimes.

Analysis of Pre-Training Dynamics

Altering the positive reward scaling factor $\beta$ exerts consistent control over global entropy during pre-training, as evidenced by stable perplexity but divergent entropy trends in both dense and MoE models.

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.

A negative $\beta$ setting ($\beta < 0$) sharpens the distribution, concentrating probability mass on the ground-truth, thus reducing entropy; conversely, positive $\beta$ produces flatter, higher-entropy policies. The explicit differentiation of negative reinforcement via $\tilde{\lambda}$ and $\hat{\lambda}$ further modulates local entropy within the output space, providing increased flexibility over token-level policy shaping.

Structural scaling properties are preserved (larger models outperform at all stages), but the gains for precision-oriented configurations become increasingly prominent at scale.

Figure 3: Changes of performance during pre-training across models with various model parameters, developed based on dense and MoE architectures under different configurations.

Mid-Training and Reasoning Probes

Incorporating reasoning-focused data in mid-training reveals a persistent and even growing advantage for the precision-oriented (low-entropy) initialization. Negative $\beta$ settings drive higher performance ceilings in both knowledge and reasoning benchmarks, and aggressive local negative shaping never impedes knowledge acquisition, but produces consistent improvements in logic and mathematics.

Figure 4: Changes of performance during mid-training across 4B dense and 10B-A0.5B MoE models, developed based on different configurations.

Reinforcement Learning Outcomes

The most pronounced effects manifest during RL finetuning—where the initialization dictated by pre-training diversity-precision balance critically determines the accessible reasoning/exploration space. Counterintuitively, initializing RL from a low-entropy, precision-oriented base (i.e., with mass concentrated on the ground-truth) leads to stronger and more stable RL learning curves across dense and MoE models.

Figure 5: Changes of performance during RL training across various actor models, developed based on a 4B dense architecture under different configurations.

Figure 6: Changes of performance during RL training across various actor models, developed based on a 10B-A0.5B MoE architecture under different configurations.

Analysis of policy entropy and sequence length during RL training indicates that high-entropy (diversity-oriented) pre-training causes a rapid entropy collapse and suppressed sequence generation at RL onset. In contrast, the precision-oriented regime yields more robust policy distributions, less prone to mode collapse, and effect smoother activation of long-chain-of-thought capabilities.

Figure 7: Changes of entropy and response length during RL training across various actor models, developed based on 4B dense and 10B-A0.5B MoE architectures under different configurations.

Diversity/Precision and Pass@k

The Pass@$k$ analysis on math and code tasks demonstrates that maximizing diversity does not translate into maximized solution coverage—rather, precision-optimized models maintain the requisite output variability without sacrificing accuracy, refuting the hypothesis that high entropy is optimal for enumerative search.

Figure 8: $\text{Pass@}k$ curve of base models on mathematics reasoning and code generation tasks, developed based on 4B dense and 10B-A0.5B MoE models under different configurations.

The data collectively indicate that, for high-level reasoning, exploration space should be tightly focused via high-probability mass on anticipated correct paths, not diffused among low-probability alternatives.

Implications and Future Directions

This work challenges the standard prescription of maximizing entropy/diversity in LLM pre-training, especially when downstream RL is expected to enable abstract or compositional reasoning. It suggests that RL’s effectiveness depends on a pre-training regime that pre-screens reasoning trajectories via a precision-biased output distribution.

Practically, reward shaping for next-token prediction can serve as a tuning axis for LLM developers to tailor models according to intended downstream RL workloads (e.g., scientific/mathematical reasoning vs. creative generation). Theoretically, this framework clarifies that pre-training and RL are not separable optimization processes—the choice of loss/reward at the pre-training stage is a critical determinant of which RL exploration strategies are feasible and effective.

Future research can extend these findings to settings with hierarchical or multi-stage reasoning, latent-variable architectures, and adaptive computation policies, further investigating how the structure of the pre-training policy interacts with more complex, multi-step reward signals.

Conclusion

This study provides a rigorous, RL-theoretic justification for precision-oriented pre-training objectives in LLMs and demonstrates with strong empirical evidence that such configurations yield superior reasoning performance following RL finetuning. These findings establish that pre-training reward shape, via explicit diversity-precision control, is a critical axis of LLM performance scaling, with direct theoretical and practical consequences for state-of-the-art LLM development and deployment (2512.22955).