复用 FLOPs:以极异策略前缀为条件扩展困难问题的强化学习

Reuse your FLOPs: Scaling RL on Hard Problems by Conditioning on Very Off-Policy Prefixes

arXiv: 2601.18795v1

论文信息

标题: Reuse your FLOPs: Scaling RL on Hard Problems by Conditioning on Very Off-Policy Prefixes

作者: Amrith Setlur, Zijian Wang, Andrew Cohen, et al.

发布日期: 2026-01-26

arXiv ID: 2601.18795v1

PDF 链接: 下载 PDF

3 分钟速览

  • 研究问题:这篇论文要解决在困难问题上强化学习训练大型语言模型时,正确轨迹稀缺导致学习停滞的问题。
  • 核心方法:通过将过去的正确推理轨迹前缀拼接到原问题上,运行在线策略强化学习,从而绕开离线策略带来的优化不稳定性。
  • 关键结果:在训练困难问题上,PrefixRL 达到同一训练奖励所需的计算量仅为最强基线(SFT+RL)的 1/2,且最终准确率提升超过 45%(见论文图 2)。
  • 主要局限:方法依赖可用的正确离线轨迹;前缀长度若大量集中在极长片段,回传到原始问题的性能提升会变慢;理论分析基于自然策略梯度,与实际使用的 REINFORCE 有所不同。
  • 适合读者:从事大语言模型强化学习、数学或代码推理后训练、以及关注计算效率与数据重用的研究者与工程师。

论文背景和研究动机

当前,强化学习(RL)已成为提升大语言模型(LLM)推理能力的主流方法,尤其在数学和编程领域。通常,这些方法是在线策略的:从当前模型多次采样推理轨迹,并根据正确与否进行更新。然而,当问题难度很高,模型几乎无法采样到正确轨迹(如 pass@2k ≈ 0)时,模型将得不到有效学习信号,训练陷入停滞。

在实践中,我们其实并非头一次面对这些困难问题。先前的推理或 RL 训练中往往已经生成过大量离线轨迹数据,其中虽然稀有,但确实包含部分正确答案。自然的问题是:如何重用这些已成规模的计算成果,引导在线策略进入更高奖励的状态,从而加速学习。

直接的思路是把离线轨迹当作监督目标,例如先用正确轨迹进行监督微调,再进行在线 RL;或者在 RL 中通过重要性采样直接使用这些轨迹。但这两种方式都会导致问题:SFT 可能引起生成多样性崩溃,伤害 RL 的探索能力;离线轨迹在在线策略下概率极低,重要性采样会导致梯度方差过大,训练不稳定(见图 3)。

因此,论文提出一种新思路:不把离线轨迹作为监督目标,而是将其作为「条件」来引导在线 RL。这就是 PrefixRL 的出发点。

核心方法和技术细节

PrefixRL 包含三个关键步骤。首先,从已有的正确离线轨迹中截取前缀(即轨迹的前 h 个 token),并将其拼接到原问题之后,形成带有引导信息的「前缀问题」。这些前缀通常揭示了解题的高层策略,大幅提升了模型在该前缀条件下给出正确答案的概率(见图 4)。

第二步,训练时同时使用原始问题和前缀问题,但在前缀部分的 token 上不计算策略梯度(梯度掩码),只对模型自行生成的后续部分进行 RL 更新。由于前缀来自正确轨迹,它们将模型「重置」到更可能走向成功的中继状态,减少了梯度方差,增强了学习信号。

第三步依赖一个经验发现:仅在前缀问题上训练就能有效提升模型在原始问题上的表现,这一现象被论文称为反传泛化。也就是说,即使训练和测试的输入分布不同(有前缀 vs. 无前缀),模型的解题能力仍然得到传递。

在理论层面,论文证明了两点(见定理 3.2、3.3):

  1. 目标一致性:若离线轨迹正确且可实现,则最大化 PrefixRL 目标的策略也同时最大化标准 RL 目标。
  2. 样本效率提升:在自然策略梯度设定下,PrefixRL 仅需更少样本即可降低次优性差距,且该差距受初始策略与生成离线数据策略间的 KL 散度控制,不会随迭代累积。

当离线轨迹来自拒绝采样时,该 KL 散度仅随采样预算对数增长,意味着额外开销有限。

创新点和贡献

PrefixRL 的核心创新在于改变了 “如何使用离线数据” 的范式:传统方法视离线数据为待模仿的目标,而 PrefixRL 视之为状态重置的条件。这与监督或重要性采样形成了本质区别——它避免了因模仿太强而导致的熵塌陷以及因概率太低而导致的训练崩溃。

另一个贡献是揭示并系统分析了反传泛化现象。论文发现,相比于直接模仿前缀中的策略,PrefixRL 甚至可以学会与给定前缀不同的、更优的解题策略(如图 6 所示,从 “Erdős–Gallai 定理” 策略转向了 “Dirichlet 定理” 策略)。这种跨状态分布的泛化只能通过函数近似(LLM 共享的内部表征)来解释,无法在表格型 RL 中实现。

此外,PrefixRL 实现了自我提升:离线数据通过从自身基模型进行大规模拒绝采样产生的,形成了 “采样 → 条件 RL → 模型提升” 的闭环,无需依赖更强的教师模型或人工标注数据。

实验结果分析

实验基于 Distilled Llama-3.1-8B-instruct 和 Qwen3-4B-instruct 两个模型,在 1000 个极端困难问题(基模型 pass@512 ≈ 0)上训练。

计算效率匹配下(包含拒绝采样的前期成本),PrefixRL 的训练奖励达到最强基线的 2 倍,最终准确率绝对提升超过 45%(Llama 上,见图 2、图 9)。在 AIME ‘25 等留出基准上,PrefixRL 的 pass@1 从 38.2 提升至 61.3,并在 k 增大时优势进一步拉大,说明模型的后缀生成分布整体质量更高(图 10)。

从训练动态看,PrefixRL 更均匀地提升了困难问题的 pass@1,而非仅集中攻克少数问题(图 13a);同时维持了更高的 token 级熵,避免了 SFT+RL 中的探索崩塌(图 13 左);并拥有更高的梯度范数、更低的梯度方差,即更高的信噪比(图 14)。

值得注意的是,即使离线轨迹来自不同模型族(如用 Qwen 的轨迹训练 Llama),PrefixRL 仍然有效且效果接近(图 12)。这说明方法对离线数据源具有较好的鲁棒性。

实践建议

PrefixRL 为 LLM 的推理后训练提供了一项高性价比的实践方案:

  1. “废物利用” 旧计算:已经运行过的大量推理性采样或历史 RL 运行所积累的离线轨迹,可通过 PrefixRL 成为后续训练的强引导信号,实现数据与算力的复用。
  2. 适合困难问题密集的场景:当模型在多数问题上采不到正奖励时,可以通过构建前缀版的同类问题,为在线 RL 创造更稠密的学习信号。
  3. 前缀选取策略:论文显示,前缀长度介于轨迹长度的 40% 到 80% 之间并随机采样,并结合少量极长前缀,可使反传泛化更快发生。实践中,建议先统计基模型在各前缀比例下的准确率,选择能大幅提升准确率的最短前缀区间(即关键策略揭示点),以平衡计算开销和训练效益。
  4. 避免 SFT 陷阱:如果同时有离线数据和在线 RL 需求,PrefixRL 可以替代 SFT 初始热身阶段,保持模型更高的输出熵和探索能力,从而避免 SFT 导致过拟合和后续 RL 收益递减。

综上,PrefixRL 提供了一种更稳定、更高效的利用离线推理数据的后训练范式,尤其适合在困难问题上追求更大增益、且已持有较多历史推理数据的团队。