
Qwen3 低成本手撕Search-R1的强化学习训练框架
OpenAI-o1和DeepSeek-R1,通过奖励驱动的学习,而无需依赖于明确的逐步监督,在逻辑推理和迭代推理方面取得了显著的成果。
使用RL训练能够更有效地搜索相关信息的策略模型。例如,Search-R1、R1-Searcher和ReSearch等模型都尝试通过强化学习来提升模型的搜索能力。
这些方法通常使用静态的本地文本语料库,如维基百科,无法捕捉到现实世界交互的复杂性。
DeepResearcher引入了与商业搜索引擎(如谷歌)的实时交互,使得模型能够在接近真实世界网络搜索的环境中进行训练。
但它面临着重大挑战:
- 一是搜索引擎返回的文档质量不可预测,这给训练过程带来了噪声和不稳定性;
- 二是强化学习训练需要频繁的rollout,可能涉及数十万次搜索请求,这导致了巨大的API费用,严重限制了系统的可扩展性。
一、ZEROSEARCH核心思想
- ZEROSEARCH利用LLMs在大规模预训练过程中所积累的丰富世界知识,使其能够根据搜索查询生成相关文档。通过轻量级的监督式微调,即使是相对较小的LLMs也能够有效地模拟真实搜索引擎的行为。
- 通过模拟搜索引擎的交互过程来训练策略模型,使其能够在没有真实搜索引擎的情况下有效地进行信息检索和推理。
- 其创新点在于,通过控制文档质量,解决搜索引擎返回文档质量不可预测的问题,并且消除了API费用,极大地提高了训练的可扩展性。
二、架构介绍
- ZEROSEARCH的整体架构围绕着一个强化学习框架展开,其核心组件包括:
模拟LLMs在训练过程中扮演着搜索引擎的角色,它接收策略模型生成的搜索查询,并返回相应的文档,这些文档可以是有用的,也可以是噪声的,具体取决于训练阶段和课程策略的设置。
奖励模型根据策略模型生成的答案与真实答案之间的匹配程度给出奖励信号,以驱动策略模型的学习过程。
参考模型则用于计算策略模型输出与参考输出之间的差异,从而为策略模型的更新提供指导。
策略模型是需要优化的对象,负责根据输入的查询生成搜索请求,并基于检索到的信息生成最终答案。
一个策略模型(policy model)
一个参考模型(reference model)
一个奖励模型(reward model)
一个模拟LLMs(simulation LLM)
三、技术细节
- 强化学习目标:
ZEROSEARCH的优化目标是最大化策略模型在给定查询下生成正确答案的期望奖励,同时最小化策略模型输出与参考模型输出之间的差异。
具体来说,优化目标可以表示为:
- 训练模板:
内部推理阶段(<think>...</think>
)
搜索查询阶段(<search>...</search>
)
答案生成阶段(<answer>...</answer>
)。
- ZEROSEARCH采用了一个多轮交互模板,将策略模型的输出分为三个阶段:
- 这种结构化的输出格式有助于模型更好地组织推理过程,并使训练过程更加透明和可控。
- 课程学习策略,是ZEROSEARCH的一个关键创新点:
通过逐渐增加生成文档的噪声比例,模拟越来越复杂的检索场景,从而逐步提升策略模型的推理能力。
这种策略类似于人类学习过程中的逐步引导,有助于模型更好地适应不同难度的任务。
- 损失掩码机制:
在训练过程中,ZEROSEARCH对检索到的文档token应用损失掩码,避免了这些外部生成的token对策略模型更新过程的干扰。
这一机制有助于稳定训练过程,并提高模型的最终性能。
四、整体流程
- 数据准备:首先,收集与真实搜索引擎交互的轨迹数据,并将其标记为正样本或负样本,用于后续的轻量级监督式微调。
- 轻量级监督式微调:使用标记好的轨迹数据对LLMs进行微调,使其能够根据搜索查询生成有用或噪声文档。
- 策略模型初始化:初始化策略模型,并将其与模拟LLMs和奖励模型连接起来,形成完整的强化学习框架。
- 训练:
在训练过程中,策略模型根据输入的查询生成搜索请求,模拟LLMs根据搜索请求生成文档,策略模型根据生成的文档进行推理并生成答案。
奖励模型根据答案的准确性给出奖励信号,策略模型根据奖励信号进行更新。
随着训练的进行,逐渐增加生成文档的噪声比例,以提升策略模型的推理能力。
- 模型评估:在训练完成后,使用各种问答数据集对模型的性能进行评估,比较其与真实搜索引擎和其他基线方法的性能差异。
举个例子:
假设问题是“谁是《冰与火之歌》系列小说的作者?”
- 策略模型首先在
<think>...</think>
阶段进行内部推理,可能会思考“我需要找到《冰与火之歌》系列小说的作者是谁”。 - 然后,在
<search>...</search>
阶段,策略模型生成一个搜索查询,如“《冰与火之歌》作者”。模拟LLMs接收到这个查询后,生成一些文档,这些文档可能是有用的,也可能是噪声的,具体取决于当前的课程阶段。 - 策略模型根据这些文档在
<answer>...</answer>
阶段生成最终答案,如“乔治·R·马丁”。 - 奖励模型根据生成的答案与真实答案之间的匹配程度给出奖励信号,策略模型根据这个奖励信号进行更新,以提升其在后续任务中的性能。
4.1 轻量级监督式微调的详细步骤
- 数据收集:
通过让LLMs与真实搜索引擎进行多轮交互,收集一系列的交互轨迹。
这些轨迹包括查询、搜索引擎返回的文档以及LLMs基于这些文档生成的答案。
- 数据标注:
根据LLMs生成的答案是否正确,将这些轨迹标注为正样本或负样本。
如果答案正确,则将该轨迹标记为正样本,表示检索到的文档是有用的;
如果答案错误,则将该轨迹标记为负样本,表示检索到的文档是噪声。
- 提取查询-文档对:
从正负样本轨迹中提取查询-文档对。
这些查询-文档对将用于后续的微调过程。
- 设计微调提示:
为了使LLMs能够生成有用或噪声文档,设计微调提示。
在提示中加入“有用”或“噪声”等关键词,并将输入问题及其对应答案纳入提示中,以拓宽LLMs的知识边界。
- 微调过程:
使用提取的查询-文档对和设计好的微调提示对LLMs进行轻量级监督式微调。
在微调过程中,调整LLMs的参数,使其能够根据搜索查询生成符合提示要求的文档。
- 评估微调效果:
通过一些验证数据集评估微调后的LLMs的性能,确保其能够有效地模拟真实搜索引擎的行为,并且能够根据提示生成有用或噪声文档。
4.2 基于课程的rollout策略的详细步骤
- 初始化课程参数:
设置初始噪声概率ps、最终噪声概率pe以及指数基数b等课程参数。
这些参数将决定在训练过程中生成文档的噪声比例如何变化。
- 策略模型生成搜索查询:
在每次rollout过程中,策略模型根据输入的查询生成一个或多个搜索查询。
这些查询将被发送到模拟LLMs。
- 模拟LLMs生成文档:
模拟LLMs接收到搜索查询后,根据当前的课程阶段生成相应的文档。
在训练初期,生成的文档主要是有用的,随着训练的进行,逐渐增加噪声文档的比例。
- 策略模型进行推理:
策略模型根据模拟LLMs返回的文档进行推理,并生成最终答案。
在推理过程中,策略模型可能会根据需要进行多轮搜索。
- 奖励模型给出奖励信号:
奖励模型根据策略模型生成的答案与真实答案之间的匹配程度给出奖励信号。
如果答案正确,奖励信号较高;如果答案错误,奖励信号较低。
- 更新策略模型:
根据奖励信号和课程参数,更新策略模型的参数。
在更新过程中,策略模型将学习如何在不同难度的检索场景中生成更准确的答案。
- 调整课程阶段:
随着训练的进行,逐步调整课程参数,增加生成文档的噪声比例。
这将使策略模型逐渐适应更具挑战性的检索场景,提升其推理能力。
4.3 训练算法的细节
- 选择合适的强化学习算法:
根据具体任务和模型的特点,选择合适的强化学习算法,如PPO、GRPO或Reinforce++等。
- rollout过程:
在rollout过程中,策略模型根据输入的查询生成搜索查询,模拟LLMs根据搜索查询生成文档,策略模型根据生成的文档进行推理并生成答案。
这个过程将产生一系列的交互数据,包括查询、文档、答案以及对应的奖励信号。
- 计算损失函数:
策略损失:衡量策略模型输出与参考模型输出之间的差异,
价值损失:衡量价值模型预测与真实奖励之间的差异,
熵损失:用于鼓励策略模型的探索行为。
根据策略模型的输出和奖励信号,计算损失函数。
损失函数通常包括策略损失、价值损失以及熵损失等。
- 更新模型参数:
根据计算得到的损失函数,使用反向传播算法更新策略模型、参考模型和奖励模型的参数。
在更新过程中,应用梯度裁剪等技术,以避免梯度爆炸或梯度消失的问题。
- 重复训练过程:
我
ZEROSEARCH优点
- ZEROSEARCH框架通过轻量级监督式微调将LLMs转变为一个检索模块,并采用基于课程的rollout策略逐步提升策略模型的推理能力,有效地解决了现有基于强化学习的LLMs搜索能力训练方法中存在的搜索引擎返回文档质量不可预测和API费用高昂的问题。
- ZEROSEARCH在多个问答数据集上均取得了优异的性能,超越了真实搜索引擎和其他基线方法。
- ZEROSEARCH还具有良好的可扩展性,能够适应不同大小和类型的LLMs,并且与多种强化学习算法兼容。
- ZEROSEARCH在部署模拟LLMs时需要一定的GPU资源,但其成本远低于商业搜索引擎API费用,具有较高的经济可行性。
- ZEROSEARCH可以应用于其他类型的自然语言处理任务,如文本摘要、机器翻译等,以充分发挥其在提升LLMs性能方面的潜力。
本文转载自CourseAI,作者:CourseAI
