Direct Pass@$k$ Optimization for Thought Pruning: Balancing Quality and Diversity in Test-Time Reasoning

LLM CoT ML ALG
2026年10月22日
测试时缩放(test-time scaling)已成为推理模型的一项关键推理策略,其核心思想是让模型生成大量候选解,以期其中至少有一个是正确的。然而,若对所有候选解均执行至最终完成,计算开销将十分巨大;因此,近期研究提出一种策略:仅生成每个候选解的较短前缀(prefix),便对候选池进行剪枝,后续计算资源仅分配给通过剪枝筛选出的幸存候选解。现有剪枝策略大体可分为两类:基于质量的方法为每个部分解(partial solution)打分,并剔除那些被预测最终将导向错误答案的候选;而基于多样性的方法则倾向于选择遵循不同推理路径的部分解,从而更可能覆盖不同的最终答案。这两类目标之间的张力本质上即“探索”与“利用”的权衡问题;然而,此前关于剪枝的研究几乎从未直接揭示这一权衡关系,更遑论对其进行显式优化。本文提出一种新型剪枝方法,通过直接优化所选解池的 pass@$k$ 指标来应对该权衡问题:具体而言,该方法旨在选取一个部分解集合,使其在全部继续生成至完成之后,其中至少有一个能得出正确答案的概率最大化。由于该优化目标天然地同时奖励高质量的部分解、并惩罚它们之间的冗余性,因此质量与多样性得以由该准则本身自然平衡,而无需依赖额外的、人为设定的折中设计。我们在八个开源权重的推理模型上进行了实验,结果表明,本方法在性能上全面超越现有的基于质量与基于多样性的剪枝策略。尤其当剪枝过程保留中等数量的部分解以供后续生成时,性能提升尤为显著,这凸显了在剪枝阶段直接平衡质量与多样性的重要价值。
Test-time scaling has emerged as a key inference strategy for reasoning models, where a model generates a large number of candidate solutions in the hope of finding at least one correct solution. Running all candidate solutions to completion can be computationally demanding, so recent work prunes the candidate pool after generating only a short prefix, committing further computation only to the surviving candidates. Existing pruning strategies divide broadly into two categories: Quality-based methods assign a score to each partial solution and discard those predicted to ultimately lead to an incorrect final solution, while diversity-based methods select partial solutions that follow different reasoning trajectories and are therefore likely to cover different final solutions. The tension between these two objectives is an instance of exploration versus exploitation, yet little to no prior work on pruning exposes this trade-off directly, let alone optimizes it. In this paper, we introduce a pruning procedure that addresses this trade-off directly by optimizing pass@$k$ over the selected solution pool. In particular, it chooses the pool to maximize the probability that at least one of its partial solutions, when continued to completion, yields a correct answer. Because this objective simultaneously rewards high-quality partial solutions and penalizes redundancy among them, quality and diversity are naturally balanced by the criterion itself rather than through separate ad-hoc design choices. On eight open-weight reasoning models, our method outperforms existing quality-based and diversity-based pruning strategies. The gains are particularly pronounced when the pruning procedure keeps a moderate number of partial solutions for continued generation, highlighting the value of directly balancing quality and diversity during pruning.
许愿