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

Mahdi Salmani; Adel Javanmard; Devansh Gupta; Meisam Razaviyayn

cs.AIOct 22, 2026

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

Authors: Mahdi Salmani, Adel Javanmard, Devansh Gupta, Meisam Razaviyayn

Organizations: University of Southern California

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, 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.

Figures & tables

Figure 1: Test-time pruning. Only

k

of the

N

rollouts (ringed) continue past the cutoff

T

, up to the budget

B

Figure 2: Conceptual illustration. The same pool of candidate rollouts, positioned by their partial rollout representations in embedding space, under three pruning strategies. Each rollout is colored by its predicted correctness and marked ✓/✗ by its actual outcome; the

k=3

kept rollouts are ringed. (a) Pure exploitation keeps the highest-scoring rollouts, but they are correlated and cluster together, so they miss the correct rollout. (b) Pure exploration keeps a maximally diverse subset yet spends its budget on low-quality rollouts, also missing it. (c) Our method optimizes

\mathrm{pass}@k

directly over the kept subset, balancing quality and diversity, and retains the correct rollout.

Figure 3: Gain in

\mathrm{pass}@k

over random pruning. For every model, the pooled

\mathrm{pass}@k

of each pruning policy minus that of Random (Pass@

k

), in percentage points, as a function of the number

k

of rollouts kept, with a band of one standard error of the repetition-by-repetition difference. The absolute numbers are in Table 1 and the per-benchmark breakdown in Appendix D .

Figure 4: Linear Probe keeps a cluster of rollouts that fail together. The

128

rollouts of one problem (Qwen3-8B, AMO-Bench), projected on the first two principal components of their representations, with correct rollouts in green and wrong rollouts in red. The

16

rollouts ranked highest by the Linear Probe (blue rings) lie close together and are all wrong, while

16

random picks (gray rings) include

7

correct rollouts.

Figure 5: RepExp keeps outliers that are unlikely to be correct. The

128

rollouts of one problem (Qwen3-8B, Omni-MATH), embedded in two dimensions by multidimensional scaling of their representations, with correct rollouts in green and wrong rollouts in red. After a random first pick, RepExp (pink rings, numbered in pick order) selects rollouts far from the rest, and all

4

of its picks are wrong, while

4

random picks (gray rings) include

1

correct rollout.

Figure 6: A later cutoff improves PKTP. Pass@

4

on Qwen3-1.7B and DeepSeek-R1-Distill-Qwen-7B as the cutoff

T

grows from

2000

8000

tokens, with

N=128

. PKTP is the best method for every

T\geq 4000

and keeps improving with longer prefixes, while the Linear Probe and RepExp gain much less, and the Linear Probe on Qwen3-1.7B even drifts below Random. Bands are one standard error.

Figure 7: PKTP leads once the pool is large enough. Pass@

4

on Qwen3-1.7B and DeepSeek-R1-Distill-Qwen-7B as the pool size

N

grows from

32

128

, with

T=4000

. PKTP is the best method for every

N

on Qwen3-1.7B and from

N=80

on DeepSeek-R1-Distill-Qwen-7B, where the Linear Probe falls as the pool grows and stays below Random. Bands are one standard error.

Appendix figures & tables6 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 8: Gain in

\mathrm{pass}@k

over random pruning with

N=32

. As in Fig. 3 , with a pool of

N=32

rollouts drawn from the

128

generated for each problem and

T=4000

Figure 9: Gain in

\mathrm{pass}@k

over random pruning with

N=48

. As in Fig. 3 , with a pool of

N=48

rollouts drawn from the

128

generated for each problem and

T=4000

Figure 10: Gain in

\mathrm{pass}@k

over random pruning with

N=64

. As in Fig. 3 , with a pool of

N=64

rollouts drawn from the

128

generated for each problem and

T=4000

Figure 11: Gain in

\mathrm{pass}@k

over random pruning with

N=80

. As in Fig. 3 , with a pool of

N=80

rollouts drawn from the

128

generated for each problem and

T=4000

Figure 12: Gain in

\mathrm{pass}@k

over random pruning with

N=96

. As in Fig. 3 , with a pool of

N=96

rollouts drawn from the

128

generated for each problem and

T=4000

Figure 13: Gain in

\mathrm{pass}@k

over random pruning with

N=112

. As in Fig. 3 , with a pool of

N=112

rollouts drawn from the

128

generated for each problem and

T=4000

Efficient Language Model Reasoning

Efficient language model reasoning studies how to improve problem-solving accuracy while reducing inference, memory, latency, and training costs.

See all 366 papers

Sep 23, 2026

Planned Test-Time Scaling with Coordinated Reasoning Paths

Test-time scaling with parallel branches is widely adopted to improve performance on challenging reasoning tasks. The predominant approach, repeated sampling, draws branches independently from a single policy, which can produce redundant attempts and thereby limit the gains from additional inference compute. To address this limitation, we propose Planned Test-Time Scaling (PTTS), which replaces independent sampling with a coordinated joint policy: a planner generates a solution outline for each branch, steering the branches toward distinct reasoning paths, and an executor produces a full solution conditioned on each outline. Formally, we show that PTTS strictly generalizes repeated sampling and, in a stylized setting, provably promotes coverage of complementary reasoning modes and yields better pass@k scaling. We instantiate PTTS on top of strong reasoning models, keeping them fixed as executors while replacing repeated sampling with PTTS inference to further enhance test-time scaling. Concretely, we develop two variants: PTTS-ZS prompts a model to jointly generate outlines for all branches in a single autoregressive pass, while PTTS-RL directly optimizes the planner against the pass@k reward using truncated execution rollouts for efficient training and a sharper reward signal. Across five mathematical reasoning benchmarks with Qwen3-1.7B and 4B, PTTS-ZS improves pass@64 over repeated sampling by up to 6.7 points, while PTTS-RL further increases the gain to up to 13.4 points. Further analysis indicates that broader coverage of distinct reasoning paths contributes to these gains. Overall, PTTS provides a general framework for improving test-time scaling by coordinating reasoning branches, with zero-shot and trainable instantiations that yield substantial performance gains.

Xueqing Wu, Langxing Bai, Hritik Bansal +5

University of California, Los Angeles · Amazon

RL for Language Model Reasoning Test-Time Scaling LLM Reasoning

Ask PDF

Aug 8, 2026

Thought-Level Beam Search for Reasoning

Test-time compute scaling is a primary driver of performance in large reasoning models (LRMs), but extreme inefficiency bounds current approaches, shifting the critical question from how much compute to spend, to where to allocate it. We formalize test-time reasoning as a constrained compute allocation problem over partial trajectories. Under a fixed hardware budget, existing paradigms fail to actively allocate the compute to the most promising partial progress: traditional parallel sampling treats traces independently and induces severe memory bottlenecks, while subtractive pruning starves hardware and fails to actively and sufficiently shift the output distribution. To overcome this dichotomy, we introduce Gambit, an inference algorithm that executes thought-level beam search. By periodically pruning unpromising trajectories and immediately branching from high-quality prefixes, Gambit dynamically concentrates compute onto the most promising reasoning traces via a light-weight scorer probing hidden states while maintaining continuous high hardware utilization. Extensive evaluations across multiple models and benchmarks demonstrate that Gambit strictly dominates existing baselines. Under identical hardware constraints, our method yields up to a +6.7% absolute accuracy gain on HMMT-24 and +3.3% on AIME-25 over pruning baselines, delivers

>2\times

higher throughput on trace completion, and reduces total token consumption by up to 68.5% relative to standard parallel sampling.

Lijie Yang, Hongyin Luo, Jiawei Zhao +2

Princeton University · MIT CSAIL · Meta AI

Inference-Time Search Efficient Language Model Reasoning Inference-Time Compute Allocation Efficient Language Model Inference Algorithmic Reasoning Large Reasoning Models

Ask PDF

May 24, 2026

Beyond the Frontier: Stochastic Backtracking for Efficient Test-Time Scaling

Test-time scaling improves language model reasoning by spending additional compute to explore multiple solution trajectories. The key challenge is to maximize accuracy while minimizing the total number of generated tokens during reasoning. Recent PRM-guided methods score intermediate prefixes to steer this search, but most are frontier-only: they keep only the current active prefixes and irreversibly prune or resample away the rest using noisy PRM scores. This can cause premature commitment, diversity collapse, and the loss of prefixes that still admit correct continuations. We introduce stochastic backtracking over a persistent pool of historical prefixes, allowing test-time compute to revisit previously generated states instead of only expanding the current frontier. To make this efficient, we propose two complementary mechanisms. Subpool Selection strengthens greedy PRM-guided search by applying Top-N selection within random subpools, giving historical prefixes a chance to bypass over-scored frontier candidates. Power Backtrack Sequential Monte Carlo extends SMC-style resampling to the persistent pool using powered PRM scores and mixture-corrected weights. Across mathematical reasoning benchmarks and model scales, our methods consistently achieve higher accuracy per token count, and the same level of accuracy using only a fraction of the token count in comparison to strong PRM-guided baselines, demonstrating that persistent-pool stochastic backtracking provides a simple and effective way to improve the accuracy-token trade-off in test-time scaling.

Dao Tran, Duc Anh Le, Ngoc Luu +3

Qualcomm AI Research†

Inference-Time Scaling Efficient Language Model Reasoning

Ask PDF

Abstract

Efficient Language Model Reasoning

Explore similar work

Planned Test-Time Scaling with Coordinated Reasoning Paths

Thought-Level Beam Search for Reasoning

Beyond the Frontier: Stochastic Backtracking for Efficient Test-Time Scaling