Direct Pass@$k$ Optimization for Thought Pruning: Balancing Quality and Diversity in Test-Time Reasoning
Test-time scaling generates many candidate reasoning traces, but completing all of them is prohibitively expensive. This paper proposes PKTP, a pruning framework that selects a subset of partial rollouts to maximize the probability that at least one of them ultimately succeeds—i.e., $\text{pass}@k$. It models joint correctness with a latent Gaussian threshold model, derives a closed-form submodular surrogate, and greedily selects the kept set with a $(1-1/e)$ approximation guarantee. The work matters because it directly exposes and optimizes the quality–diversity trade-off that existing pruning methods handle only indirectly.
The paper offers a principled, set-level formulation of thought pruning and shows consistent empirical gains at moderate $k$ across eight open-weight models. However, the strongest caveat is that the $(1-1/e)$ greedy guarantee of Theorem 2 applies only to a local second-order surrogate, not to the true $\text{pass}@k$ objective. As the authors themselves note, "The expansion is local, so we use Eq 9 as a tractable surrogate rather than as an exact identity or a bound." Additionally, the reliance on an external Qwen3-32B judge for answer verification and the absence of any released code or training scripts raise barriers to independent validation.
Casting pruning as exact $\text{pass}@k$ maximization in Eq. 3 is elegant and directly targets the metric practitioners care about. A particular strength is the analytical unification in Section 3.3: switching off the correlation term recovers a quality-only probe, while the equal-quality limit yields the RepExp diversity criterion, showing concretely how PKTP interpolates between the two extremes. The empirical protocol is also unusually thorough, using eight models, multiple math benchmarks, and 20 cross-validation repetitions to reduce variance.
The core limitation is the uncontrolled approximation gap between the surrogate $\widetilde{L}(S)$ and the true failure probability $L(S)$. The greedy guarantee is formally vacuous for the actual objective because the surrogate is neither a bound nor proven uniformly close. A second concern is the strong, largely untested modeling assumption that correctness arises from thresholding a latent Gaussian score that is linear in fixed EMA embeddings. The authors acknowledge this indirectly by noting that "the weak correctness signal in the embeddings of the partial rollouts" limits large-$k$ performance, which raises the deeper question of whether the chosen representation space suffices to support reliable joint modeling at all. Finally, the paper does not discuss calibration of the probability estimates produced by the latent model, so it is unclear whether the optimized surrogate correlates meaningfully with true success probabilities.
The evidence supports the central claim that direct $\text{pass}@k$ optimization dominates pure quality and pure diversity at moderate $k$. For instance, "On DeepSeek-R1-Distill-Qwen-1.5B at k=4, for example, PKTP is 3.0 points above the Linear Probe and 3.2 points above RepExp," and PKTP stays above Random for every model at $k \geq 2$. That said, the baselines exhibit strange behavior—RepExp is "sometimes outperformed by Random" and the Linear Probe drifts below Random on some models—which suggests that the embedding quality is itself a limiting factor rather than a neutral testbed. The comparison does not include recent process-reward-model-based pruners or stronger confidence-based early-exit methods, leaving open whether the gains are specific to the chosen baseline tier.
The paper reports many key hyperparameters—$N=128$, $T=4000$, $B=32768$, temperature $0.8$, top-p $0.95$, EMA rate $10^{-3}$, and $\tau=1$—along with a clear cross-validation scheme of eight folds and 20 repetitions. However, no code, data, or training configurations (learning rate, optimizer, number of epochs) are released or even mentioned. Reproduction would require re-implementing the second-order surrogate likelihood of Eq. 9 and the rank-one greedy updates of Appendix A.2 from scratch, which is nontrivial given that small implementation details can affect submodular greedy behavior. The use of a Qwen3-32B judge for correctness verification means exact numbers may be sensitive to the judge version and prompt template, neither of which are provided.
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.
Pick a starting point or write your own. Challenges run in the background, so you can keep reading while the AI investigates.
No challenges yet. Disagree with the review? Ask the AI to revisit a specific claim.