Token-Level Policy Optimization: Linking Group-Level Rewards to Token-Level Aggregation via Sequence-Level Likelihood

📄 arXiv: 2604.12736v1 📥 PDF

作者: Xingyu Lin, Yilin Wen, Du Su, Jinchang Hou, En Wang, Wenbin Liu, Chenfu Bao, Zhonghou Lv

分类: cs.CL

发布日期: 2026-04-14

💡 一句话要点

提出TEPO，通过序列似然和KL散度约束优化LLM的token级策略，提升数学推理能力。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: token级策略优化 序列似然 KL散度约束 数学推理 大型语言模型

📋 核心要点

  1. 现有GRPO方法在token级稀疏奖励下易出现熵崩溃或模型退化，影响CoT推理。
  2. TEPO通过序列似然将组级奖励与token关联，并使用KL散度掩码约束优化策略。
  3. 实验表明，TEPO在数学推理任务上达到SOTA，并显著提升训练稳定性，收敛速度提升50%。

📝 摘要（中文）

群相对策略优化(GRPO)显著提升了大型语言模型(LLM)的推理能力，尤其是在数学推理方面。然而，GRPO和相关的熵正则化方法仍然难以处理token级别的稀疏奖励，这是思维链(CoT)推理中固有的挑战。这些方法通常依赖于未区分的token级熵正则化，这容易导致熵崩溃或模型在稀疏token奖励下性能下降。本文提出了TEPO，一种新颖的token级框架，它(1)利用序列级似然通过token级聚合将组级奖励与单个token联系起来，并且(2)引入了token级KL散度掩码约束，该约束针对具有正优势和熵减少的token，以减轻突发的策略更新。实验表明，TEPO不仅在数学推理基准测试中实现了最先进的性能，而且显著提高了训练稳定性，与GRPO/DAPO相比，收敛时间减少了50%。

🔬 方法详解

问题定义：现有基于群相对策略优化（GRPO）的方法在处理链式思考（CoT）推理中token级别的稀疏奖励时存在困难。这些方法依赖于未区分的token级熵正则化，容易导致熵崩溃或模型性能下降，无法有效利用token级别的奖励信号。

核心思路：TEPO的核心思路是将组级别的奖励通过序列级别的似然性与token级别的聚合联系起来，从而将奖励信号更精细地分配到每个token上。此外，通过引入token级别的KL散度掩码约束，防止策略的剧烈更新，提高训练的稳定性。

技术框架：TEPO框架包含两个主要部分：1) 使用序列级别的似然性来连接组级别的奖励和token级别的聚合。具体来说，通过计算生成序列的概率，将序列级别的奖励反向传播到每个token上。2) 引入token级别的KL散度掩码约束。该约束只针对那些具有正优势（即对结果有积极贡献）并且熵值正在减少的token，从而避免对所有token进行无差别的约束，减少策略更新的突变。

关键创新：TEPO的关键创新在于将序列级别的似然性与token级别的奖励联系起来，并引入了有选择性的KL散度掩码约束。与传统的熵正则化方法不同，TEPO能够更精确地控制每个token的策略更新，避免了熵崩溃和模型退化的问题。

关键设计：TEPO的关键设计包括：1) 序列似然的计算方式，需要选择合适的语言模型和解码策略。2) KL散度掩码的阈值设置，需要根据具体的任务和数据集进行调整，以平衡策略更新的幅度和稳定性。3) 优势函数的计算方式，需要选择合适的基线模型，以准确评估每个token的贡献。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

TEPO在数学推理基准测试中取得了SOTA性能，并且显著提高了训练的稳定性。与GRPO/DAPO相比，TEPO的收敛时间减少了50%，表明其在处理token级别稀疏奖励问题上的有效性。实验结果验证了TEPO在提升LLM推理能力方面的优越性。

🎯 应用场景

TEPO方法可应用于各种需要复杂推理和决策的自然语言处理任务，例如数学问题求解、代码生成、知识图谱推理等。该方法能够提升模型在稀疏奖励环境下的学习能力，提高生成结果的准确性和可靠性，具有广泛的应用前景。

📄 摘要（原文）

Group Relative Policy Optimization (GRPO) has significantly advanced the reasoning ability of large language models (LLMs), particularly in their mathemat ical reasoning performance. However, GRPO and related entropy regularization methods still struggle with token-level sparse-rewards, which is an inherent chal lenge in chain-of-thought (CoT) reasoning. These approaches often rely on undifferen tiated token-level entropy regularization, which easily leads to entropy collapse or model degradation under sparse token rewards. In this work, we propose TEPO, a novel token-level framework that (1) leverages sequence-level likelihood to link group-level rewards with individual tokens via token-level aggregation, and (2) introduces a token-level KL-Divergence mask constraint that targets tokens with positive advantages and decreasing entropy to mitigate abrupt policy updates. Experiments demonstrate that TEPO not only achieves state-of-the-art performance on mathematical reasoning benchmarks but also markedly enhances training stability, reducing convergence time by 50% compared with GRPO/DAPO.