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

清华大学与加州理工学院联合发布TDRM框架:通过时间差分学习优化奖励模型的时间一致性与平滑性

奖励模型(RM)作为大语言模型(LLM)在强化学习(RL)和推理验证中的核心组件,已在数学问题求解、代码生成以及指令跟随等多种任务中展现出卓越的性能表现。

然而,现有奖励模型面临一个关键挑战——缺乏时间一致性,这导致了策略更新效率低下和强化学习训练过程不稳定等一系列问题。

具体来说,LLM推理轨迹中单个步骤的奖励往往与相邻步骤的奖励相互独立,从而在训练中产生不一致且容易误导的信号,同时在推理阶段也难以提供有效的指导。特别是在长思维链(CoT)场景下,这一问题更加显著——模型在完成一长串推理步骤之前无法获得任何奖励,难以准确判断“哪些步骤有用、哪些步骤多余”。

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

值得一提的是,所有相关代码、数据集以及语言模型检查点均已通过GitHub平台开源

清华大学与加州理工学院联合发布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框架:通过时间差分学习优化奖励模型的时间一致性与平滑性 奖励模型 时间一致性 强化学习 思维链 第2张

图|TDRM的整体框架示意图

1.平滑性

平滑性是推理过程中有效奖励建模的一个重要特性,因为它反映了中间步骤中价值更新的一致性和稳定性,确保推理轨迹的微小变化不会导致价值估计出现不成比例的偏差。为评估平滑性,团队采用两种互补方法来对比ScalarPRM和TDRM的表现。

  • 局部普希茨连续(The local Lipschitz constant):用于量化奖励对相邻状态变化的敏感度。分析表明,相比于ScalarPRM(0.3331),TDRM(0.2741)在连续步骤间的平均普希茨连续更小,这表明其奖励过渡更为平滑且时间一致性更好;
  • TD误差(TD error):通过计算连续推理步骤间的TD误差,并结合推理步骤间价值差异,从双重维度评估估计值函数的连续性和一致性。

清华大学与加州理工学院联合发布TDRM框架:通过时间差分学习优化奖励模型的时间一致性与平滑性 奖励模型 时间一致性 强化学习 思维链 第3张

图|奖励模型的平滑度对比

然而,以往研究表明,CoT的长度在LLM推理过程中并非总是稳定增加。结合以上对奖励平滑度的分析,研究团队认为奖励塑造是稳定这种涌现的长度缩放行为的关键机制。

2.奖励建模

根据论文描述,在基于时间差分的PRM框架中,奖励塑造具有双重目的:一方面通过提供结构化反馈优化时间差分更新,另一方面缓解不同推理长度下奖励信号的波动性。包括:

  • 余弦奖励(Cosine Reward):实现了一个基于余弦的奖励函数,以适应每个推理步骤的正确性及其相对长度。它为正确和不正确的步骤分配不同的奖励范围。奖励从最大值开始,随着推理长度接近最大长度而逐渐衰减至最小值。
  • 时间差分:将计算出的余弦奖励与时间差分框架相结合,从而更新过程奖励模型。
  • TD-λ:相比于n步时间差分,TD-λ是一种具有更高灵活性的在线算法。由于其在线特性,TD-λ允许过程奖励模型在观察到奖励后立即将信息传播到更早的状态。
  • 损失函数:为优化过程奖励模型,采用交叉熵损失,将钳位后的时间差分目标作为每个推理步骤的软标签,使模型能够从奖励的时间一致性中学习。

3.强化学习

在强化学习方面,研究团队将其设计为在线算法,在训练过程中动态地使用即时(on-the-fly)状态值计算时间差分目标。与依赖预先计算的状态值的离线算法不同,这一方法能够适应不断变化的轨迹,利用已见轨迹来估计未见轨迹的状态值。这种适应性确保了更准确的价值预测,从而增强了奖励模型的一致性和鲁棒性。

清华大学与加州理工学院联合发布TDRM框架:通过时间差分学习优化奖励模型的时间一致性与平滑性 奖励模型 时间一致性 强化学习 思维链 第4张

图|算法TDRM的处理过程

在TDRM中,通过线性组合的方式将可验证奖励和基于过程的奖励相结合,以发挥两者的互补优势。这一组合后的奖励信号被用于训练GRPO目标函数,从而提升了学习过程的整体性能和数据效率。

更多技术细节详见论文。

实际效果如何?

为验证TDRM的有效性,研究团队测试了TDRM在推理时验证和训练时在线强化学习两种场景下的性能。

  • 针对推理时验证,通过两种关键设置对比不同奖励模型。Best-of-N采样法首先生成包含N个潜在输出的池,随后应用奖励模型确定单一最佳候选方案,旨在平衡输出结果的多样性与最优性。Greedy Search通过迭代选择得分最高的序列来生成输出。
  • 针对训练时在线强化学习,在5个有挑战性的数据集MATH-500Minerva MathOlympiad BenchAIME24AMC23)上,对TDRM与主流方法进行对比。参照SimpleRL方法,使用Pass@1指标结合贪心解码评估最终任务的性能。

1.奖励建模

通过观察不同模型与数据集的Best-of-N采样结果,研究团队为TDRM的优越性提供了实证依据。

首先,在MATH-500数据集上,随着采样预算从Best-of-128增加到Best-of-1024,TDRM的表现明显优于ScalarPRM和ScalarORM。这充分证明TDRM具有更强的可靠性,能在更大采样预算下持续识别最优响应。

清华大学与加州理工学院联合发布TDRM框架:通过时间差分学习优化奖励模型的时间一致性与平滑性 奖励模型 时间一致性 强化学习 思维链 第5张

表|MATH-500测试结果;GSM8K上的Best-of-128结果

在树搜索评估中,TDRM再次展现出更优的性能,并提供了更准确的推理轨迹验证。并且,TDRM的准确性随着搜索分支数量的增加而提升,体现出它在复杂决策空间中导航的有效性。

清华大学与加州理工学院联合发布TDRM框架:通过时间差分学习优化奖励模型的时间一致性与平滑性 奖励模型 时间一致性 强化学习 思维链 第6张

图|树搜索结果

2.强化学习

TDRM在仅使用2500条MATH Level-3提示的有限数据集上,成功超越了8种主流模型,实现了最高的平均准确率,凸显了其在强化学习训练中的可靠性。

TDRM通过结合可验证奖励和基于过程的奖励,确保了稳定的性能和更优的数据效率,即使在训练样本有限的情况下也能实现持续学习

清华大学与加州理工学院联合发布TDRM框架:通过时间差分学习优化奖励模型的时间一致性与平滑性 奖励模型 时间一致性 强化学习 思维链 第7张

表|在5个系列8个基础模型进行强化学习训练之后在数学基准上的评测结果

以上结果表明,将时间一致性整合到奖励模型不仅有助于提升RL训练的稳定性,还为构建更具可扩展性的RLHF流程、实现更高质量的推理搜索,以及推动LLM在复杂目标结合中的广泛应用提供了新的可能性。