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

Forum topic · 小凯 · 2026-10-24

Summary

This paper introduces a candidate-solution pruning method for test-time scaling in reasoning LLMs. Reasoning models generate many candidate solutions to increase the chance of at least one correct answer, but running all candidates is computationally expensive, so the candidate pool is often pruned after a short prefix. Prior pruning strategies are either quality-based (discarding partial solutions predicted to end up wrong) or diversity-based (selecting solutions following different reasoning paths), reflecting an exploration-vs-exploitation trade-off. The proposed approach directly optimizes pass@k over the selected pool, choosing solutions to maximize the probability that at least one partial solution, when continued, yields a correct answer. Because this objective rewards high-quality partial solutions while penalizing redundancy, quality and diversity are balanced naturally by the criterion itself rather than by ad hoc design choices. Experiments on eight open-weight reasoning models show the method outperforms existing quality-based and diversity-based pruning strategies. Paper: arXiv 2610.31667.

Paper Overview

Field: NLP Authors: Mahdi Salmani, Adel Javanmard, Devansh Gupta, Meisam Razaviyayn Published: 2026-10-22 arXiv: 2610.31667

Summary (translated from Chinese forum abstract)

Test-time scaling has become a key inference strategy for reasoning models: the model generates a large number of candidate solutions, hoping at least one is correct. Running all candidates to completion is computationally costly, so recent work prunes the candidate pool after generating only a short prefix. Existing pruning strategies fall into two categories: quality-based methods score each partial solution and discard those predicted to produce wrong final answers; diversity-based methods select partial solutions that follow different reasoning paths. The tension between these two objectives is essentially an exploration-vs-exploitation trade-off, but almost no prior work directly exposes or optimizes it.

This paper proposes a pruning method that handles the trade-off by directly optimizing pass@ $$k$$ over the pool of selected solutions — choosing a pool that maximizes the probability that at least one partial solution leads to a correct answer when generation continues. 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. Across eight open-weight reasoning models, the method outperforms existing quality-based and diversity-based pruning strategies.

Original Abstract

> 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. Existing pruning strategies divide into quality-based and diversity-based methods. We introduce a pruning procedure that addresses this trade-off directly by optimizing pass@k over the selected solution pool. Because this objective simultaneously rewards high-quality partial solutions and penalizes redundancy among them, quality and diversity are naturally balanced by the criterion itself. On eight open-weight reasoning models, our method outperforms existing pruning strategies.

--- *Auto-collected on 2026-10-24*