Diversity or Precision? A Deep Dive into Next Token Prediction
作者: Haoyuan Wu, Hai Wang, Jiajia Wu, Jinxiang Ou, Keyao Wang, Weile Chen, Zihao Zheng, Bei Yu
分类: cs.CL
发布日期: 2025-12-28 (更新: 2026-02-05)
💡 一句话要点
提出基于策略梯度的预训练目标，提升LLM在强化学习中的探索能力
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 大型语言模型 强化学习 预训练 策略梯度 奖励塑造 探索空间 token预测
📋 核心要点
- 现有方法依赖预训练模型的token输出分布进行强化学习探索，但缺乏对该分布如何影响探索空间的系统研究。
- 论文提出一种广义预训练目标，通过奖励塑造策略平衡多样性和精确性，优化预训练分布，从而改善强化学习的探索空间。
- 实验表明，与直觉相反，面向精确性的预训练先验能够为强化学习提供更优越的探索空间，提升端到端推理性能。
📝 摘要（中文）
近期的研究表明，强化学习(RL)能够显著提升大型语言模型(LLM)的推理能力。然而，这种RL训练的有效性关键取决于预训练模型token输出分布所定义的探索空间。本文重新审视了标准的交叉熵损失，将其解释为单步episode中策略梯度优化的一种特定实例。为了系统地研究预训练分布如何塑造后续RL的探索潜力，我们提出了一种广义的预训练目标，该目标将on-policy RL原则应用于监督学习。通过将next-token预测构建为随机决策过程，我们引入了一种奖励塑造策略，该策略显式地平衡了多样性和精确性。我们的方法采用正奖励缩放因子来控制ground-truth token上的概率集中，并采用一种rank-aware机制，不对称地处理高排名和低排名的负token。这使我们能够重塑预训练的token输出分布，并研究如何为RL提供更有利的探索空间，最终提高端到端的推理性能。与更高的分布熵有利于有效探索的直觉相反，我们发现施加面向精确性的先验可以为RL产生更好的探索空间。
🔬 方法详解
问题定义：现有的大型语言模型在进行强化学习微调时，其探索空间很大程度上依赖于预训练模型的token输出分布。如何有效地利用和优化这个预训练分布，使其能够为后续的强化学习提供更有利的探索空间，是当前面临的一个重要问题。传统的交叉熵损失训练方式可能无法提供最佳的探索空间，导致强化学习训练效率低下。
核心思路：论文的核心思路是将next-token预测任务视为一个随机决策过程，并借鉴强化学习中的策略梯度优化思想来改进预训练过程。通过引入奖励塑造策略，显式地平衡token预测的多样性和精确性，从而优化预训练模型的token输出分布，使其更适合后续的强化学习探索。
技术框架：该方法的核心在于修改预训练目标函数。它基于标准的交叉熵损失，并引入了奖励塑造机制。具体来说，它包含以下几个关键模块：1) 正奖励缩放因子：用于控制模型对ground-truth token的概率集中程度。2) Rank-aware机制：不对称地处理高排名和低排名的负token，鼓励模型区分不同负token的重要性。整体流程是，首先使用修改后的预训练目标函数训练语言模型，然后使用强化学习方法对模型进行微调。
关键创新：该论文的关键创新在于将强化学习的on-policy原则引入到监督学习的预训练阶段，并提出了一种显式平衡多样性和精确性的奖励塑造策略。这种方法不同于传统的交叉熵损失训练，它能够更精细地控制预训练模型的token输出分布，从而为后续的强化学习提供更优的探索空间。
关键设计：关键设计包括：1) 正奖励缩放因子：通过调整该因子的大小，可以控制模型对ground-truth token的关注程度。2) Rank-aware机制：该机制根据负token的排名，给予不同的惩罚权重，鼓励模型区分不同难度的负token。3) 损失函数：修改后的损失函数结合了交叉熵损失和奖励塑造项，能够同时优化模型的预测精度和探索能力。
🖼️ 关键图片
📊 实验亮点
实验结果表明，与传统的交叉熵损失训练相比，该方法能够为强化学习提供更优越的探索空间，从而显著提升端到端的推理性能。具体来说，实验发现，与直觉相反，施加面向精确性的先验能够带来更好的强化学习效果。具体的性能提升数据未知，但论文强调了其优于传统方法的结论。
🎯 应用场景
该研究成果可应用于各种需要通过强化学习微调的大型语言模型任务，例如对话生成、文本摘要、代码生成等。通过优化预训练模型的token输出分布，可以提高强化学习的训练效率和最终性能，从而降低训练成本，提升模型在实际应用中的表现。该方法还有潜力应用于其他序列生成任务。
📄 摘要（原文）
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.