跳转至

OptScale: Probabilistic Optimality for Inference-time Scaling

会议: AAAI 2026
arXiv: 2506.22376
代码: GitHub
领域: LLM推理优化
关键词: 推理时缩放, 并行采样, 概率最优, Best-of-N, 计算效率

一句话总结

提出概率最优框架 OptScale,通过建模验证器分数的概率分布推导出最优采样数量的理论下界,动态决定每个问题所需的最少采样次数,在保持推理准确率的同时大幅减少计算开销。

研究背景与动机

  1. 领域现状:推理时缩放(inference-time scaling)已成为提升 LLM 推理能力的主流技术,通过并行生成多个候选答案并选取最优,可显著提高数学推理等任务的性能。
  2. 现有痛点:现有方法面临严重的效率瓶颈——候选解数量与计算成本呈线性关系,token 消耗随采样数 N 线性增长,导致计算资源浪费。现有动态分配策略依赖启发式规则,缺乏数学基础。
  3. 核心矛盾:推理准确率要求尽可能多地采样,但计算预算有限——简单问题不需要那么多样本,困难问题可能无论怎么采样都无法解决,如何在理论上确定最优采样数量?
  4. 本文要解决什么? 建立推理时缩放的概率理论基础,推导最优采样数量的闭式解,并实现一个高效实用的自适应采样算法。
  5. 切入角度:将 Best-of-N 选择过程形式化为概率分布,通过极值分布理论推导最优采样数 N*。
  6. 核心idea一句话:验证器分数服从可估计的概率分布,利用极值理论可以推导出满足目标性能和置信度所需的最小采样数量。

方法详解

整体框架

给定输入问题 q,LLM 生成 N 个候选答案,验证器(如 PRM)对每个答案打分,选择最高分答案。OptScale 的核心是动态决定 N 值:先估计验证器分数的概率分布参数 (μ, σ),然后通过理论公式计算满足性能阈值和置信度的最小 N,达到 N 后立即停止采样。

关键设计

  1. 概率最优理论框架:假设验证器分数 {s_i} 是 i.i.d. 的随机变量,服从概率密度函数 f_S(s|θ,q)。对于 N 个独立采样,最大值 Y = max{s_1,...,s_N} 的 CDF 为 F_Y(s) = [F_S(s)]^N。要求 P(Y >= s_min) >= α,即至少有概率 α 能找到一个超过质量阈值 s_min 的答案。对不等式求解可得最优采样数的闭式下界:
  2. N* >= ceil(log(1-α) / log F_S(s_min))
  3. 这一公式优雅地将"需要多少样本"转化为验证器分数分布尾部概率的函数
  4. 当 F_S(s_min) 接近 1(即大多数样本质量很高),N* 很小,自动减少采样
  5. 当 F_S(s_min) 接近 0(即几乎不可能达到阈值),N* 趋向无穷,提示放弃
  6. 这是首次为并行推理时缩放提供有数学保证的闭式理论指导

  7. OptScale_t(可训练版本):将验证器分数建模为截断正态分布(约束在 [0,1]),核心参数为 (μ, σ)。利用离线数据训练两个 MLP 分别预测先验参数 μ_bar 和 σ_bar。推理时先用 MLP 给出先验预测,再通过 MAP 估计结合实际观测的验证器分数 D = {s_k} 动态更新参数。MAP 目标函数同时优化似然项(拟合观测)和先验项(靠近 MLP 预测),平衡了全局趋势和实例特异性。

  8. OptScale_0(免训练版本):完全不依赖预训练模型。对新问题采用均匀不确定性的启发式策略初始化 (μ_0, σ_0),随后通过 MLE 从观测分数中直接估计参数。随着更多验证器分数被观测到,参数估计不断精化。该变体无额外模型、无训练数据依赖,即插即用,适合快速部署。实验表明它在多数场景下甚至优于可训练版本。

  9. 自适应终止机制:两种变体最终都通过更新后的 (μ, σ) 代入 F_S(s_min) 计算 N。每生成一个新样本就更新参数并重新计算 N。一旦当前样本数 N >= N*,立即终止采样。这实现了三重自适应效果:

  10. 简单问题:分布参数估计显示高质量概率大,N* 很小,快速早停
  11. 中等难度问题:需要适量采样直到分布估计稳定
  12. 极困难问题:N* 达到预设上限时截断,避免无意义的过度采样

损失函数 / 训练策略

OptScale_t 训练阶段使用离线数据拟合 MLP 预测器,学习问题到分布参数的映射。推理阶段无需额外训练。OptScale_0 完全免训练,纯推理时运行。

实验关键数据

主实验

在 Deepseek-R1-Distill-Qwen-7B 上的结果(N=8):

方法 MATH-500 Acc. MATH-500 Toks(↓) GSM8K Acc. GSM8K Toks(↓) AIME24 Acc. AIME24 Toks(↓)
Best-of-N 94.8 22135 92.4 3582 70.0 79367
Self-Consistency 93.4 22135 90.1 3582 60.0 79367
OptScale_0 94.8 11354 92.4 1687 70.0 49505
OptScale_t 94.8 18236 92.4 3492 70.0 53855

N=64 时在 MATH-500 上:OptScale_0 仅用 110001 tokens(vs BoN 的 174693),准确率从 94.0% 提升到 94.6%。

消融实验

在 QwQ-32B(N=60)上的结果:

方法 MATH-500 Acc. Toks(↓) AMC23 Acc. Toks(↓)
Best-of-N 94.8 230402 97.5 420481
Self-Consistency 95.4 230402 92.5 420481
OptScale_0 95.8 107720 100.0 190633
OptScale_t 95.8 106412 100.0 202603

OptScale 在 AMC23 上达到了 100% 准确率,同时 token 消耗降低了约 54%。

在 Llama-3.1-8B(N=64)上,OptScale_0 将 GSM8K 的 token 从 14697 降到 5720(降 61%),准确率从 89.3 提升到 89.5。

关键发现

  • OptScale_0 在多数场景下比 OptScale_t 更高效,因为免训练版本在简单问题上能更快做出早停决策
  • 在 Llama-3.1-8B 这种较弱模型上,OptScale 同样有效,token 节省比例甚至更加显著
  • OptScale 对简单问题自动少采样,对困难问题主动控制计算上限,避免无效消耗
  • MR-Thinking(自我反思方法)在多个基准上准确率反而下降,说明增加每样本推理深度不如优化采样策略

亮点与洞察

  • 首次为推理时并行缩放建立了严格的概率理论基础,推导出了闭式最优采样数下界
  • 方法显著实用:OptScale_0 完全免训练,即插即用,部署成本为零
  • 在保持甚至提升准确率的情况下,token 消耗可降低 37%~61%
  • 理论框架具有普适性,不依赖特定 LLM 或验证器,可与任何 Best-of-N 流程组合

局限性 / 可改进方向

  • 理论假设验证器分数是 i.i.d.,实际上同一问题的不同采样可能存在相关性(如 temperature 较低时)
  • 截断正态分布假设可能对某些任务不够精确,更灵活的分布族(混合模型等)可能带来提升
  • 主要在数学推理任务上验证,代码生成、逻辑推理等场景尚需测试
  • OptScale_t 需要离线数据训练预测器,对全新任务的泛化能力有待评估
  • 与树搜索(MCTS)等顺序推理方法的结合未探索

相关工作与启发

  • vs Self-Consistency: SC 使用多数投票选答案,固定 N 无法自适应;OptScale 动态调整 N 且用 verifier 选优
  • vs SEAL: SEAL 动态分配计算但依赖启发式不确定性指标;OptScale 从概率理论出发有闭式解
  • vs MR-Thinking: 自我反思型方法增加了每个样本的 token 开销,实验显示准确率反而可能下降;OptScale 通过减少无效样本数节省更多

评分

  • 新颖性: ⭐⭐⭐⭐ 首次将极值理论引入推理时缩放,提供了理论下界,但概率建模本身的技术创新有限
  • 实验充分度: ⭐⭐⭐⭐ 覆盖多个 benchmark 和多个模型(7B/8B/32B),但缺乏非数学推理任务的验证
  • 写作质量: ⭐⭐⭐⭐⭐ 理论推导清晰完整,从框架到实现逻辑自洽
  • 价值: ⭐⭐⭐⭐ 理论+实践双重贡献,免训练版本实用价值高,但 i.i.d. 假设限制了理论的严格适用范围