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.
Problem formulation, related work, and the pruning trade-off
Large Language Models (LLMs) have demonstrated remarkable capabilities in complex reasoning, showcasing a strong aptitude for solving complex problems that require multi-step logic. To further enhance these capabilities, test-time scaling has emerged as a prominent strategy, where generating multiple candidate solutions for a single problem allows the model to explore diverse reasoning strategies and achieve substantially better final performance. However, a significant challenge in scaling test-time compute is the prohibitive cost of generating many full solutions, which we refer to as rollouts moving forward. Beyond generation, there is also the cost of evaluating the finished rollouts in best-of-N methods, which score each rollout with a reward model, and this is computationally demanding on a large pool of long reasoning traces. This computational bottleneck, both in generation and in evaluating all rollouts in best-of-N methods, necessitates test-time pruning of partial rollouts before they reach a final answer and terminate. In other words, we generate all N rollouts only up to a cutoff of T tokens, keep the k most promising of them and run those to the full token budget B, as illustrated in Figure 1.
One family of methods examines the confidence dynamics of the model during generation, utilizing metrics such as logit entropy or confidence. The underlying intuition is that if a model is overly confident at an intermediate decoding step, extending the reasoning process is a waste of compute; conversely, if the model is entirely lacking in confidence, it is unlikely to reach a correct solution regardless of the additional compute provided, allowing for early termination. Another line of work leverages the model's internal state, specifically the hidden states of its transformer blocks, to identify which rollouts are worth continuing. By probing these internal states, one can predict whether the model will ultimately arrive at a correct solution, effectively framing the pruning problem as a classification task. This allows for the early discarding of partial rollouts that are unlikely to yield a correct final answer. In parallel, a separate line of research prioritizes representation exploration (RepExp) to select a subset of rollouts that cover as many distinct reasoning strategies as possible. By selecting partial rollouts whose corresponding embeddings are dissimilar, RepExp aims to explore a diverse set of reasoning trajectories.
These approaches, however, sit at the two ends of a spectrum, and neither directly optimizes the quantity we actually care about. The problem with probing methods is that the rollouts they rank highest tend to be similar to one another. As a result, the probe may be confidently wrong about all of the kept rollouts at once, and none of them reaches a correct answer. This is the problem of pure exploitation. RepExp, on the other hand, finds diverse rollouts, but because it ignores the quality of the rollouts, it can spend much of its budget on rollouts that are unlikely to succeed. These limitations motivate selecting the rollouts by the metric we actually want to optimize, choosing the set S of k rollouts that maximizes pass@k, that is, max|S|=k pass@k(S), which we formalize in equation 3. Optimizing this objective directly keeps the selected rollouts diverse enough to avoid shared failures while still favoring high-quality ones.
To this end, we propose PKTP (Pass@k Thought Pruning), a pruning framework that chooses the kept rollouts by maximizing pass@k directly. The main difficulty is that, even if the rollouts are sampled independently, their outcomes are not independent (conditioned on the partial observation of the rollout). Hence, pass@k depends on their joint distribution (i.e., rollouts with similar partial rollouts tend to succeed or fail together). We model this joint distribution with a latent Gaussian model, in which the correctness of each rollout is obtained by thresholding a Gaussian score. Evaluating pass@k under this model requires an expectation over the latent variables, so we instead derive a closed-form second-order approximation of it. We show that this approximation is submodular, enabling us to approximately solve it using a greedy procedure to find a subset of size k. An important aspect of the proposed framework is that the latent model is trained only once and does not depend on k, and a single greedy pass yields a nested sequence of subsets, so the budget k can be chosen online at inference time without retraining.
Probabilistic model and submodular approximation
Let H=[h1,…,hN]∈ℝd×N collect the representations of the N generated partial rollouts, Qθ:ℝd→ℝ be a learnable mapping predicting the rollout quality, and define Qθ(H)=[q1,…,qN]⊤ with qi:=Qθ(hi). We model the latent scores of the rollouts as Z=Qθ(H)⏟Rollout Quality+H⊤f⏟Rollouts Correlation+m⏟Outcome Noise, m∼𝒩(0,IN), f∼𝒩(0,Id), where m and f are independent. Intuitively, Zi is a score of how promising rollout i is, so the higher its quality qi, the higher its score. On top of qi, the term hi⊤f is a random effect shared by rollouts with similar representations, which makes them succeed or fail together. Indeed, for distinct rollouts, Cov(Zi,Zj∣H)=hi⊤hj, so rollouts with aligned hi have positively correlated outcomes. Finally, mi is noise specific to rollout i.
The correctness Ci of rollout i is then modeled by thresholding its score at zero. More specifically, we define the random variable Ci=2×1{Zi≥0}−1. In other words, a positive score implies that the rollout is correct, while a negative score implies that it is incorrect. Since the scores are independent conditionally on f, the likelihood of an observed correctness pattern c∈{−1,1}N is Pr[C=c∣H]=𝔼f[∏i=1NΦ(ci(qi+hi⊤f))], where Φ is the standard Gaussian CDF. Suppose we have M training questions, where question m consists of the representations H(m) of a pool and its observed outcomes c(m)∈{−1,1}N. The mapping Qθ is trained by maximizing their likelihood under this equation, that is, by solving maxθ(1/M)∑m=1M log Prθ[C=c(m)∣H(m)], where Prθ is the likelihood of this equation with Hi=[h1(m),…,hN(m)] and qi=Qθ(hi(m)). Solving this directly is hard, since the expectation over f has no closed form. We therefore expand each log-factor to second order in hi⊤f, which turns the expectation into a Gaussian integral that can be evaluated exactly.
Replacing each log-factor of Pr[C=c∣H] in this equation by this expansion at ε=hi⊤f gives Pr∼[C=c∣H]=(∏i=1NΦ(ciqi))det(Id+Σ(c))−1/2exp(12b(c)⊤(Id+Σ(c))−1b(c)), where Σ(c)=∑i=1Nκi(ci)hihi⊤ and b(c)=∑i=1Nciλi(ci)hi, λi(ci)=ϕ(qi)/Φ(ciqi), κi(ci)=λi(ci)(λi(ci)+ciqi), and ϕ is the standard Gaussian density. The expansion is local, so we use this equation as a tractable surrogate rather than as an exact identity or a bound. We therefore replace Prθ in the training equation by Pr∼θ.
Instead of using hi directly in the correlation term, one could learn an encoder of the representations and use ℰψ(hi) in its place, but its helpfulness depends on the amount of data. Since our data is limited, we keep hi fixed, so that Qθ is the only learned component, and for the same reason we use a simple linear model Qθ(h)=θ⊤h+θ0 for Qθ, as in the Linear Probe. Unlike the Linear Probe, which is trained on the correctness of each rollout separately, Qθ is trained on the joint outcomes of all rollouts of a question, so the quality it learns already accounts for the correlation between them. Since this equation is in closed form and differentiable in θ, we maximize the surrogate likelihood by stochastic gradient ascent over the training questions.
Maximizing pass@k in equation 3 is the same as minimizing the probability that every kept rollout fails. Under the latent Gaussian model, and suppressing the conditioning on H, the failure probability of a set S of kept rollouts is L(S)=Pr[Ci=−1 ∀i∈S]=𝔼f∼𝒩(0,Id)[∏i∈SΦ(−(qi+hi⊤f))], and we seek the set of k rollouts that minimizes it. Since L(S) has no closed form, we approximate it with the surrogate in the selection equation. For a set S⊆{1,…,N}:={1,…,N}, let λi=ϕ(qi)/Φ(−qi), κi=λi(λi−qi), ΣS=∑i∈Sκihihi⊤ and bS=∑i∈Sλihi. Applying the second-order surrogate to the rollouts in S with ci=−1 for every i∈S gives a surrogate L∼(S) of L(S) with log L∼(S)=∑i∈S log Φ(−qi)+(−12 log det(Id+ΣS)+12bS⊤(Id+ΣS)−1bS). The first term favors large qi, because the baseline failure factor Φ(−qi) decreases with qi. The negative log-determinant rewards a collection of hi's that covers several directions instead of repeatedly selecting the same direction. The last term is nonnegative and penalizes sets whose weighted vectors λihi add up in the same direction, which the determinant alone cannot detect since it treats hi and −hi identically.
However, minimizing log L∼(S) exactly is hard. Minimizing log L∼(S) over all S⊆{1,…,N} with |S|=k is NP-hard, by a reduction from the maximum independent set problem. When the objective is monotone and submodular, however, a greedy optimizer that adds one rollout at a time is guaranteed to achieve at least a (1−e−1) fraction of the optimal value. The exact pass probability P(S):=1−L(S) is known to be monotone and submodular. An approximation of a submodular function, though, is not necessarily submodular, so we prove that its surrogate P∼(S):=1−L∼(S) is also monotone and submodular, and hence the greedy optimizer and its approximation guarantee apply to our problem.
The surrogate makes the quality–diversity trade-off explicit and connects to the two baseline criteria in limiting regimes. If the correlation term in the selection surrogate is switched off, then log L∼(S)=∑i∈S log Φ(−qi), so minimizing it keeps the k largest quality scores and recovers a quality-only, probe-style ranking; it coincides with Linear Probe when Qθ is the corresponding linear probe. At the other extreme, suppose that the quality scores are equal and that the penalty term is negligible throughout selection. Then every κi equals a common value κ, so ΣS=κ∑i∈Shihi⊤ and the relevant set-dependent term is −12 log det(Id+ΣS). By the matrix determinant lemma, its greedy marginal ranks a candidate j by hj⊤(Id+ΣS)−1hj, the residual-variance criterion used by RepExp on the same representations. Between these limits, the modular quality, log-determinant diversity, and penalty terms jointly determine the selected set.
To see how PKTP moves from Linear Probe to RepExp as k grows, we expand log L(S) and its surrogate when the correlations are weak. As ρ→0, log L(S)=∑i∈S[log Φ(−qi)+ρ2/2 λiqi‖hi‖2]+ρ2∑i
At small k, if the gap between the kth and (k+1)th largest qualities exceeds a constant multiple of kρ2, PKTP keeps exactly the k candidates with the largest qi, as Linear Probe does, and in the same order when this holds at every budget up to k. At large k, when a few candidates of high quality stand out from a bulk of lower quality, the objective of PKTP, averaged over the bulk qualities, reduces to the RepExp objective plus a term for these candidates. Its minimizers therefore keep them and nearly maximize the RepExp objective. Finally, the conditions of these two results do not conflict. We construct a single model, with one quality distribution and one value of ρ, in which PKTP follows Linear Probe at small k, while its averaged objective reduces to the RepExp objective as k grows.
Experiments and results
We evaluate on eight open-weight reasoning models in two size tiers. The small tier is Qwen3-0.6B, Qwen3-1.7B, OpenMath-Nemotron-1.5B and DeepSeek-R1-Distill-Qwen-1.5B, and the large tier is Qwen3-8B, OpenMath-Nemotron-7B, DeepSeek-R1-Distill-Qwen-7B and Skywork-OR1-Math-7B. Every model generates its rollouts with temperature 0.8, top-pp of 0.95, and a full token budget of B=32,768 tokens.
We pool problems from several benchmark sources rather than a single one, since a pruning policy has to be trained on part of the problems and evaluated on the rest, and no single competition benchmark is large enough for both. Both tiers use AIME 2024, 2025 and 2026 and HMMT 2024, 2025 and 2026, together with RIMO-N. The small tier adds BeyondAIME, OlymMATH (Easy) and Omni-MATH levels 7 and 8, while the large tier replaces these by AMO-Bench, OlymMATH (Hard) and Omni-MATH levels 8 and 9. This gives around 600 problems per model. For each model we then remove the problems on which it generates all correct or all wrong solutions in 128 trials.
A completed rollout counts as correct when its final answer matches the reference answer of the problem. We take that answer from the \boxed{} expression of the rollout and ask a Qwen3-32B judge whether it is equivalent to the reference answer. We use a judge rather than a symbolic checker, because benchmarks like AMO-Bench have descriptive answers that a symbolic checker cannot handle.
Only a few hundred problems are available for each model, so a single split into training and test problems would leave too few on either side, and we evaluate with cross-validation instead. For each model we partition its problems into eight folds, with each benchmark spread evenly across the folds. Each fold serves once as the test set while the learned components are trained on the other seven, so every problem is scored exactly once by a model that was not trained on it. One such pass over all eight folds forms a repetition. Because the outcome of a repetition depends on how the problems happen to be divided, we run 2020 repetitions, each with a new random partition obtained by shuffling the problems of every benchmark before assigning them to folds, and with fresh random seeds for training. Every reported number is the average over the 2020 repetitions, and its standard error is computed across them. For RepExp we set the regularizer to τ=1. For all methods, the representation hi of a partial rollout is the EMA of its hidden states with rate 10−3.
The experimental results show that PKTP outperforms existing quality-based and diversity-based pruning strategies across all eight models. 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. Compared to Linear Probe, PKTP achieves significant improvements because it accounts for the correlation between rollouts with similar representations. Compared to RepExp, PKTP achieves comparable or better results because it also accounts for the quality of each rollout. The nested subset property of the greedy procedure means that the budget k can be chosen online at inference time without retraining, which is a practical advantage over methods that require re-optimization for each budget.
Relation to baseline criteria in limiting regimes
The surrogate makes the quality–diversity trade-off explicit and connects to the two baseline criteria in limiting regimes. If the correlation term in the selection surrogate is switched off, then log L∼(S)=∑i∈S log Φ(−qi), so minimizing it keeps the k largest quality scores and recovers a quality-only, probe-style ranking; it coincides with Linear Probe when Qθ is the corresponding linear probe. At the other extreme, suppose that the quality scores are equal and that the penalty term is negligible throughout selection. Then every κi equals a common value κ, so ΣS=κ∑i∈Shihi⊤ and the relevant set-dependent term is −12 log det(Id+ΣS). By the matrix determinant lemma, its greedy marginal ranks a candidate j by hj⊤(Id+ΣS)−1hj, the residual-variance criterion used by RepExp on the same representations. Between these limits, the modular quality, log-determinant diversity, and penalty terms jointly determine the selected set.
To see how PKTP moves from Linear Probe to RepExp as k grows, we expand log L(S) and its surrogate when the correlations are weak. As ρ→0, log L(S)=∑i∈S[log Φ(−qi)+ρ2/2 λiqi‖hi‖2]+ρ2∑i
At small k, if the gap between the kth and (k+1)th largest qualities exceeds a constant multiple of kρ2, PKTP keeps exactly the k candidates with the largest qi, as Linear Probe does, and in the same order when this holds at every budget up to k. At large k, when a few candidates of high quality stand out from a bulk of lower quality, the objective of PKTP, averaged over the bulk qualities, reduces to the RepExp objective plus a term for these candidates. Its minimizers therefore keep them and nearly maximize the RepExp objective. Finally, the conditions of these two results do not conflict. We construct a single model, with one quality distribution and one value of ρ, in which PKTP follows Linear Probe at small k, while its averaged objective reduces to the RepExp objective as k grows.
Limitations and future directions
The framework assumes a latent Gaussian model for the rollout outcomes, which may not capture all sources of dependence among rollouts. The representation used for correlation is the EMA of hidden states, and other pooling operations or encoders may yield different results. The experiments cover a finite set of benchmarks and model families, and broader generalization to other reasoning domains remains to be tested. Additionally, the greedy optimizer requires computing rank-one updates to the determinant and inverse, which adds a computational overhead that scales as O(k(Nd+d2)), and further engineering may be needed for very large pools.
Future work might explore learned encoders of the representations, alternative probabilistic models for rollout dependence, and extension of PKTP to other test-time scaling objectives such as best-of-N. The framework also suggests that directly optimizing the metric one cares about, rather than relying on separate quality and diversity heuristics, is a general principle that may apply beyond test-time pruning.