arXiv:2503.01422cs.CLcs.AI2025-03NeurIPS被引 63

不生成全部样本也能高效提升模型推理效果

Sampling-Efficient Test-Time Scaling: Self-Estimating the Best-of-N Sampling in Early Decoding

  • 通过早期解码状态判断最优路径,提前剪枝低质量分支
  • 内存占用降低80%以上,延迟减少50%,计算成本节省70%-80%
  • 无需奖励模型,适合资源受限的部署场景

测试时扩展通过在推理阶段分配更多计算资源来提升大语言模型性能。基于采样的最佳N选一(Best-of-N, BoN)方法通过并行扩展搜索空间,从模型分布中寻找更优解。然而其成本-性能权衡仍不明确。现有挑战主要有:(1) 生成N个完整样本需大量GPU内存,限制了有限资源下的推理能力;(2) 奖励模型带来额外内存与延迟开销,且训练强奖励模型存在数据成本。尽管已有研究尝试优化效率,但尚未同时解决上述问题。为此,本文提出自截断最佳N选一(Self-Truncation Best-of-N, ST-BoN),一种无需生成全部N个样本且无需奖励模型的解码方法。该方法利用模型内部状态在早期采样中的一致性,识别最具潜力的路径并剪枝次优分支。实验显示,ST-BoN将动态GPU内存使用降低超80%,推理延迟减少50%;在相同成本下,性能达到全量BoN水平,且可将准确率提升3-4个百分点。

原文摘要 · Abstract (English)

Test-time scaling enhances large language model performance by allocating additional compute resources during inference. Best-of-N (BoN) sampling serves as a common sampling-based scaling technique, broadening the search space in parallel to find better solutions from the model distribution. However, its cost-performance trade-off is still underexplored. Two main challenges limit the efficiency of BoN sampling: (1) Generating N full samples consumes substantial GPU memory, reducing inference capacity under limited resources. (2) Reward models add extra memory and latency overhead, and training strong reward models introduces potential training data costs. Although some studies have explored efficiency improvements, none have addressed both challenges at once. To address this gap, we propose Self-Truncation Best-of-N (ST-BoN), a decoding method that avoids fully generating all N samples and eliminates the need for reward models. It leverages early sampling consistency in the model's internal states to identify the most promising path and truncate suboptimal ones. In terms of cost, ST-BoN reduces dynamic GPU memory usage by over 80% and inference latency by 50%. In terms of cost-performance trade-off, ST-BoN achieves the same performance as Full-BoN while saving computational cost by 70%-80%, and under the same cost, it can improve accuracy by 3-4 points.

推理优化采样策略大模型

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。