论文

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

Small RL Controller, Large Language Model: RL-Guided Adaptive Sampling for Test-Time Scaling

模型推理测试时计算扩展

摘要

测试时间缩放提高了 大语言模型 的推理性能,但会在总计算和延迟方面产生大量成本。现有的自适应采样方法通过动态决定何时停止采样来部分缓解这个问题,但它们通常依赖于启发式规则或依赖于分布假设。在这项工作中,我们将自适应采样制定为马尔可夫决策过程(MDP)。我们使用强化学习(RL)训练轻量级采样控制器,以共同平衡答案正确性、延迟和计算成本。在每一轮中,控制器决定停止采样或获取额外的样本。我们的方法是轻量级的,仅依赖于最终答案的统计,并且可以在CPU上进行训练和部署。我们进一步表明,所得框架允许将其解释为具有显式预算约束的约束优化问题的拉格朗日松弛。针对 ASC 和 ESC 等强基线的实验表明,我们的方法在答案正确性、采样轮次和所需样本总数之间实现了改进的权衡。