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

First, we cast pruning as maximizing the $\mathrm{pass}@k$ of the kept set, and we model how similar rollouts fail together. Second, we derive a closed-form, submodular approximation of $\mathrm{pass}@k$ that links the probe and RepExp. Third, through experiments on eight reasoning models, we show that PKTP improves over existing pruning methods, especially for moderate values of $k$.

Given a question \mathbfit{q}, we draw $N$ candidate reasoning rollouts y^{(i)}\sim\pi(\,\cdot\mid\mathbfit{q}\,) from a model $\pi$. Each rollout is observed only through its first $T$ tokens, the prefix $y_{1:T}^{(i)}$. We encode this prefix as a representation vector $h_{i}=h\!\left(y_{1:T}^{(i)}\right)$, a hidden-state embedding of the reasoning up to token T. Pruning then keeps $k\leq N$ rollouts, runs them to the full token budget $B$, and discards the rest. A completed rollout $y_{1:B}$ is either correct or not, recorded by the sign $C(y_{1:B})\in\{-1,+1\}$,

Here Verify is a task-specific checker that returns true exactly when a completed rollout reaches a correct final answer, for instance by matching a reference solution on a math problem or passing the unit tests on a coding task. We write $C(y^{(i)})$ for the outcome of the $i$-th rollout once it is run to the budget $B$.

To prune partial rollouts, we need to map each prefix to a vector that captures its essential information. One line of work relies on hand-designed features of the rollout and the model’s internal states, such as token confidence, the geometry of the hidden-state trajectory, or the convergence of intermediate answers. These features are usually very low dimensional [13, 40, 38, 19, 28]. A second line of work instead keeps the hidden states of the transformer and pools them into a single vector [42, 39, 29], and this is the class we use. In particular, we take $h_{i}$ to be an exponential moving average (EMA) of the hidden states, which has proven an effective pooling operation [29]. Our improvements are thus shown for this class of embeddings, although our framework can be applied to other embeddings as well. Appendix F gives the details of the representation we use.

At pruning time the outcomes $C(y^{(i)})$ are still unknown, since each depends on a (random) continuation we have not yet generated. How, then, should we choose the set $S\subseteq\{1,\ldots,N\}$ of $|S|=k$ rollouts to keep? Two extreme choices are as follows.

Neither extreme balances quality against diversity. What we actually want is the set of $k$ rollouts with the best $\mathrm{pass}@k$. Mathematically, we seek the set $S$ most likely to contain a correct rollout,

Here, $\mathbf{1}\{\cdot\}$ is the indicator function, i.e., $\mathbf{1}\{True\}=1$ and $\mathbf{1}\{False\}=0$. The objective in Eq. 3 is exactly $\mathrm{pass}@k$ of the kept pool, the probability that at least one surviving rollout, once run to the budget $B$, yields a correct answer. We also assume that the final outcomes are binary, although the same approach generalizes to continuous outcomes, where each completed rollout is scored by a reward model and the goal is to maximize the best of the kept rollouts [10], in which case Eq. 3 becomes a max@$k$ objective [41].

Existing pruning methods score each rollout on its own, using the model’s confidence [30, 13], a reward model [36], or a probe on its hidden states [54, 25], or they select a diverse set, as RepExp does [39]. Pass@$k$ has also been used as a training objective for reasoning models [41, 8], but not to decide which partial rollouts to keep. Outside of language models, choosing a set so that at least one item succeeds has been studied in information retrieval and Bayesian optimization [6, 34, 46]. Appendix G gives an extended discussion. In what follows, we discuss how such an objective can be solved.

Let $H=[h_{1},\ldots,h_{N}]\in\mathbb{R}^{d\times N}$ collect the representations of the $N$ generated partial rollouts, $Q_{\theta}:\mathbb{R}^{d}\to\mathbb{R}$ be a learnable mapping predicting the rollout quality, and define $Q_{\theta}(H)=[q_{1},\ldots,q_{N}]^{\top}$ with $q_{i}:=Q_{\theta}(h_{i})$. We model the latent scores of the rollouts as

where $\mathbf{m}$ and $\mathbf{f}$ are independent. Intuitively, $\mathbf{Z}_{i}$ is a score of how promising rollout $i$ is, so the higher its quality $q_{i}$, the higher its score. On top of $q_{i}$, the term $h_{i}^{\top}\mathbf{f}$ is a random effect shared by rollouts with similar representations, which makes them succeed or fail together. Indeed, for distinct rollouts, $\operatorname{Cov}(\mathbf{Z}_{i},\mathbf{Z}_{j}\mid H)=h_{i}^{\top}h_{j}$, so rollouts with aligned $h_{i}$ have positively correlated outcomes. Finally, $\mathbf{m}_{i}$ is noise specific to rollout $i$.

The correctness $C(y^{(i)})$ of rollout $i$ is then modeled by thresholding its score at zero. More specifically, we define the random variable $\mathbf{C}_{i}=2\times\mathbf{1}\{\mathbf{Z}_{i}\geq 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 $\mathbf{f}$, the likelihood of an observed correctness pattern $c\in\{-1,1\}^{N}$ is

where $\Phi$ is the standard Gaussian CDF, and the derivation can be found in Section A.1.

Suppose we have $M$ training questions, where question $m$ consists of the representations $H^{(m)}$ of a pool and its observed outcomes $c^{(m)}\in\{-1,1\}^{N}$. The mapping $Q_{\theta}$ is trained by maximizing their likelihood under Eq. 5, that is, by solving

where $\Pr_{\theta}$ is the likelihood of Eq. 5 with $H=H^{(m)}=[h_{1}^{(m)},\ldots,h_{N}^{(m)}]$ and $q_{i}=Q_{\theta}(h_{i}^{(m)})$. Solving Eq. 6 directly is hard, since the expectation over $\mathbf{f}$ in Eq. 5 has no closed form. We therefore expand each log-factor to second order in $h_{i}^{\top}\mathbf{f}$, which turns the expectation into a Gaussian integral that can be evaluated exactly (Theorem 1).

The proof is in Section B.1.1. The expansion is local, so we use Eq. 9 as a tractable surrogate rather than as an exact identity or a bound. We therefore replace $\Pr_{\theta}$ in Eq. 6 by $\widetilde{\Pr}_{\theta}$.

Instead of using $h_{i}$ directly in the correlation term, one could learn an encoder $\mathcal{E}_{\psi}$ of the representations and use $\mathcal{E}_{\psi}(h_{i})$ in its place, but its helpfulness depends on the amount of data. Since our data is limited, we keep $h_{i}$ fixed, so that $Q_{\theta}$ is the only learned component, and for the same reason we use a simple linear model $Q_{\theta}(h)=\theta^{\top}h+\theta_{0}$ for $Q_{\theta}$, as in the Linear Probe [54]. Unlike the Linear Probe, which is trained on the correctness of each rollout separately, $Q_{\theta}$ 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 Eq. 9 is in closed form and differentiable in $\theta$, we maximize the surrogate likelihood by stochastic gradient ascent over the training questions.

Maximizing $\mathrm{pass}@k$ in Eq. 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

and we seek the set of $k$ rollouts that minimizes it. Since $L(S)$ has no closed form, we approximate it with Theorem 1.

The proof is in Section B.1.2. We therefore minimize $\log\widetilde{L}(S)$ over $|S|=k$. Each term of Eq. 11 has a direct interpretation. The first term favors large $q_{i}$, because the baseline failure factor $\Phi(-q_{i})$ decreases with $q_{i}$. The negative log-determinant rewards a collection of $h_{i}$’s that covers several directions instead of repeatedly selecting the same direction. The last term is nonnegative and penalizes sets whose weighted vectors $\lambda_{i}h_{i}$ add up in the same direction, which the determinant alone cannot detect since it treats $h_{i}$ and $-h_{i}$ identically. For example, two rollouts of the same quality fail together more often with $h_{2}=h_{1}$ than with $h_{2}=-h_{1}$, yet both cases give the same $\Sigma_{S}$ and hence the same determinant, while $b_{S}$ is $2\lambda_{1}h_{1}$ in the first case and $0$ in the second.

However, minimizing $\log\widetilde{L}(S)$ exactly is hard.

The proof is in Section B.3, 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 [32]. The exact pass probability $P(S):=1-L(S)$ is known to be monotone and submodular [46]. An approximation of a submodular function, though, is not necessarily submodular, so we prove that its surrogate $\widetilde{P}(S):=1-\widetilde{L}(S)$ is also monotone and submodular, and hence the greedy optimizer and its approximation guarantee apply to our problem.

The proof is in Section B.2.

The surrogate makes the quality–diversity trade-off explicit and connects to the two baseline criteria in limiting regimes. If the correlation term in Eq. 11 is switched off, then $\log\widetilde{L}(S)=\sum_{i\in S}\log\Phi(-q_{i})$, so minimizing it keeps the $k$ largest quality scores and recovers a quality-only, probe-style ranking; it coincides with Linear Probe when $Q_{\theta}$ 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 $\kappa_{i}$ equals a common value $\kappa$, so $\Sigma_{S}=\kappa\sum_{i\in S}h_{i}h_{i}^{\top}$ and the relevant set-dependent term is $-\tfrac{1}{2}\log\det(I_{d}+\Sigma_{S})$. By the matrix determinant lemma, its greedy marginal ranks a candidate $j$ by $h_{j}^{\top}(I_{d}+\Sigma_{S})^{-1}h_{j}$, the residual-variance criterion used by RepExp on the same representations (see Section A.3 for the derivation). Between these limits, the modular quality, log-determinant diversity, and penalty terms jointly determine the selected set.

The above two extremes do not show how PKTP moves from Linear Probe to RepExp as $k$ grows. To see this, we expand $\log L(S)$ and its surrogate when the correlations are weak.

Proposition 2 shows how the correlation strength $\rho$ controls where PKTP sits between the two baselines. At $\rho=0$ the rollouts are independent, the expansion reduces to $\sum_{i\in S}\log\Phi(-q_{i})$, and PKTP ranks the rollouts by quality, as Linear Probe does. For $\rho>0$, every pair of kept rollouts adds a penalty $\rho^{2}\lambda_{i}\lambda_{j}h_{i}^{\top}h_{j}$, which is large when their representations are similar. A rollout added to a set of $t$ rollouts pays $t$ such penalties, so the first pick is decided by quality alone, and as the set grows the penalties add up until they outweigh the differences in quality, after which PKTP spreads its picks out, as RepExp does. The larger $\rho$ is, the sooner this happens, so $\rho$ sets how fast PKTP moves from Linear Probe to RepExp as $k$ grows.

Appendix C makes the two ends of this behavior precise for a fixed $\rho$. At small $k$, if the gap between the $k$th and $(k+1)$th largest qualities exceeds a constant multiple of $k\rho^{2}$, PKTP keeps exactly the $k$ candidates with the largest $q_{i}$, 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 $\rho$, in which PKTP follows Linear Probe at small $k$, while its averaged objective reduces to the RepExp objective as $k$ grows.

Figure 3 reports our main setting, in which each method selects $k\in\{1,2,4,8,16\}$ of the $N=128$ rollouts of a problem from their prefixes, and we set the cutoff threshold to $T=4000$ tokens. For every model it shows the $\mathrm{pass}@k$ of each pruning policy minus that of Random (Pass@$k$), pooled over all test problems. Random keeps $k$ rollouts chosen uniformly at random, which is equivalent to sampling $k$ rollouts without any pruning, so it is the natural reference point. The absolute numbers are given in Table 1 and the breakdown by benchmark in Appendix D.

As evident in Fig. 3, the Linear Probe performs strongly at $k=1$, but as $k$ increases it gets closer and closer to Random, as in the Qwen models, and in some cases it even becomes worse than Random at large $k$. This can be attributed to a bias towards a cluster of rollouts that the model predicts to be correct but that may not be, in the case where the classifier is wrong. In other words, ranking by the probe sharpens the distribution of the kept rollouts around what the probe considers correct. This mirrors the effect observed in reinforcement learning, where models trained with GRPO outperform their base models at small $k$ but are outperformed by them at large $k$ [53]. As an example, Fig. 4 shows the $128$ rollouts of one problem (Qwen3-8B, AMO-Bench), projected on the first two principal components of their representations. The $16$ rollouts ranked highest by the probe are all close to each other and all wrong, while $16$ random picks include $7$ correct ones.

Figure 3 also shows that RepExp behaves in the opposite way to the Linear Probe. At $k=1$ it shows, as expected, a performance similar to Random, since its first pick is random. As $k$ increases, the general trend is that RepExp gets better and better than Random, especially at $k=16$ (e.g., the DeepSeek-R1-Distill models). Interestingly, however, for moderate $k$ it is sometimes outperformed by Random, for instance on OpenMath-Nemotron-1.5B and Qwen3-8B. We attribute this to a tendency of RepExp to favor outliers at small and moderate $k$, since the rollouts farthest from those already selected are often unusual rollouts that are unlikely to be correct. A similar effect is reported by Tuyls et al. 39, where RepExp stops helping under high-temperature sampling, since the less coherent responses look novel in representation space yet do not necessarily contain correct answers. As an example, Fig. 5 shows the $128$ rollouts of one problem (Qwen3-8B, Omni-MATH), embedded in two dimensions with multidimensional scaling on their representations. After its first, random pick, RepExp selects outliers far from the other rollouts, and all $4$ of its picks are wrong, while $4$ random picks include $1$ correct one.

Finally, Fig. 3 shows that PKTP outperforms both the Linear Probe and RepExp, and its advantage is clearest at moderate $k$, where each baseline suffers from its own bias. 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. On Qwen3-0.6B it leads both baselines by more than $2$ points at every $k\geq 2$, and on DeepSeek-R1-Distill-Qwen-7B it is $3.3$ points above the Linear Probe at $k=4$. Unlike the two baselines, PKTP also stays above Random on every model for $k\geq 2$, since it keeps the correctness signal of the probe while spreading the kept rollouts over outcomes that are unlikely to fail together. At large $k$, however, the performance of PKTP and RepExp gets closer. We believe this is due to the weak correctness signal in the embeddings of the partial rollouts, which raises the question of how a quality embedding can be obtained from a partial rollout, and whether this is possible at all.

Results for the other pool sizes, $N\in\{32,48,64,80,96,112\}$ with the same cutoff $T=4000$, are given in Appendix E.

An important aspect of pruning is the cutoff threshold $T$, and how the algorithm behaves as the cutoff grows. Figure 6 varies $T$ on Qwen3-1.7B and DeepSeek-R1-Distill-Qwen-7B. For every $T\geq 4000$, PKTP outperforms all other methods on both models. An interesting observation is that a later cutoff in general improves the performance of PKTP. On Qwen3-1.7B its pass@$4$ rises from $63.5\%$ at $T=2000$ to $65.9\%$ at $T=8000$, and on DeepSeek-R1-Distill-Qwen-7B from $57.4\%$ to $61.7\%$. Intuitively, the more tokens are generated, the more is revealed about the quality of a rollout and where it is heading, so there is more signal to exploit. Interestingly, however, the same does not necessarily happen for the Linear Probe and RepExp. On Qwen3-1.7B, for example, the Linear Probe drifts down from $64.2\%$ to $62.3\%$ and ends below Random. One reason can be the representation we use. Since the EMA with rate $10^{-3}$ effectively averages over the last $1000$ tokens, it may lose the signal in the earlier tokens.

Another important aspect of pruning is the size $N$ of the pool the $k$ rollouts are chosen from. Figure 7 varies $N$ on Qwen3-1.7B and DeepSeek-R1-Distill-Qwen-7B. PKTP outperforms all other methods for every $N$ on Qwen3-1.7B, and on DeepSeek-R1-Distill-Qwen-7B once the pool reaches $N=80$. A larger pool offers more candidates to choose from, and on Qwen3-1.7B both baselines benefit from it, with RepExp rising from $62.9\%$ at $N=32$ to $63.9\%$ at $N=128$ as the pool becomes more diverse and the Linear Probe from $62.1\%$ to $63.1\%$. Interestingly, however, the same does not hold on DeepSeek-R1-Distill-Qwen-7B, where the Linear Probe falls from $58.1\%$ to $56.4\%$ and stays below Random for every $N$, and RepExp falls from $59.5\%$ at $N=48$ to $58.6\%$, while PKTP rises from $58.8\%$ to $59.7\%$. One reason can be that a larger pool also contains more outliers, which raises the chance that extreme cases like RepExp and the Linear Probe select them at a moderate $k$ such as $k=4$.

Write $\mathbf{Z}_{i}=q_{i}+\mathbf{m}_{i}+h_{i}^{\top}\mathbf{f}$. Conditioned on $\mathbf{f}$, the $\mathbf{Z}_{i}$ are independent of each other, so

Since the standard Gaussian is symmetric about $0$, $-c_{i}\mathbf{m}_{i}$ follows the same distribution as $\mathbf{m}_{i}$, and with $\Phi(x):=\Pr_{Z\sim\mathcal{N}(0,1)}[Z\leq x]$ each factor equals $\Phi\bigl(c_{i}(q_{i}+h_{i}^{\top}\mathbf{f})\bigr)$. Taking the expectation with respect to $\mathbf{f}$ on both sides gives Eq. 5.

As in Corollary 1, let $\Sigma_{S}=\sum_{i\in S}\kappa_{i}h_{i}h_{i}^{\top}$ and $b_{S}=\sum_{i\in S}\lambda_{i}h_{i}$. Write $\Omega_{S}=(I_{d}+\Sigma_{S})^{-1}$, $r_{j}=h_{j}^{\top}\Omega_{S}h_{j}$ and $u_{j}=h_{j}^{\top}\Omega_{S}b_{S}$. Adding a candidate $j\notin S$ changes $\Sigma_{S}$ by the rank-one matrix $\kappa_{j}h_{j}h_{j}^{\top}$. The matrix determinant lemma gives

while Sherman–Morrison gives

Substituting Eq. 18 and Eq. 19 into Eq. 11, and using $b_{S\cup\{j\}}=b_{S}+\lambda_{j}h_{j}$, the corresponding change in the unhalved quadratic form is

Half of this, together with the change $-\tfrac{1}{2}\log(1+\kappa_{j}r_{j})$ in the determinant term and the new quality term $\log\Phi(-q_{j})$, yields

The greedy step adds the candidate with the smallest $\Delta_{j}(S)$. After selecting $j$, the maintained state updates are

The remaining scalars $r_{\ell}$ and $u_{\ell}$ can be updated recursively as well. If $g_{\ell}=h_{\ell}^{\top}\Omega_{S}h_{j}$, then

With all $h_{i}$ retained in memory, each greedy step costs $O(Nd+d^{2})$, so one forward pass costs $O\bigl(k(Nd+d^{2})\bigr)$ time and $O(Nd+d^{2})$ storage. A naive implementation that instead rebuilds and factorizes $I_{d}+\Sigma_{S}+\kappa_{j}h_{j}h_{j}^{\top}$ for every candidate costs $O(d^{3})$ per candidate, hence $O(kNd^{3})$ time, with the same $O(Nd+d^{2})$ storage.

Suppose that every quality score equals a common value $q$, so that $\lambda_{i}=\lambda$ and $\kappa_{i}=\kappa$ for every $i$, and that the penalty term is dropped. Then Eq. 11 reduces to

Adding a candidate $j\notin S$ changes $\Sigma_{S}$ by the rank-one matrix $\kappa h_{j}h_{j}^{\top}$. The matrix determinant lemma, $\det(A+uv^{\top})=\det(A)\bigl(1+v^{\top}A^{-1}u\bigr)$ for an invertible $A$, applied with $A=I_{d}+\Sigma_{S}$ and $u=v=\sqrt{\kappa}\,h_{j}$, gives

Subtracting Eq. 24 at $S$ from Eq. 24 at $S\cup\{j\}$ and using Eq. 25, the factor $\det(I_{d}+\Sigma_{S})$ cancels and

The first term of Eq. 27 is the same for every candidate, and $x\mapsto\log(1+\kappa x)$ is increasing since $\kappa>0$, so the candidate with the smallest $\Delta_{j}(S)$ is the one with the largest $h_{j}^{\top}(I_{d}+\Sigma_{S})^{-1}h_{j}$. Moreover, $(I_{d}+\Sigma_{S})^{-1}=\kappa^{-1}\bigl(\kappa^{-1}I_{d}+\sum_{i\in S}h_{i}h_{i}^{\top}\bigr)^{-1}$, so this ranking is the RepExp rule of Section 2 with regularizer $\tau=1/\kappa$.