强化递归语言模型（18分钟阅读）

# 强化递归语言模型 | alphaXiv 来源: https://www.alphaxiv.org/blog/reinforcement-learning-for-rlms *我们通过强化学习对小型（4B）模型进行微调，使其成为原生递归语言模型（RLM），具体方法是在单一共享策略下训练父RLM和子RLM。借助强化学习，小型模型能够学习到特定任务的RLM行为，这些行为无法通过提示工程甚至监督微调（SFT）来激发。本文假设读者对RLM有基本了解。关于RLM的优秀学习资源可参考原始的RLM博客文章。* 目录 我们研究了对4B模型进行强化学习微调，以便在生产环境中用作RLM。虽然RLM是一种强大的推理策略，但它们可能存在不可预测的延迟，并且需要大量的提示调优才能获得一致的行为。强化学习微调使我们能够训练出成本低廉的专用RLM。 我们没有为父RLM和子RLM分别训练策略模型，而是训练一个模型来同时扮演父分解器和子代理的角色。我们采用简单的强化学习训练设置，其中子RLM的轨迹继承自生成它们的父RLM轨迹的优势，这使我们能够训练单一策略。这消除了为单个子RLM轨迹提供额外奖励信号的需要。 在涉及多篇科学文档的证据选择任务上，我们展示了经过强化学习微调的4B模型在相同的RLM框架和REPL环境下，性能与Claude Sonnet 4.6相当，而模型大小和成本却减少了许多数量级。 我们的代码（https://github.com/NovaSky-AI/SkyRL/pull/1596）包含训练脚本、RLM支架的实现以及证据选择环境，可在SkyRL（https://github.com/NovaSky-AI/SkyRL）上获取。 RLM在程序化环境中生成语言模型，该环境存储较长的用户提示（上下文），传统上这些提示直接馈入LM的上下文窗口。在此环境中，上下文是一个外部对象，LM可以通过程序化操作检查它，并通过递归调用自身来分解它，最终目标是回答关于该上下文的用户查询。 与原始RLM论文一样，我们使用Python的读取-评估-打印循环（REPL）作为环境。RLM不将代码执行仅仅视为另一种工具，而是在REPL中让代码成为模型检查和转换数据的主要接口。每一轮，模型编写要执行的代码，REPL执行代码，RLM编排器将执行代码的结果（主要是`print()`语句）作为用户消息返回给模型，用于下一轮。 REPL向模型暴露一组内置函数： - **`FINAL(answer)`** / **`FINAL\_VAR(variable\_name)`** — 标记轨迹结束：`FINAL(...)` 返回字面字符串作为最终答案，而 `FINAL_VAR(...)` 则按名称查找现有REPL变量并返回其值。 - **`rlm_query(prompt, context=None)`** — 使用给定的提示（和可选的上下文覆盖）生成单个子RLM轨迹，在同一策略下运行新的代理循环，并将子的最终答案字符串返回给父级的REPL。 - **`rlm_query_batched(prompts, context_list=None)`** — 与 `rlm_query` 相同，但并行调度多个子（每个提示一个，并与相应的上下文配对），并按顺序返回它们的最终答案列表。 以编程方式与上下文交互并生成子RLM或子LM的能力，使得RLM成为处理长上下文问题的非常强大的推理策略。 在本博客中，我们将聚焦于从科学文档中选择证据的任务。给定一个问题和一组合集论文，目标是返回论文中回答该问题的片段。对于此任务，存储在REPL中的上下文是给定问题对应的所有论文的全文。 以单篇论文为例，考虑关于原始RLM论文的以下问题：*本文使用了哪些基线方法？* RLM可以使用的一种简单策略是，在上下文中使用关键字“baseline”调用关键字搜索函数，分析匹配结果，然后返回最相关的片段。 当扩展到多篇论文时，显然任务可以通过子RLM进行分解和并行化。一个直观的方法是由根级RLM识别哪些论文值得进一步探索，并生成子RLM从单篇论文中提取片段。这里值得强调的一点是，RLM的好处不仅仅局限于长上下文任务。这个特定任务的上下文一开始并不特别长（我们想尝试处理几百篇论文的情况），但它高度可并行化。拥抱任务的并行性不仅能改善墙钟时间，实际上还能提升任务性能。许多研究表明，顺序推理会将LLM引入“前缀”陷阱，即模型最终只探索它最初探索的内容。 与其他RLM通常评估的任务的一个重要区别是，根级RLM生成的是真正的子RLM，它们可以访问REPL环境。对于其他数据集（如OOLONG），RLM仅被训练为对外部冻结模型进行子LM调用。 为了抽象某些基本操作，我们用四个预定义函数初始化REPL环境： - `list_papers(ctx)` — 遍历上下文字典中的每篇论文，并输出每篇论文的ID、标题和摘要。返回标题列表。 - `search(text, keyword, window=300, bidirectional=True)` — 不区分大小写的关键字搜索，返回每个匹配项周围的对齐句子片段。如果 `text` 是字典，则搜索每篇论文并按论文ID分组结果；否则搜索单个字符串。 - `extract_section(snippet, start_phrase, end_phrase)` — 提取 `snippet` 中从 `start_phrase` 开始到 `end_phrase` 结束的子字符串（包含，不区分大小写）。如果未找到任一短语，则回退到片段开头或运行到结尾。 - `get_paper_abstract(ctx, paper_id)` — 在上下文字典中查找一篇论文，并返回格式化的“论文ID / 标题 / 摘要”字符串，用于标记子提示，以便工作模型知道它正在处理哪篇论文。 这些预定义函数类似于传统多轮ReAct代理循环中的工具调用。 我们合成生成科学查询和支持证据。我们首先在alphaXiv上选择随机论文，然后检索最多9篇语义相似的论文组成一个组，优先选择高点赞的论文。然后，OCR模型将组中的每篇论文分成段落，我们让前沿模型生成问题，并从该问题的论文子集中选择相关段落。并非组中每篇论文都有与问题相关的段落。OCR模型仅用于构建干净的真实证据。在测试时，模型从PDF解析库接收噪声文本，模拟在生产环境中对每篇新文档运行OCR既繁琐又昂贵的情况。总共，我们合成了1000个查询，覆盖最多10篇论文的组，每组最多三个查询。 另一点要明确的是，RAG不太适合此任务。模型需要动态返回不同长度和数量的逐字文本跨度，而不是固定大小的top-k。不过，RAG和RLM并非互斥，索引块可以作为另一个REPL工具与search和extract_section一起暴露出来。 原始RLM论文对Qwen3-8B进行了SFT，使用了来自Qwen3-Coder-480B-A35B-Instruct的RLM式推理轨迹。SFT的目的是教模型如何导航REPL环境，包括提交答案和进行子LM调用的RLM特定语法。然而，SFT并未用于针对特定任务或数据集训练RLM。 对于推理模型和传统代理框架，强化学习已被证明是一种稳健的方法，可以在不伴随纯SFT方法所带来的灾难性遗忘的情况下，提高模型在特定任务上的能力。我们如何将其扩展到RLM？ ### 单层RLM（无子调用） 在考虑RLM的递归元素之前，我们首先想确认可以在单篇论文的证据选择任务上对Qwen3.5-4B进行强化学习微调，该任务不涉及生成子RLM。这些初步运行中有几个值得分享的要点。 **任务策略。** 特定的任务策略必须放在提示中。这包括预定义函数的描述，用于指导RLM并加速轨迹生成。对于单篇论文任务，策略是使用多个关键字变体通过`search`搜索论文，用精细化的`search`调用扩展最有希望的片段，然后通过`extract_section`提交最终答案，整个过程在5-7轮内完成。如果没有明确的策略，即使是像Sonnet 4.6这样的前沿模型也需要90秒才能生成一个轨迹，而在给定策略时只需30秒。 **冷启动SFT。** 虽然对于更大的模型可能不需要，但RLM框架的复杂性使得即使有好的提示和预定义工具，Qwen3.5-4B的pass@16分数也会为0。基于RLM的任务对于大多数小型模型来说超出了其能力范围。 在没有SFT阶段进行训练时，我们观察到了原始RLM论文中提到的许多相同失败模式，包括使用错误的语法提交最终答案，以及过多的推理和工具调用导致模型上下文窗口在轮次间膨胀。 对于SFT阶段本身，我们使用与原始RLM论文相同的方法。我们使用Qwen3.5-397B-A17B在相同的证据选择任务上生成教师RLM轨迹，并过滤掉包含REPL错误或检索到的证据片段与真实证据片段之间的字符跨度F1分数为零的轨迹。我们的SFT数据集是RL训练数据集中一个小的、预留的部分，包含几十个示例。在初始运行中，我们发现对整个RL训练数据集（尽管它不大）上的轨迹进行SFT导致了熵坍缩和随后的不稳定。 **逐步训练。** RLM框架的另一个细微之处是连续轮次不共享前缀。用户提示不是持久聊天历史的一部分，而是在每一轮之前附加。这样做的目的是确保原始用户查询不会埋藏在模型上下文窗口的深处，并且RLM会记住其任务。 由于每轮的用户提示是重写而不是累积的，因此不能将整个轨迹用作单个训练样本。每一轮必须是一个单独的训练样本，因此一个N轮的轨迹产生N个样本。对于优势计算，只有对应轨迹最后一步的轨迹才包含在GRPO组中，并且优势值广播到先前轮次的轨迹。更多关于逐步强化学习训练的信息，请参见此处。 **LLM评估器。** 我们使用基于规则的LLM评估器进行奖励分配。我们最初尝试了可验证的奖励，如所选片段的F1分数，但这被证明噪声很大。像“哪种方法在X基线上得分最高？”这样的问题可以用多个文本选择来回答，其中一些未包含在我们的标签中。我们在之前关于检索代理的工作中也遇到过类似问题。为了规避这一点，我们使用了基于规则的LLM评估器，它们被提供了原始查询、真实文本和预测文本。基于规则的评估器已被证明对奖励黑客攻击更具鲁棒性。Cameron Wolfe关于基于规则奖励的优秀概述可在此处找到。 有了这些考虑，强化学习微调在单篇论文任务上取得了显著改进，评估器分数从Qwen3.5-4B的约0.6跃升至0.8。 请注意，由于我们在SFT模型之上进行强化学习，训练从0.6的奖励开始，而不是从零开始，这表明了冷启动SFT阶段的重要性。 ### 多层RLM（带子调用） 现在让我们考虑多篇论文的情况，这需要真正的具有递归子调用的RLM。我们希望根级RLM识别哪些论文值得调度子RLM，而子RLM应该从它们分配的论文中提取相关段落。 分别训练根级和子级RLM的策略可能很诱人。然而，这将需要两组奖励信号和一个笨重的训练流程来训练两个策略。或者，我们可以使用冻结模型进行子调用，只训练根级RLM的模型，这是原始RLM论文采用的方法。然而，考虑到经过SFT的4B模型无法成功执行子任务，某种形式的在策略训练是必要的。我们设计了一个训练目标和强化学习训练设置，可以泛化到不同的RLM任务，并使单一策略能够有效地学习根级和子级RLM的角色。 **简单来说。** 使用标准GRPO计算根级RLM轨迹的优势。每个子轨迹继承其父根轨迹的优势，并且子轨迹对损失的贡献被平均（除以`k_g`，即该轨迹中的子数量），这样任何单个根轨迹都不会因为生成了更多子调用而被过度加权。 **或者更正式地说。** 对于批次中的每个查询`x`，采样一组`G`个根轨迹： \{y_g\}_{(g=1)}^G \sim \pi_{\theta_{\text{old}}}(\cdot \mid x) 每个根RLM轨迹`y_g`通过其`rlm_query`或`rlm_query_batched`调用确定性地诱导一组子提示`\phi(y_g) = \{x_{g,i}\}_{i=1}^{k_g}`，并且子轨迹从同一策略中采样： y_{g,i} \sim \pi_{\theta_{\text{old}}}(\cdot \mid x_{g,i}), \quad i=1,\dots,k_g **优势计算。** 只有根RLM轨迹获得可验证的奖励`r_g = R(x, y_g)`。在`G`个根轨迹上计算组相对优势： A_g = \dfrac{r_g - \text{mean}\! \left( \{r_{g'}\}_{g'=1}^G \right)}{\text{std}\! \left( \{r_{g'}\}_{g'=1}^G \right)} **优势继承。** 为简单起见，`y_g`的每个子轨迹继承`y_g`的优势： A_{g,i} := A_g \quad \text{for all } i=1,\dots,k_g **GRPO目标。** 令 \rho_\theta(y \mid x) = \dfrac{\pi_\theta(y \mid x)}{\pi_{\theta_{\text{old}}}(y \mid x)} 表示逐token的重要性比率（为清晰起见省略token索引）。 扩展到RLM树的GRPO目标为： \mathcal{J}(\theta) = \mathbb{E}_{x, \{y_g\}, \{y_{g,i}\}} \Bigg[ \dfrac{1}{G} \sum_{g=1}^G \Bigg( \underbrace{\mathcal{L}_g^{\text{root}}(\theta)}_{\text{根轨迹}} + {\color{red}\dfrac{1}{k_g} \sum_{i=1}^{k_g} \underbrace{\mathcal{L}_{g,i}^{\text{child}}(\theta)}_{\text{子轨迹}}} \Bigg) \Bigg] - \beta D_{\text{KL}}[\pi_{\theta} \| \pi_{\text{ref}}]