arXiv:2505.16826cs.AIcs.CL2025-05NeurIPS被引 23

提出细粒度的关键词优势估计方法,提升数学推理模型训练效果

KTAE: A Model-Free Algorithm to Key-Tokens Advantage Estimation in Mathematical Reasoning

  • 通过统计分析采样结果,量化每个词元对最终答案的贡献
  • 在5个数学推理数据集上表现更优,响应更短且准确率更高
  • 无需额外模型,可无缝集成到现有强化学习框架中

近期研究显示,将强化学习与基于规则的奖励结合,可在无需监督微调的情况下显著提升大语言模型的推理能力。然而,主流强化学习算法如GRPO及其变体DAPO存在粒度粗糙的问题:它们计算的是序列级别的优势值,使同一序列中的所有词元获得相同优势,无法捕捉词元级的贡献,阻碍有效学习。为此,我们提出关键词优势估计(KTAE)——一种无需引入额外模型的新型算法。KTAE利用采样轨迹的正确性,通过统计分析量化序列内每个词元对最终结果的重要性,再将其与序列级优势结合,得到更精细的词元级优势估计。实验表明,使用GRPO+KTAE和DAPO+KTAE训练的模型在五个数学推理基准上均优于基线方法。尤为突出的是,其生成响应更短,准确率更高,甚至超越了使用相同基础模型的R1-Distill-Qwen-1.5B。

原文摘要 · Abstract (English)

Recent advances have demonstrated that integrating reinforcement learning with rule-based rewards can significantly enhance the reasoning capabilities of large language models, even without supervised fine-tuning. However, prevalent reinforcement learning algorithms such as GRPO and its variants like DAPO, suffer from a coarse granularity issue when computing the advantage. Specifically, they compute rollout-level advantages that assign identical values to every token within a sequence, failing to capture token-specific contributions and hindering effective learning. To address this limitation, we propose Key-token Advantage Estimation (KTAE) - a novel algorithm that estimates fine-grained, token-level advantages without introducing additional models. KTAE leverages the correctness of sampled rollouts and applies statistical analysis to quantify the importance of individual tokens within a sequence to the final outcome. This quantified token-level importance is then combined with the rollout-level advantage to obtain a more fine-grained token-level advantage estimation. Empirical results show that models trained with GRPO+KTAE and DAPO+KTAE outperform baseline methods across five mathematical reasoning benchmarks. Notably, they achieve higher accuracy with shorter responses and even surpass R1-Distill-Qwen-1.5B using the same base model.

强化学习数学推理优势估计大模型

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。