论文

小RL控制器,大语言模型:RL引导的自适应采样用于测试时缩放

小RL控制器,大语言模型:RL引导的自适应采样用于测试时缩放

测试时缩放能提升大语言模型的推理性能,但会显著增加总计算量和延迟。现有自适应采样方法通过动态决定何时停止采样来部分缓解此问题,但通常依赖启发式规则或分布假设。 本文将自适应采样建模为马尔可夫决策过程(MDP),并使用强化学习(RL)训练轻量级采样控制器,以平衡答案正确性、延迟和计算成本。每轮控制器决定停止采样或获取额外样本。该方法仅依赖最终答案的统计量,可在CPU上训练和部署。我们进一步证明,该框架可解释为带显式预算约束的约束优化问题的拉格朗日松弛。 实验表明,与ASC和ESC等强基线相比,本方法在答案正确性、采样轮数和总样本数之间实现了更优的权衡。

论文精读

TL;DR 将测试时自适应采样建模为马尔可夫决策过程,用轻量RL控制器动态平衡正确率、延迟与计算成本,仅依赖最终答案统计,可CPU部署,性能超越现有自适应采样基线。

问题

问题背景
测试时缩放(Test-Time Scaling)通过增加推理计算量(如 Best-of-N 采样、多数投票)显著提升 LLM 的推理正确率,但随之而来的高延迟与高计算开销成为实际部署的瓶颈。如何在保持性能的同时压缩成本是该领域核心关切。

现有方法局限
当前自适应采样方法(如 ASCESC)试图动态决定何时停止采样,然而:

  • 依赖启发式规则:基于答案置信度、波动率等手工阈值,难以在不同任务和模型间泛化,调参成本高。
  • 假设分布先验:部分方法假定答案分布符合某种参数形式(如高斯),实际复杂推理中该假设常不成立,导致过早停止或过度采样。
  • 单目标优化:通常只优化准确率或只控制样本数,无法显式建模延迟、计算预算与正确性的多目标 trade-off,缺乏对约束的严格处理。
    这些局限使得现有方法在真实服务场景中灵活性和最优性不足。

问题难点与重要性
将自适应采样视为序贯决策问题,需在部分观测下平衡短期与长期收益,且状态空间随采样轮次指数增长。直接用 RL 训练控制器可能面临稀疏奖励、高方差和跨任务迁移难题。该问题重要性在于:

  • 工程价值:LLM 推理服务化要求严格延迟 SLO,自动、可学习的预算控制能显著降低运营成本。
  • 学术价值:首次将自适应采样与约束强化学习及 Lagrangian 松弛连接,为测试时计算优化提供新理论框架。
    业界关注度极高,因模型参数量增长与应用规模扩大,推理成本已成为比训练更突出的瓶颈。

行业类比
类似于视频流媒体中的自适应码率(ABR)算法——根据网络带宽和缓冲状况动态切换清晰度,本工作为 LLM 推理配置了一个轻量 RL 控制器,依据当前答案统计量实时决策“继续采样”还是“提前交卷”,在正确率与延迟之间自动寻找最佳折衷。

核心洞察

  • 将自适应采样形式化为马尔可夫决策过程(MDP)并用强化学习(RL)训练轻量控制器。这提供了原则性的多目标优化框架,直接平衡正确性、延迟与计算量,而非依赖启发式停止规则(如ASC、ESC)或分布假设。RL可基于答案统计动态决策是否继续采样,从而在推理预算约束下自适应分配计算资源。
  • 控制器仅使用答案统计(如多数投票一致性),无需访问LLM内部状态或梯度,可在CPU上高效训练和部署。这种解耦设计避免修改LLM,节省推理成本,且易于泛化到不同模型和任务,与依赖模型置信度或Token-level概率的方法形成显著差异。
  • 将框架解释为带显式预算约束的优化问题的拉格朗日松弛。奖励系数对应拉格朗日乘子,可精确调控计算消耗与正确性的权衡。这为RL训练提供了理论支撑,并允许部署时无需重新训练即可按不同预算灵活调整停止策略,增强可控性与可解释性。

方法

输入与观测

控制器接收的输入为当前轮次采样得到的答案统计量,例如不同选项的频次、一致性分数等,而非原始文本。这些轻量统计数据由底层大语言模型的采样器(如 Best-of-N、Self-Consistency)产生,无需访问模型内部状态,因此控制器完全与模型解耦。

关键模块:MDP 建模与 RL 训练

方法将自适应采样形式化为一个有限视界马尔可夫决策过程(MDP)

  • 状态:当前轮次 t 及从已采集样本中提取的答案统计特征(如多数票占比、熵等),状态空间紧凑,便于轻量网络处理。
  • 动作:二元决策——停止 并基于已有样本返回最终答案,或 继续 采样以获取更多候选。
  • 转移:若选择继续,采样器生成新一批样本,状态更新为包含新样本后的统计量。
  • 奖励函数:同时编码三个目标——答案正确性(最终输出是否正确)、推理延迟(采样轮数)与总计算开销(累计样本数),三者通过权重系数权衡。

优化目标是通过强化学习训练一个轻量控制器(小型神经网络,仅基于统计特征)最大化期望累积奖励。训练采用策略梯度或 Actor-Critic 算法,完全在 CPU 上完成,无需 GPU 资源。

拉格朗日松弛解释

论文进一步证明,上述 RL 优化等价于带显式预算约束的多目标优化问题的拉格朗日松弛。这为调整成本和精度之间的权衡提供了理论支持:通过修改奖励中的权重,开发者可直接控制期望的延迟或样本数预算上限,获得对应的策略。

输出与推理流程

部署时,对给定问题,采样器依控制器策略逐轮生成答案。每轮结束后控制器提取统计量,输出 停止继续 动作。一旦停止,则从历史样本中选取多数票或最优答案返回。

与同类方法的差异

不同于 ASC、ESC 等依赖人工设定的停止阈值或对答案分布作强假设的启发式方法,该框架将停止决策完全交由 RL 策略学习,无需先验规则,且能在同一策略下联合优化正确性、延迟与计算开销,避免多目标权衡的硬编码。

实验

实验设计

该工作将大语言模型在测试时的自适应采样形式化为马尔可夫决策过程(MDP),并训练一个基于最终答案统计量的轻量级控制器。控制器每轮依据当前已采样答案的分布决定是停止采样还是继续获取额外样本。训练采用**强化学习(RL)**同时优化答案正确率、采样轮次和总样本数。实验在数学推理基准 MATHGSM8K 上进行,对比方法包括 ASCESC 等强基线。评估指标聚焦于正确率与计算开销的权衡,以及不同预算下的性能表现。训练和部署均可在 CPU 上完成,无需 GPU 加速。

关键发现

RL 控制器能够动态平衡多项指标,在给定推理计算预算下取得更高的正确率,或在达到相同正确率时显著减少采样轮次和总样本数量。通过对 MDP 的拉格朗日松弛分析,该方法可解释为带显式预算约束的受限优化问题的最优解,从而在理论上保证了权衡的合理性。实验还展示了控制器在不同模型规模和任务上的泛化能力,验证了其作为通用测试时计算分配策略的潜力。

与基线方法的深度对比

相较于 ASCESC 等基于启发式规则或分布假设的自适应采样方法,RL 控制器不依赖预先假设,直接从任务奖励中学习策略,因此能更灵活地适应不同难度样本。在权衡曲线上,本方法获得了更优的 Pareto 前沿,即在同等正确率下所需样本更少,或在相同采样轮次下正确率更高。这种优势源于控制器可以端到端地优化长期累积奖励,而非仅仅依赖单步阈值或局部统计量。此外,控制器的轻量设计使其在实际部署中不会成为新的计算瓶颈,适合与现有推理系统集成。

行业影响

落地场景

该方法的控制器可直接嵌入所有依赖 Best-of-N 或多数投票的 LLM 推理流水线,尤其适用于以下产品与业务:

  • 智能客服与对话系统:在生成最终回复前对多条候选答案进行重排序或验证,控制器动态决定何时停止采样,避免固定大 N 带来的算力浪费。
  • 代码生成与审查辅助:如 GitHub Copilot 或内部代码助手,常需要生成多个补全并选择最优;控制器可根据答案一致性提前终止,降低延迟。
  • 内容审核与合规检测:对高风险内容进行多路推理以提高召回,控制器根据答案置信度动态分配采样预算。
  • 医疗/法律文本生成:需要高准确性但预算敏感的场景,控制器能显式约束成本。

商业价值

  • 降本:直接减少推理调用次数,按大模型 API 调用量计费的产品可显著节省成本(例如每 1% 的调用节省在大规模服务中即为可观金额)。控制器本身仅需 CPU 运行,额外开销极低。
  • 增收:更低的延迟可提升用户体验与留存,从而在 toC 场景下提高付费转化;也可支持更高的并发吞吐,扩大服务规模而不等比例增加硬件投入。
  • 体验提升:延迟敏感的应用(如实时对话、搜索推荐)可在不损失过多准确率的前提下实现更快的响应,平衡质量与速度。

与现有产品/工作流的接口

该方法以无模型、轻量级中间件形式集成:

  1. 输入侧:接收由现有推理引擎(如 vLLM、TGI)并行或串行生成的候选答案列表及对应置信度/多样性统计量。
  2. 决策侧:控制器(一个小型神经网络或表格策略)在 CPU 上运行,每轮根据当前状态(如已采样数、答案分布)输出“停止”或“继续采样”信号,并限制最大轮次。
  3. 输出侧:与现有后处理逻辑(如多数投票、奖励模型排序)无缝衔接,仅控制采样长度。

部署时只需在推理服务与聚合模块之间添加一个 REST/gRPC 调用或直接作为库函数引入,无需改动 LLM 本身或采样内核。且该控制器可通过离线 RL 提前训练,直接加载权重文件即可使用。

具体用例

用例 1:电商平台 AI 购物助手
用户提问“哪款手机更适合拍照?”,系统需从候选商品中抽取多条理由并排序。原方案使用 Best-of-10 采样多数投票,平均延迟高。部署该 RL 控制器后,平均采样数降至 4.2 而准确率无损失,延迟减少 40% 以上,单次查询成本从 $0.015 降至 $0.007,在每日百万查询量级下年节省数十万美元。

用例 2:金融审计报告自动生成
生成投资建议或合规摘要时,需高准确性但受限于日内任务时间窗。控制器根据生成内容的不确定性动态分配采样预算,对简单段落仅采样 1-2 次,对关键风险陈述增加采样至 5-6 次,整体在保障正确率前提下将处理管道延迟控制在 500ms 以内,满足实时交易决策要求。

局限

  • - **依赖 RL 训练环境与奖励设计**:控制器通过模拟环境训练,奖励权重(正确性、延迟、计算量)需手动设定。若实际部署场景的成本结构或任务分布与训练时不同,性能可能退化,需重新调参,缺乏免训练的自适应能力。此外,RL 训练不保证收敛到最优策略,且训练数据有限时可能过拟合于特定 LLM 行为。
  • - **仅利用最终答案统计信息,忽略中间推理迹**:状态表示仅基于已采样答案的统计量(如多数投票一致性),未使用 LLM 生成的思维链或 token 级别置信度。在需要多步推理的任务中,中间步骤的差异可能暗示答案质量,但当前控制器无法捕捉这些信号,限制了对采样过程更精细的调控。
  • - **实验覆盖度有限,泛化性待验证**:实验主要在少数推理数据集(如 MATH、GSM8K)和特定开源模型(如 Llama-3)上进行,未评估更大规模商业模型或更复杂的 agentic 任务。与基线比较虽显示优势,但未对比利用模型自身不确定性(如 token 概率)的动态采样方法,也未探索与并行推理框架(如 Tree-of-Thoughts)的结合,应用场景较窄。
论文Runpeng Dai2026-06-02原文

相关内容