Directly Optimizing Pass@k: How Thought Pruning Balances Reasoning Quality and Diversity

Reasoning models in test-time scaling improve the probability of finding correct answers by generating a large number of candidate solutions, but running all candidates through the full reasoning chain is extremely computationally expensive. A compromise is to prune the candidate pool after generating short prefixes, letting only promising candidates continue computing. Existing pruning strategies fall into two broad categories: quality-based methods score each partial solution and discard those likely to end up wrong; diversity-based methods select partial solutions from different reasoning trajectories to cover more final answers. The tension between these two objectives is essentially the exploration-exploitation trade-off, but few prior works have directly exposed and optimized this trade-off.

From the Divide Between Quality and Diversity to Unified Optimization

According to the paper (source: rss), existing pruning strategies focus either only on quality or only on diversity, lacking a principled criterion to balance both simultaneously. The paper proposes a new pruning pipeline that directly optimizes pass@k for the selected solution pool—that is, maximizing the probability that at least one partial solution in the pool yields a correct answer after continued generation. This objective simultaneously rewards high-quality partial solutions and penalizes redundancy among them, so quality and diversity are naturally balanced by the criterion itself rather than through separate ad-hoc design choices.

pass@k is a metric commonly used to evaluate code generation and reasoning models, defined as the probability that among k complete solutions independently sampled from the model, at least one is correct. In the context of test-time scaling, it measures the "coverage capability" of the candidate pool. For partial solutions (incomplete reasoning prefixes), pass@k can be understood as: if every prefix in the pool were allowed to continue generating until completion, the probability of ultimately obtaining at least one correct answer. This metric naturally penalizes both low-quality prefixes (unlikely to lead to correct solutions) and redundant prefixes (leading to the same answer as existing prefixes), so optimizing it is equivalent to finding a balance point between exploration and exploitation.

Test-time scaling refers to investing more computation at inference time (e.g., generating more candidates, extending chain-of-thought) to improve model performance, rather than relying solely on scaling during training. Its core challenge is how to allocate compute when the budget is limited: should a few candidates run through the full reasoning chain, or should a large number of candidates generate only short prefixes and then be pruned? The quality of the pruning strategy directly determines the cost-effectiveness of test-time scaling.

Method Core: Pass@k as the Pruning Objective

The directness of this method lies in treating pruning as a combinatorial optimization problem: select a subset from all partial solutions such that the probability of at least one correct solution after continued generation is maximized. Since the definition of pass@k naturally considers the coverage relationships among multiple solutions, optimizing it is equivalent to finding the best trade-off point between quality (the likelihood that an individual solution is correct) and diversity (non-redundancy among solutions). The paper emphasizes that this balance is not achieved through manual tuning but is endogenously determined by the optimization objective.

Modeling pruning as a combinatorial optimization problem means selecting an optimal solution from an exponentially large number of possible subsets. In practice, greedy or approximate algorithms are typically used: at each step, select the partial solution that most increases the estimated pass@k of the pool, until the budget is reached. This requires estimating the probability that each partial solution is "correct after continued generation," with common approaches including using the model's own log-probabilities, training lightweight scorers, or calibrating with a small number of complete trajectories. The accuracy of the estimation directly affects pruning effectiveness and is also where the method's potential computational and implementation burden lies.

Experimental Performance: Outperforming Baselines on Eight Models

On eight open-weight reasoning models, the method outperforms existing quality-based and diversity-based pruning strategies. The gains are particularly pronounced when the pruning process retains a moderate number of partial solutions for continued generation, highlighting the value of directly balancing quality and diversity at the pruning stage. The paper does not disclose specific model names and datasets, but this result suggests that the method generalizes well across different reasoning models.

Implications: Pruning Should Be an Objective-Driven Decision

The significance of this work lies in elevating pruning from "heuristic filtering" to "objective-driven decision-making." In the past, quality scores and diversity metrics were often designed separately and tuned independently, whereas directly optimizing the downstream task metric (pass@k) aligns the pruning strategy with the ultimate goal. For teams deploying reasoning models in practice, this means that without significantly increasing computation, smarter pruning can retain more potentially correct solutions, thereby improving the efficiency of test-time scaling.

However, the method also requires estimating the success probability of partial solutions for each model and task, which may itself introduce additional overhead. The paper's abstract does not detail the specific approach to probability estimation, which may be an implementation detail that future practitioners need to pay attention to.

Overall, this research provides a clear principle for pruning in test-time reasoning: do not pursue quality or diversity in isolation, but let the pruning criterion directly reflect the metric you ultimately want.

One-line Summary

The paper proposes a thought pruning method that directly optimizes pass@k, unifying the balance between quality and diversity.

Paragraph Summary

Reasoning models in test-time scaling improve accuracy by generating a large number of candidate solutions, but full-chain computation is costly, necessitating pruning at the prefix stage. Existing pruning strategies fall into quality-based and diversity-based categories, which are essentially an exploration-exploitation trade-off that is rarely directly optimized. This paper proposes a new pipeline that treats pruning as a combinatorial optimization problem, directly maximizing the pass@k of the candidate pool—the probability that at least one partial solution in the pool yields a correct answer after continued generation. This objective naturally rewards quality and penalizes redundancy, making quality and diversity endogenously balanced by the criterion. On eight open-weight reasoning models, the method outperforms existing baselines, with notable gains especially when retaining a moderate number of partial solutions. Its significance lies in elevating pruning from heuristic filtering to objective-driven decision-making, though it requires estimating the success probability of partial solutions, which may introduce additional overhead.