小型 RL 控制器与大语言模型:面向测试时扩展的强化学习引导自适应采样 Small RL Controller, Large Language Model: RL-Guided Adaptive Sampling for Test-Time Scaling
把自适应采样建模为MDP,用四层MLP强化学习控制器平衡答案精度、延迟与算力开销。
前置知识
测试时计算扩展 (Test-Time Scaling)
指在推理阶段而非训练阶段投入更多算力来提升模型表现,典型做法是对同一问题多次采样回答再用多数投票(self-consistency)或 best-of-N 选最优,无需重训模型但代价是算力与延迟大幅上升。
本文要解决的就是这类方法成本过高的痛点,理解其成本-性能权衡是读懂全文的基础。
自适应采样 (Adaptive Sampling)
在多次采样过程中根据已有答案动态决定继续采样或提前停止的策略,如 ASC 靠后验概率阈值停止、ESC 靠批内一致性停止,目的是在精度损失可控前提下减少总采样数与延迟。
本文方法是这类策略的全新设计,需先了解现有方案及其启发式规则的局限才能理解改进点。
马尔可夫决策过程 (MDP) 与 PPO
MDP 用状态、动作、转移、奖励四元组建模序贯决策;PPO(近端策略优化)是稳定的 on-policy 策略梯度强化学习算法,通过裁剪重要性采样比率更新策略网络,本文用它训练采样控制器。
本文把自适应采样建模为 MDP 并用 PPO 求解,这是方法的核心形式化基础。
拉格朗日松弛与约束优化
把带约束的优化问题转化为无约束问题:将约束以惩罚项(乘以对偶变量/拉格朗日乘子)加入目标函数。本文证明其 RL 目标等价于'延迟与算力预算约束下最大化精度'问题的拉格朗日松弛。
Proposition 1 用这一视角为惩罚权重设计提供理论依据,是新颖性的关键。
多数投票与答案熵
多数投票取采样答案池中频次最高的答案作为最终预测;答案熵衡量答案分布的不确定性(熵高表示答案分散、难收敛)。本文状态同时用 top-K 频次向量与熵来刻画采样收敛程度。
理解状态如何从答案统计构造、以及停止判断依据,才能看懂控制器决策逻辑。
研究动机
测试时计算扩展(如 self-consistency、best-of-N、tree-of-thoughts)通过在推理时分配更多算力来提升 LLM 的最终答案质量,但代价昂贵——既推高总计算量又增加延迟。现有自适应采样方法(如 ASC、ESC)试图通过动态决定何时停止采样来缓解,却存在三类关键局限:(i) 它们严重依赖人工设计的启发式规则——ASC 用 0.95 的 p 值阈值、ESC 用固定 chunk size 5——或分布假设,而非显式推导一个在性能-成本权衡上最优的策略;(ii) 一些方法需要额外辅助信号,例如模型内部置信度、隐藏状态或问题难度,这些在闭源模型或标准推理流水线下往往不可用;(iii) 某些技术高度侵入性,会打断模型自然推理过程甚至要求重训 LLM。实际表现上,ASC 虽能把总样本数从 32 降到约 15-27,却因逐个顺序采样导致轮次暴涨十倍以上,延迟极差;ESC 用粗粒度批处理压轮次,却要以更多总样本或牺牲精度为代价。
本文的目标是本文的目标是设计一个既轻量又有原则的自适应采样框架,能够显式优化性能-成本权衡,同时灵活兼顾延迟(采样轮次)和计算(总样本数)这两个对实际部署至关重要的目标。框架必须满足三个硬约束:第一,完全非侵入,不干扰 LLM 的自然推理过程,也不需要对底层模型做任何额外训练;第二,不依赖任何辅助信号——不需要模型置信度、隐藏状态或问题难度估计,只使用每轮已采样答案的统计量;第三,控制器本身要足够轻量,能够在 CPU 上训练和部署,从而可即插即用地接入任何标准推理流水线,包括闭源 API 模型。
与已有工作不同的是,本文的独特切入角度是把自适应采样严格建模为有限时域马尔可夫决策过程(MDP),用一个四层 MLP 作为控制器,通过 PPO 强化学习显式学习最优的停止/续采样策略,让 ASC 与 ESC 退化为该 MDP 框架下的固定规则策略特例。状态、动作、奖励全部只依赖已采样答案池的纯统计量(top-K 频次向量 $V_t$、样本数 $|D_t|$、熵 $H(V_t)$),完全不触碰查询语义、模型置信度或隐藏状态。更关键的是,作者用 Proposition 1 证明该 RL 目标可解释为'在延迟预算 $C_{lat}$ 和计算预算 $C_{comp}$ 约束下最大化精度'这一约束优化问题的拉格朗日松弛,惩罚权重 $\lambda_{lat}, \lambda_{comp}$ 正是对应的对偶变量——这把工程化的 RL 设计与有理论根基的约束优化视角优雅地统一起来。
核心方法
直觉上,每当 LLM 采样出一批候选答案后,我们其实可以从'答案分布是否已经收敛'来判断是否值得继续花钱采样——若多数答案已高度集中,再采几条边际收益很小;若答案仍分散,则应追加样本。作者把这种判断交给一个极小的神经网络来做。技术路线如下:给定查询 $x$ 和最大采样预算 $N=32$,在每一轮 $t$,控制器观察当前答案池 $D_t$ 的统计特征构成状态 $s_t$,据此从动作集 $A=\{0,1,2,4\}$ 选择动作——$0$ 表示停止并返回多数投票答案,$k>0$ 表示再并行生成 $k$ 个候选;环境更新答案池、重算状态进入下一轮,直到主动停止或触及预算上限。整个过程用 PPO 训练,最大化期望累计奖励。
核心创新点在于用'学习到的最优策略'替代'固定的人工规则策略'。在作者构建的 MDP 里,ASC(每次采 1 个、靠后验阈值停)和 ESC(固定 batch、靠批内一致性停)本质上都是在同一状态-动作空间上运行的固定规则策略;而 RL-Guided 用 PPO 动态学习一个显式最大化期望累计奖励的策略。它与已有方法的本质区别有三:第一,纯统计驱动,状态只用 top-K 频次向量 $V_t$、样本数 $|D_t|$ 和熵 $H(V_t)$ 拼成 7 维向量,不引入任何辅助特征;第二,多目标联合优化,奖励同时惩罚延迟(轮次)和计算(样本数);第三,完全非侵入,不修改 LLM 推理过程。
方法步骤详情
完整流程分五步。(1) 状态构建:收集答案池 $D_t$,统计出现频次最高的前 $K=5$ 个答案计数 $V_t$,加上总样本数 $|D_t|$ 与熵 $H(V_t)$,拼成 7 维状态 $s_t$。(2) 动作选择:四层 MLP 策略网络读入 $s_t$ 输出分布,从 $A=\{0,1,2,4\}$ 采样动作 $a_t$。(3) 环境转移:若 $a_t=0$ 则终止,用 $\text{MajorityVote}(D_t)$ 输出预测 $\hat{y}_t$;若 $a_t=k>0$ 则并行生成 $k$ 个候选加入 $D_t$ 并重算 $s_{t+1}$,若 $n_t+a_t\geq N$ 则强制终止。(4) 奖励:续采样扣分 $r^{step}_t=-\lambda_{lat}-\lambda_{comp}a_t$;终止时按多数投票是否等于 $y^*$ 给 $r^{final}_t=\pm1$。(5) 训练:用 PPO 在 DAPO 随机采样的 200 道题上优化 $J(\pi_\theta)=\mathbb{E}\sum_t r_t$。
技术新颖性
技术新颖性体现在三方面。其一,首次把自适应采样严格建模为可被任意 RL 算法(PPO/DQN/策略梯度)求解的 MDP,使 ASC/ESC 退化为该框架下固定规则策略的特例,统一了这一类问题。其二,Proposition 1 证明目标是约束优化 $\max J_{acc}(\pi_\theta)\ \text{s.t.}\ J_{comp}\leq C_{comp}, J_{lat}\leq C_{lat}$ 的拉格朗日松弛,惩罚权重充当对偶变量——为权重选择提供理论锚点,并自然指向未来用约束 RL 直接按硬预算优化的方向。其三,奖励设计刻意与问题特定信号解耦:用 Running Majority(继续采到 $N=32$ 时的多数投票)作为目标 $y^*$ 而非 ground-truth,与纯统计状态保持一致;Table 2 消融证明这种解耦反而带来最佳精度-效率权衡与跨模型泛化,用真实标签反而使精度掉到 53.7%。
实验结果
核心发现按表逐一分析。Table 1 在三个数学推理基准(AIME24/AIME25/HMMT25)×四个采样器上对比。以 Qwen3-4B-Instruct 平均指标为例:SC@32 为 54.8% 精度/1.0 轮/32 样本,ASC 达 54.8%/15.3 轮/15.3 样本,ESC 为 54.4%/4.1 轮/20.3 样本,而 RL-guided 达 54.6%/2.8 轮/10.9 样本——精度持平 SC 同时总样本压缩 65.8%。总体相比 ASC,RL-guided 把轮次降约 3-4×、总样本降约 30%;相比 ESC 轮次降约 10%、总样本降约 35%。GPT-4.1-nano 上为 27.7%/6.9 轮/17.4 样本,显著优于 ASC(21.7 轮)与 ESC(26.7 样本)。Figure 2 两维度稳定占优;Figure 3 显示样本分配与答案熵正相关;Figure 4 证明 0.6B 训练的控制器可跨模型迁移;Table 2 表明 Running Majority 目标远胜 Real Label。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| 数学推理(AIME24/25 + HMMT25),Qwen3-4B-Instruct 采样器 | 平均精度 / 平均采样轮次 / 平均总样本数 | 54.6% / 2.8 轮 / 10.9 样本 (5 个随机种子均值) | ASC 54.8% / 15.3 轮 / 15.3 样本;ESC 54.4% / 4.1 轮 / 20.3 样本;SC@32 54.8% / 1.0 / 32 | 对比 ASC 轮次降约 3-4×、总样本降约 30%;对比 ESC 轮次降约 10%、总样本降约 35%,精度基本持平 SC |
| 跨模型迁移评测(GPT-4.1-nano 闭源模型) | 平均精度 / 采样轮次 / 总样本数 | 27.7% / 6.9 轮 / 17.4 样本(用 Qwen3-0.6B 训练的控制器直接迁移) | ASC 27.8% / 21.7 轮 / 21.7 样本;ESC 27.5% / 5.3 轮 / 26.7 样本;SC@32 27.8% / 1.0 / 32 | 精度持平前提下总样本比 ASC 少 45.5%、轮次少约 3 倍,且无需在目标模型上重新训练控制器 |
| 奖励信号消融(Qwen3-4B-Instruct) | 平均精度 / 轮次 / 总样本 | Running Majority 目标 54.6% / 2.8 / 10.9 | Real Label 53.7% / 4.1 / 17.0;Full Majority 54.6% / 2.8 / 11.7 | 纯统计目标比真实标签精度更高、总样本少约 36%,验证奖励与问题特定信号解耦的设计优越性 |
局限与改进
作者承认:当前状态表示刻意只用简单统计量,未来可纳入答案置信度、答案平均长度等更丰富信号;奖励设计目前用固定惩罚权重,未来可直接用真实成本(时间、金钱)作为惩罚。我自己进一步观察到几点局限:(1) 方法只在数学推理任务上验证,依赖从回答中抽取 \boxed{} 终答案来做频次统计与多数投票,对自由生成、代码、长文等无明确离散终答案的任务无法直接套用;(2) Running Majority 目标隐含'继续采到 N 会收敛'的假设,对真正极难、答案长期不收敛的题可能误导停止时机;(3) 惩罚权重 $\lambda_{lat}, \lambda_{comp}$ 是固定超参,作者虽用 Proposition 1 把它解读为对偶变量,却并未做动态对偶更新,只能扫超参得到不同 Pareto 点(Figure 2),无法严格满足给定硬预算约束;(4) 评测集中在 AIME/HMMT 竞赛数学,未覆盖工具调用、代码生成、多模态等场景。
独立分析的弱点
独立分析主要有五点弱点,各附改进方向。(1) 离散答案依赖:方法依赖 \boxed{} 等终答案做频次统计与多数投票,对开放生成、代码、长文任务不适用,可用 embedding 语义聚类替代严格答案匹配来构造状态。(2) 固定动作集 $A=\{0,1,2,4\}$:动作粒度需人工设定且固定,对不同预算/模型未必最优,可改为连续动作或自适应动作空间。(3) 静态权重 $\lambda$:只能通过扫权重得到不同权衡点,无法在给定硬预算下严格最优,可引入 Constrained Policy Optimization 类动态对偶更新(作者也提到)。(4) 评测域单一:仅 AIME/HMMT 竞赛数学,未覆盖多步工具调用、代码生成、多模态推理。(5) 训练数据依赖 DAPO 数学题,向非数学域的迁移尚未验证。
未来方向
作者明确提出的方向:纳入更丰富的状态信号(答案置信度、答案平均长度)、用真实时间/金钱成本替代固定惩罚作为奖励、用约束 RL 在给定硬预算下直接优化策略。基于成果可进一步延伸:(1) 把控制器做成条件策略 $\pi(a|s,w)$,输入偏好向量 $w$,训练单一统一策略覆盖整条 Pareto 前沿,省去扫多组权重的开销;(2) 把同一 MDP 框架推广到树搜索/beam search 的剪枝与扩展决策,用轻量控制器替代启发式搜索预算分配;(3) 让控制器同时决定采样温度、top-p 等推理超参,而不仅是样本数;(4) 与 step-wise verifier 或 confidence 估计结合,在保留统计驱动可迁移性的同时提升停止决策质量。
复现评估
复现友好度较高。代码已开源(github.com/RunpengDai/RL-Guided-Adaptive-Sampling)。基准 AIME24/25、HMMT25 均公开,训练集 DAPO 公开,采样器为开源 Qwen3 系列(0.6B/1.7B/4B)加闭源 GPT-4.1-nano。控制器仅四层 MLP,作者强调可在 CPU 上训练与部署,算力门槛很低——训练只需 200 题 × 32 样本的 rollout 数据。RL 采用标准 PPO。附录 A 给出完整实现细节,附录 C 给出 Proposition 1 的形式化证明。主要成本在于生成训练 rollout 的 LLM 推理,以及为得到 Figure 2 不同 Pareto 点而训练多组权重控制器。整体复现难度中等偏低,对具备基本 GPU 资源的研究组即可上手。
论文图表