当前位置:首页 > 科技资讯 > 正文

TDRM框架:优化奖励模型,提升大语言模型性能

奖励模型(RM)在基于大语言模型(LLM)的强化学习(RL)和推理时验证中占据核心地位,已在数学求解、代码生成和指令遵循等任务中展现出卓越性能。然而,现有奖励模型存在关键缺陷——缺乏时间一致性,导致“策略更新效果不佳”和“强化学习训练不稳定”等问题。

具体而言,LLM推理轨迹中某一步的奖励往往与相邻步骤无关,从而在训练过程中出现不一致、容易误导的信号,并在推理阶段难以提供有效引导。尤其长思维链(CoT)场景中,这些问题更为突出——模型在完成一长串推理步骤前无法获得任何奖励,难以判断“哪步有用、哪步多余”。

针对这一痛点,清华大学团队联合加州理工学院提出了TDRM框架——通过训练过程中最小化时间差分(TD)来学习更平滑、更可靠的奖励模型。

值得一提的是,所有代码、数据和语言模型检查点已在GitHub上开源。

TDRM框架:优化奖励模型,提升大语言模型性能 TDRM框架 奖励模型 时间一致性 强化学习 第1张

论文链接:https://arxiv.org/abs/2509.15110

GitHub地址:https://github.com/THUDM/TDRM

研究结果显示,经过TD训练的过程奖励模型(PRM)在Best-of-N和树搜索场景中,分别能够取得最高6.6%和23.7%的性能提升。

当与可验证奖励强化学习(RLVR)结合时,经过TD训练的过程奖励模型能够实现更高数据效率的强化学习——仅用2.5k数据就能达到基线方法需要50.1k数据才能达到的相当性能,并在Qwen2.5-(0.5B, 1.5B)、GLM4-9B-0414、GLM-Z1-9B-0414等8种模型变体上得到更高质量的语言模型策略。

打造更平滑、更可靠的奖励模型

与以往将时间差分用于构建中间奖励信号离线数据集的方法不同,TDRM采用时间差分学习来构建用于强化学习训练的可靠奖励模型,生成更平滑的奖励空间、更密集的奖励信号。

据论文描述,TDRM框架包含以下三个核心模块:

  • 过程奖励模型:通过n步时间差分学习结合奖励塑造训练得到过程奖励模型。
  • 强化学习:在训练好的过程奖励模型的引导下进行在线强化学习,从而优化策略更新。
  • TDRM集成:将过程奖励与可验证奖励进行有效的线性组合,并应用于不同策略模型系列和规模的演员-评论家(Actor-Critic)式在线强化学习中。

TDRM框架:优化奖励模型,提升大语言模型性能 TDRM框架 奖励模型 时间一致性 强化学习 第2张

图|TDRM的整体框架示意图