PRIME:通过隐式奖励实现过程强化的可扩展强化学习解决方案

摘要

近年来,大规模语言模型(LLMs)的推理能力在数据驱动的模仿学习(Imitation Learning)下不断提升,但其扩展性仍面临严峻挑战。为了突破这些瓶颈,PRIME团队提出了一种新颖的强化学习方法——PRIME(Process Reinforcement through Implicit Rewards),通过隐式过程奖励(Implicit Process Reward)有效地解决了奖励稀疏和可扩展性的问题。本文将详细介绍PRIME算法的具体实现,包括完整的算法流程、关键设计以及实验结果,以展示其在提升语言模型复杂推理能力方面的卓越性能。


友情链接: 借一步  背多分   ACEJoy


 

引言

大规模语言模型在自然语言处理任务中展现出强大的能力,但随着模型规模的增加,基于模仿学习的方法逐渐暴露出扩展性不足的问题。更强的推理能力通常需要更多高质量的数据进行训练,而高质量数据资源的获取成本高且数量有限,导致持续改进变得困难。为了解决这一问题,PRIME团队提出了将数据驱动的方法转化为基于探索的强化学习方法,通过强化学习(Reinforcement Learning, RL)实现模型能力的进一步提升。

PRIME方法基于隐式过程奖励模型(Implicit Process Reward Model, PRM),无需任何过程标签,仅通过结果奖励模型(Outcome Reward Model, ORM)进行训练,隐式地建模过程奖励。这一方法不仅提升了模型的推理能力,还在在线强化学习训练中展现出三大显著优势:密集奖励(Dense Reward)、可扩展性(Scalability)和简洁性(Simplicity)。本文将深入探讨PRIME算法的实现细节及其在多个数学基准测试中的表现。

PRIME算法概述

PRIME(Process Reinforcement through Implicit Rewards)是一种在线强化学习算法,旨在通过隐式过程奖励有效提升语言模型的推理能力。其核心思想是利用隐式过程奖励模型为每个生成的token提供奖励,从而解决奖励稀疏的问题,并通过在线更新机制实现奖励模型的可扩展性和简洁性。

隐式过程奖励模型

隐式过程奖励模型(Implicit PRM)通过结果奖励模型(ORM)进行训练,仅依赖答案的最终对错信息,自动学习过程奖励。这一模型直接学习一个Q函数,为每个token提供奖励估计,避免了额外价值模型的需求。此外,隐式PRM能够仅通过结果标签在线更新,使其更具可扩展性,并能够有效缓解分布偏移问题。

PRIME算法流程

PRIME算法的具体实现流程如下所示:

  1. 初始化:使用SFT(Supervised Fine-Tuning)模型作为策略模型(Policy Model, πθ\pi_\thetaπθ​)和隐式PRM(πψ\pi_\psiπψ​)的初始模型。
  2. 生成Rollouts:策略模型生成一批rollouts,即多个任务的回答序列。
  3. 奖励计算
    • 结果奖励(Outcome Reward, ror_oro​):通过结果验证器对生成的回答进行评分,确定答案的最终对错。
    • 过程奖励(Process Reward, rpr_prp​):隐式PRM基于每个生成的token提供的奖励估计。
  4. 更新隐式PRM:利用生成的rollouts和结果奖励,在线更新隐式PRM。
  5. 优势估计(Advantage Estimation):分别计算结果奖励和过程奖励的回报,通过RLOO(Reward Leave-One-Out)方法估计优势(Advantage)。
  6. 策略更新:结合结果奖励和过程奖励,通过PPO(Proximal Policy Optimization)损失函数更新策略模型。
  7. 迭代训练:重复上述步骤,直至模型性能达到预期目标。

以下是PRIME算法的伪代码示意:

PRIME算法伪代码:

1. 初始化策略模型 πθ 和隐式PRM πψ 使用SFT模型
2. 重复直到收敛:
    a. 使用策略模型 πθ 生成rollouts
    b. 使用隐式PRM πψ 和结果验证器对rollouts评分,获得 r_o 和 r_p
    c. 在线更新隐式PRM πψ 基于 r_p 和实际结果 r
    d. 计算结果奖励 R_o 和过程奖励 R_p
    e. 估计优势 A = R_o + R_p
    f. 使用PPO损失函数更新策略模型 πθ

具体实现细节

数据准备与预处理

PRIME算法依赖高质量的数学数据进行训练。数据预处理步骤主要包括:

  • 数据收集:从多个公开数据集(如NuminaMath、APPS、CodeContests、TACO、Codeforces)收集数学问题和代码竞赛题目。
  • 数据清洗:过滤掉无关或错误的数据,确保数据的质量和多样性。
  • 数据格式化:将收集到的数据转换为适合模型训练的格式,便于后续的训练和评估。

模型训练

PRIME的训练过程包括初始的SFT训练和在线的强化学习训练。

  1. 初始SFT训练
    • 使用Qwen2.5-Math-7B作为基座模型,进行监督微调(SFT),得到初始的策略模型Eurus-2-7B-SFT。
  2. 隐式PRM训练
    • 隐式PRM不需要额外的过程标签,仅通过结果奖励模型进行训练。
    • 利用生成的rollouts和对应的结果奖励,在线更新隐式PRM,使其能够为每个token提供准确的奖励估计。
  3. 强化学习训练
    • 使用PRIME算法的完整流程,结合PPO优化策略模型,使其在数学推理任务上表现更为优秀。
    • 通过在线更新隐式PRM,不断提升模型在复杂推理任务中的表现。

算法关键步骤详解

Prompt Filtering

在每次强化学习迭代中,首先对生成的rollouts进行过滤,只保留那些模型性能在一定范围内(准确率在0.2至0.8之间)的样本。这一过滤步骤旨在确保模型在有挑战性的任务上进行训练,避免过于简单或过于复杂的问题影响训练效果。

计算隐式过程奖励

使用隐式PRM对每个token计算过程奖励(rtr^trt),这一奖励直接反映了每个生成的token对最终结果的贡献。隐式PRM通过学习Q函数,为每个token提供密集的奖励信号,解决了奖励稀疏的问题。

更新隐式PRM

根据预测的隐式过程奖励(rtr^trt)和实际的结果奖励(rrr),在线更新隐式PRM。这一更新过程仅依赖于结果标签,避免了对额外过程标签的需求,提高了训练的可扩展性。

优势估计(Advantage Estimation)

优势估计是强化学习中的关键步骤,用于评估当前策略相对于基准策略的改进程度。在PRIME算法中,优势估计分为两部分:

  1. 结果奖励的回报:直接采用RLOO方法计算结果奖励的回报。
  2. 过程奖励的回报
    • 使用平均隐式过程奖励计算leave-one-out基线(baseline)。
    • 对每个时间步的过程奖励进行归一化处理(即减去基线)。
    • 计算每个响应的折扣回报。

最终,优势(A)被设定为结果奖励和过程奖励回报的组合。

策略更新

利用PPO(Proximal Policy Optimization)损失函数,结合优势估计,对策略模型进行更新。PPO通过限制策略更新的步幅,确保训练的稳定性和效率。

使用示例

PRIME算法提供了针对编码和数学任务的定制化提示(prompt),以下分别展示其应用方式:

编码任务提示

{question} + "\n\nWrite Python code to solve the problem. Present the code in \n```python\nYour code\n```\nat the end."

数学任务提示

{question} + "\n\nPresent the answer in LaTex format: \boxed{Your answer}"

以下是一个具体的Inference代码示例,用于生成数学问题的解答:

import os
from tqdm import tqdm
import torch
from transformers import AutoTokenizer
from vllm import LLM, SamplingParams

os.environ["NCCL_IGNORE_DISABLED_P2P"] = "1"
os.environ["TOKENIZERS_PARALLELISM"] = "true"

def generate(question_list, model_path):
    llm = LLM(
        model=model_path,
        trust_remote_code=True,
        tensor_parallel_size=torch.cuda.device_count(),
        gpu_memory_utilization=0.90,
    )
    sampling_params = SamplingParams(max_tokens=8192, temperature=0.0, n=1)
    outputs = llm.generate(question_list, sampling_params, use_tqdm=True)
    completions = [[output.text for output in output_item.outputs] for output_item in outputs]
    return completions

def make_conv_hf(question, tokenizer):
    # 针对数学问题
    content = question + "\n\nPresent the answer in LaTex format: \boxed{Your answer}"
    msg = [{"role": "user", "content": content}]
    chat = tokenizer.apply_chat_template(msg, tokenize=False, add_generation_prompt=True)
    return chat

def run():
    model_path = "PRIME-RL/Eurus-2-7B-PRIME"
    all_problems = [
        "which number is larger? 9.11 or 9.9?"
    ]
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    completions = generate([make_conv_hf(problem_data, tokenizer) for problem_data in all_problems], model_path)
    print(completions)
    # 输出示例:
    # [['[ASSESS]\n\n# The problem asks us to compare two decimal numbers, 9.11 and 9.9, to determine which one is larger.\n# We need to compare the whole parts and the decimal parts of the numbers.\n\nNext action: [ADVANCE]\n\n# Compare the whole parts of the numbers: both 9.11 and 9.9 have the same whole part, which is 9.\n# Compare the decimal parts of the numbers: 0.11 (from 9.11) is less than 0.9 (from 9.9).\n\nNext action: [ADVANCE]\n\n# Since the whole parts are the same and the decimal part of 9.9 is greater than the decimal part of 9.11, we can conclude that 9.9 is larger than 9.11.\n\nNext action: [OUTPUT]\n\nThe final answer is \boxed{9.9}.\n\n']]
if __name__ == "__main__":
    run()

实验结果

PRIME算法在多个数学推理基准测试中相较于监督微调(SFT)模型和其他先进模型表现出显著提升。以下是主要实验结果:

Eurus-2-7B-PRIMEEurus-2-7B-SFTQwen-2.5-Math-7B-InstructLlama-3.1-70B-InstructGPT-4o
AIME 202426.7 (+23.3)3.313.316.79.3
MATH-50079.2 (+14.1)65.179.864.676.4
AMC57.8 (+27.7)30.150.630.145.8
Minerva Math38.6 (+5.9)32.734.635.336.8
OlympiadBench42.1 (+12.3)29.840.731.943.3
Avg.48.9 (+16.7)32.243.835.743.3

数据与资源对比

PRIME方法在数据和模型资源的需求上相比于Qwen-Math展示了更高的效率:

Eurus-2-7B-PRIMEQwen2.5-Math-7B-Instruct
Base ModelQwen2.5-Math-7BQwen2.5-Math-7B
SFT Data230K (开源)2.5M (开源及内部)
RM Data0618K (内部)
RMEurus-2-7B-SFTQwen2.5-Math-RM (72B)
RL Data150K queries × 4 samples66K queries × 32 samples

结论

PRIME(Process Reinforcement through Implicit Rewards)通过隐式过程奖励模型,有效地解决了大规模语言模型在强化学习中的奖励稀疏和可扩展性问题。该方法不仅提升了模型在多个数学推理基准测试中的表现,还在数据和计算资源的利用效率上表现出色。未来,结合更强大的基座模型,PRIME有望进一步缩小与最先进模型(如OpenAI o1)的性能差距,推动大规模语言模型在复杂推理任务中的应用前景。

参考文献

@misc{cui2024process,
  title={Process Reinforcement through Implicit Rewards},
  author={Ganqu Cui and Lifan Yuan and Zefan Wang and Hanbin Wang and Wendi Li and Bingxiang He and Yuchen Fan and Tianyu Yu and Qixin Xu and Weize Chen and Jiarui Yuan and Huayu Chen and Kaiyan Zhang and Xingtai Lv and Shuo Wang and Yuan Yao and Hao Peng and Yu Cheng and Zhiyuan Liu and Maosong Sun and Bowen Zhou and Ning Ding},
  year={2025},
  howpublished={\url{https://curvy-check-498.notion.site/Process-Reinforcement-through-Implicit-Rewards-15f4fcb9c42180f1b498cc9b2eaf896f}},
  note={Notion Blog}
}
@article{yuan2024implicitprm,
  title={Free Process Rewards without Process Labels},
  author={Lifan Yuan and Wendi Li and Huayu Chen and Ganqu Cui and Ning Ding and Kaiyan Zhang and Bowen Zhou and Zhiyuan Liu and Hao Peng},
  journal={arXiv preprint arXiv:2412.01981},
  year={2024}
}

致谢

我们在算法实现上扩展自veRL的强化学习算法。推理过程中我们使用vLLM,并基于EurusQwen2.5-Math以及LiveCodeBench开发了评估脚本。我们的数据来源主要包括NuminaMathAPPSCodeContestsTACOCodeforces。感谢这些数据集的贡献者们!

GitHub Star 历史

[Star History Chart](https://star-history.com/#PRIME-RL/PRIME&Date)

版权声明

本文基于PRIME项目的公开资料撰写,内容详实,旨在全面介绍PRIME算法及其实现细节,供学术交流与研究参考。

联系我们

更多信息请访问我们的GitHub仓库或阅读我们的Notion博客

免责声明

本文内容基于现有公开资料整理,所涉及的观点和结论仅代表作者个人理解,与项目团队观点一致。任何未经授权的转载和使用均属侵权,敬请谅解。

结束语

PRIME算法通过创新性的隐式过程奖励机制,成功地提升了大规模语言模型在复杂推理任务中的表现,展示了强化学习在语言模型训练中的巨大潜力。随着算法的进一步优化和数据资源的丰富,PRIME有望在未来的人工智能发展中扮演重要角色,推动语言模型能力的持续进步。

评论

发表回复

您的邮箱地址不会被公开。 必填项已用 * 标注

人生梦想 - 关注前沿的计算机技术 acejoy.com 🐾 步子哥の博客 🐾 背多分论坛 🐾 借一步网
Page Stats: PV: 5 | UV: 5
Last updated: 2025-04-20 15:54:36