PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Mahdi Salmani, Adel Javanmard, Devansh Gupta, Meisam Razaviyayn

arXiv:2610.31667v1cs.AIcs.CL

TL;DR

Pruning is needed because generating and evaluating many complete reasoning rollouts is costly, while existing quality- and diversity-based strategies do not directly optimize the desired pool outcome. The paper selects partial rollouts by optimizing pass@k with a model of joint outcomes, and reports stronger performance than existing pruning methods across eight reasoning models, especially at moderate k.

  • Problem

    Generating many full solutions and evaluating long rollouts is costly, while existing quality- and diversity-based pruning strategies do not directly optimize pass@k for the selected pool.

    Source: §1 · pp. 1–2

  • Method

    PKTP models joint rollout outcomes with a latent Gaussian model and selects partial rollouts by maximizing a closed-form approximation to pass@k.

    Source: §1 · p. 2, §3.1 · p. 6

  • Results

    PKTP outperforms existing quality-based and diversity-based pruning methods across eight reasoning models, with the largest gains at moderate k.

    Source: §1 · p. 3, §5 · p. 13

  • Takeaways & Limitations

    Directly optimizing the kept pool’s pass@k balances rollout quality and diversity within the selection criterion.

    Source: §1 · p. 2

    The formulation assumes binary final outcomes, though the authors state it can generalize to continuous reward-model outcomes with a max@k objective.

    Source: Setup · p. 4

Abstract

from arXiv · show

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.

1 Introduction

Generating and evaluating many complete reasoning rollouts is costly, motivating pruning after a short prefix. PKTP selects a pool by directly optimizing pass@k, balancing rollout quality and diversity, and outperforms existing pruning strategies, especially at moderate k.

1 Introduction

  • Test-time pruning addresses the high cost of generating many full solutions and evaluating long reasoning traces in best-of-N methods.The procedure generates N rollouts to cutoff T, then continues only k of them to budget B.Source: §1 · pp. 1–2, Figure 1 · p. 2
  • Quality-based probing can keep similar rollouts that fail together, while diversity-only RepExp can spend its budget on low-quality candidates.Figure 2 illustrates both extremes missing a correct rollout while direct pass@k selection retains one.Source: §1 · p. 2, Figure 2 · p. 3
  • PKTP chooses the k-rollout subset that maximizes pass@k, the probability that at least one continuation is correct.This objective favors quality while penalizing redundancy among selected rollouts.Source: §1 · p. 2
  • PKTP models correlated rollout outcomes with a latent Gaussian model and uses a closed-form second-order approximation with greedy subset selection.A single trained latent model supports nested subsets, allowing k to be chosen online without retraining.Source: §1 · pp. 2–3
  • PKTP outperforms existing pruning methods across eight reasoning models, with gains especially pronounced at moderate k.The paper frames its method as directly balancing quality and diversity during selection.Source: §1 · p. 3

2 Problem Formulation and Related Work

The formulation treats pruning as selecting partial rollouts whose continuations are likely to include at least one correct answer. It contrasts individual quality scoring with diversity selection and specifies the representations, outcomes, and objective used in the paper.

Setup.

  • Each of N candidate rollouts is observed through its first T tokens, represented by a hidden-state embedding, and selected rollouts continue to budget B.A completed rollout receives a binary correctness outcome from a task-specific checker.Source: Setup · pp. 3–4
  • The paper uses pooled transformer hidden states, specifically an exponential moving average, as rollout representations; its improvements are shown for this embedding class.The framework can also be applied to other embeddings.Source: Setup · p. 4
  • Linear Probe keeps the k rollouts with highest predicted success, whereas RepExp greedily selects dissimilar representations to cover diverse trajectories.These correspond to pure exploitation and pure exploration, respectively.Source: Setup · p. 4
  • The objective selects a size-k set most likely to contain a rollout that becomes correct when continued to completion.The paper calls this probability pass@k for the kept pool.Source: Setup · p. 4
  • The formulation assumes binary final outcomes, while the authors note that it generalizes to continuous reward-model scores using a max@k objective.A task-specific Verify checker can, for example, match a reference solution or test code.Source: Setup · p. 4
  • Prior pruning work scores rollouts individually using confidence, reward models, or hidden-state probes, or selects diverse sets as RepExp does.The passage also distinguishes pass@k as a prior training objective from its use here to choose partial rollouts.Source: Setup · p. 5

3 Method

The method estimates how partial-rollout representations relate to joint correctness outcomes, then selects a subset using a tractable pass@k surrogate. Its probabilistic model captures shared failures, and its greedy procedure connects quality-based and diversity-based selection.

3 Method

  • The method learns a joint outcome distribution to estimate a candidate set’s pass@k, then approximately solves the size-k subset-selection problem.The exact selection problem is NP-hard, motivating a tractable surrogate and greedy optimization.Source: §3 · p. 5, §3.2 · p. 7

3.1 Probabilistic model

  • A latent Gaussian score combines predicted rollout quality, a shared random effect tied to representations, and rollout-specific noise; thresholding the score models correctness.Aligned representations induce positively correlated scores and hence outcomes that can succeed or fail together.Source: §3.1 · p. 5
  • Because the exact likelihood requires an intractable expectation, the authors use a local second-order closed-form surrogate and train the quality mapping on joint outcomes.With limited data, they keep representations fixed and use a linear quality mapping.Source: §3.1 · pp. 5–6

3.2 Subset selection

  • The selection surrogate rewards high quality and broad representation coverage while penalizing aligned weighted vectors; it is monotone and submodular, so greedy selection has an approximation guarantee.The exact selection objective is NP-hard, while the surrogate enables the stated greedy guarantee.Source: §3.2 · p. 7

3.3 Relation to Linear Probe and RepExp

  • When quality dominates, PKTP behaves like Linear Probe; as the selected set grows, correlation penalties can shift selection toward RepExp-style diversity.The correlation strength controls how quickly this transition occurs as k grows.Source: §3.3 · p. 8

4 Experiments

The experiments compare pruning policies across eight reasoning models, using held-out evaluation to measure pass@k. PKTP outperforms the Linear Probe and RepExp most clearly at moderate k, with results also depending on cutoff length and pool size.

4.1 Setup

  • The evaluation covers eight open-weight reasoning models and uses eight-fold cross-validation repeated 20 times, with each problem scored by a model not trained on it.Models use temperature 0.8, top-p 0.95, and a 32,768-token full budget; correctness is judged by answer equivalence to the reference.Source: §4.1 · pp. 8–9

4.2 Experimental Evaluation and Comparison

  • The Linear Probe’s advantage at k = 1 fades as k grows, while RepExp generally improves with k but can trail Random at moderate k.The probe can cluster similar rollouts that are wrong together; RepExp can favor unusual outliers unlikely to succeed.Source: §4.2 · p. 10, Figure 4 · p. 10, Figure 5 · p. 11
  • At moderate k, PKTP outperforms both the Linear Probe and RepExp; on DeepSeek-R1-Distill-Qwen-1.5B at k = 4, it leads by 3.0 and 3.2 points, respectively.PKTP also stays above Random on every model for k ≥2.Source: Figure 3 · p. 9, §4.2 · pp. 10–11

4.3 Effect of the cutoff threshold

  • For T ≥4000, PKTP outperforms all other methods on both tested models; its pass@4 rises from 63.5% to 65.9% on Qwen3-1.7B and from 57.4% to 61.7% on DeepSeek-R1-Distill-Qwen-7B as T grows from 2000 to 8000.The Linear Probe on Qwen3-1.7B instead drifts from 64.2% to 62.3% and ends below Random.Source: §4.3 · pp. 11–12, Figure 6 · p. 12

4.4 Effect of the pool size

  • PKTP leads on Qwen3-1.7B at every tested pool size and on DeepSeek-R1-Distill-Qwen-7B once N reaches 80.On the latter model, as N grows, the Linear Probe falls from 58.1% to 56.4% and remains below Random, while PKTP rises from 58.8% to 59.7%.Source: §4.4 · pp. 12–13, Figure 7 · p. 12

5 Conclusion

PKTP selects partial rollouts by directly optimizing pass@k, using a latent Gaussian model and a greedy approximation. Across eight reasoning models, it outperforms quality- and diversity-based pruning, with the largest gains at moderate k.

  • PKTP models the joint distribution of rollout outcomes with a latent Gaussian model, approximates the likelihood in closed form, and greedily selects the kept set.The greedy procedure has an approximation guarantee on the surrogate objective.Source: §5 · p. 13
  • Across eight reasoning models, PKTP outperforms both quality-based and diversity-based pruning, with the largest gains at moderate k.The conclusion attributes this pattern to the biases of the two baseline families.Source: §5 · p. 13
  • Because PKTP can use any representation of partial rollouts, better representations of rollout correctness are a natural direction for future work.Source: §5 · p. 13

AI use statement

The authors report using generative AI tools for editing and formatting, figure scripts, implementation assistance, and parts of proofs. They state that no data or experimental results were AI-generated and that authors reviewed the assisted content.

  • Generative AI tools helped edit and polish the write-up and format the LaTeX source.Source: AI use statement · p. 13
  • The tools also helped write figure-plotting scripts and implement parts of the code.Source: AI use statement · p. 13
  • The tools assisted with parts of the proofs, but not with generating data or experimental results; authors reviewed and verified all assisted content.Source: AI use statement · p. 13

A Derivations

The derivations connect the latent Gaussian outcome model to a closed-form pass-probability surrogate and establish properties that support greedy selection. They also relate the equal-quality limit to RepExp and give a greedy approximation guarantee.

A.1 The likelihood; B.1.1 Proof of Theorem 1

  • Conditioning on the latent variable makes rollout outcomes independent; Gaussian symmetry yields likelihood factors involving Φ, while a quadratic approximation and Gaussian integration give a closed form.The likelihood derivation takes an expectation over the latent variable, and the proof of Theorem 1 drops the accumulated Taylor remainder.Source: §A.1 · p. 19, §B.1.1 · p. 21

A.2 Rank-one greedy updates

  • Rank-one matrix updates support greedy selection, with each step costing O(Nd + d^2) when all h_i are retained in memory.Source: §A.2 · pp. 19–20

A.3 Greedy marginal in the equal-quality limit

  • When every quality score equals a common value and the penalty term is dropped, the candidate with smallest Δj(S) ranks highest, yielding the RepExp rule with regularizer τ = 1/κ.Source: §A.3 · p. 20

B.1.2 Proof of Corollary 1

  • The selected-set failure surrogate follows by integrating out outcomes outside the set and expressing the result through the set’s aggregate terms.Source: §B.1.2 · p. 22

B.2.1 Submodularity of Exact pass@k Probability

  • Exact pass probability is normalized, monotone, and submodular for any joint distribution of correctness events, without a latent-variable or conditional-independence assumption.Source: §B.2.1 · p. 22

B.2.2 Submodularity of the Approximate Pass Probability

  • The Gaussian-model surrogate is also shown to be monotone and submodular, a property not guaranteed merely by approximating a submodular function.Source: §B.2.2 · pp. 23–25

B.2.3 Greedy Guarantee

  • Since the surrogate pass probability is normalized, monotone, and submodular, the greedy solution satisfies the stated guarantee of Theorem 2.Source: §B.2.3 · p. 25

B.3 Proof of Proposition 1

The proof establishes NP-hardness by reducing independent set to minimizing the surrogate objective, under equal integer quality scores and binary feature vectors. Its construction makes independent sets the minimizers when such sets of size k exist.

  • The optimization is NP-hard even when all quality scores are the same integer and all feature vectors are binary, via a reduction from independent set.Source: §B.3 · p. 26
  • The construction encodes graph adjacency in vector overlaps, making the Gram matrix for a selected set equal to ΔI_k + A_S.The determinant term is maximized exactly by independent sets of size k.Source: §B.3 · p. 26
  • A uniform bound on the penalty term preserves the independent-set separation, so minimizers correspond to independent sets when one of size k exists.The constructed instance has polynomial size, completing the reduction.Source: §B.3 · pp. 26–27

B.4 Proof of Proposition 2

Under weak correlations, the failure-probability objective and its surrogate have the same expansion. Appendix C then characterizes how the selection rule moves from quality-based selection at small budgets toward diversity-based selection at large budgets under stated assumptions.

B.4 Proof of Proposition 2

  • Scaling representations by ρ while holding qualities fixed makes the true failure probability and its surrogate share the same expansion as correlations weaken.The argument fixes quality scores and scales representations; ρ = 1 recovers the Section 3.1 model.Source: §B.4 · p. 28

C Selection Behavior Across Budgets

  • For fixed ρ, surrogate greedy selection moves between the Linear Probe and RepExp behaviors as k grows, with quality, log-determinant diversity, and penalty terms jointly determining intermediate selections.At one end it selects highest-quality candidates; at the other, RepExp greedily increases a log-determinant objective.Source: C Selection Behavior Across Budgets · p. 29

C.2 Assumptions

  • The analysis assumes higher-quality fixed spikes above an independently drawn bounded-quality bulk and representations h_i = ρv_i independent of bulk qualities.Large-budget results additionally require a family of sets whose representations cover directions rather than concentrate in a few.Source: §C.1 · p. 29

C.1 Overview of Appendix C

  • At small budgets, PKTP selects the top-k quality candidates when the quality gap exceeds a constant multiple of kρ^2; at large budgets, its averaged objective approaches RepExp plus a term for high-quality spikes.The large-budget conclusion concerns objective values over diverse sets, not identical selected sets.Source: §C.1 · p. 29
  • One explicit model satisfies both regimes: the first three greedy selections are the ordered spikes, while the large-budget conclusions hold as the budget grows.The construction uses one fixed quality distribution and one fixed ρ.Source: §C.1 · p. 29

C.3 Small budget k: a quality gap yields the probe-style rule

A sufficiently large quality gap makes the surrogate objective select the same candidates as the probe-style rule at small budgets. When the gaps persist across budgets, greedy selection follows the probe-style order; separated high-quality spikes provide such gaps even as the pool grows.

C.3 Small budget k: a quality gap yields the probe-style rule

  • The quality gap measures the loss from replacing the weakest top-k candidate with the strongest outside candidate.Adding a candidate changes the objective by its own failure term, up to an error that grows with set size and vanishes as ρ approaches zero.Source: §C.3 · p. 30
  • Under (A1) and (A2), if kρ^2 ≤ 1 and Γ_k > (1 + 3Λ^2)kρ^2, then Stop_k is the unique size-k minimizer of log eL, and surrogate greedy returns Stop_k after k steps.Source: §C.3 · p. 31
  • When the required gap holds at every budget through k, surrogate greedy adds the top-quality candidates in the probe-style order.The spikes create persistent gaps, so the first r selections are the ordered spikes for every bulk-quality realization when the stated Corollary 2 condition holds.Source: §C.3 · p. 32
  • Separated spikes keep this result applicable as pools grow, whereas with bulk qualities alone the kth-to-next quality gap typically shrinks and eventually fails the theorem’s condition.The spike gaps do not depend on the bulk pool size.Source: §C.3 · p. 32

C.4 Large budget k: the expected objective reduces to the RepExp objective

For large budgets, averaging PKTP’s objective over bulk qualities reduces its set-dependent part to a spike term and a log-determinant diversity objective, up to an error shrinking with k. Under an explicit construction, the paper exhibits both this large-budget behavior and small-budget selection of the highest-quality spikes.

C.4 Large budget k: the expected objective reduces to the RepExp objective

  • The large-budget result assumes fixed ρ, bounded bulk qualities, and diverse size-k sets whose representations cover directions and whose bulk contributions nearly cancel.The claims hold for each fixed ρ as k grows, not necessarily when ρ shrinks with k.Source: §C.4 · p. 32
  • As k grows, the averaged objective over diverse sets depends on the kept spikes and log det(Id + κG_S), with other terms constant and an error that shrinks.The result is about averaging over bulk qualities for a fixed set, not the objective for a single bulk-quality draw.Source: §C.4 · pp. 32, 35
  • For sufficiently large k, every minimizer of the averaged objective keeps all spikes when the candidate-set family permits a bulk candidate to be replaced by any missing spike.This closure condition is not automatic for families defined only by the diversity assumptions.Source: §C.4 · p. 34
  • Once all spikes are kept, minimizing the averaged objective nearly maximizes the RepExp objective among diverse sets containing them.D_κ is the RepExp objective with τ = 1/κ.Source: §C.4 · p. 35

C.5 An explicit model with both regimes

  • An explicit model uses three fixed spikes and a cyclic sequence of antipodal bulk directions to construct diverse candidate sets with the required closure property.Pairing opposite directions cancels the vector sum while preserving covariance along each axis.Source: §C.5 · pp. 35–37
  • In that model, one fixed quality distribution and one fixed ρ yield ordered spike selections for every k ≤ 3 and the large-budget result as k(N) grows while k(N)/N vanishes.The bulk qualities retain fixed nonzero variance as the pool grows.Source: §C.5 · pp. 36, 38

D Benchmark-Stratified Results from Section 4.2

This section reports the pass@k results behind Figure 3 and breaks them down by benchmark, with one table per model. Entries are average pass@k percentages across repeated cross-validation runs, with standard errors.

E Results on different pool sizes

The appendix varies the rollout-pool size, specifies the representation used by the method, and situates PKTP among related pruning and pass@k objectives.

E Results on different pool sizes

F Representation details

  • For each rollout, the representation is the normalized exponential moving average of last-layer hidden states after T tokens, or all tokens if the rollout ends earlier.The EMA uses α = 10^-3.Source: F Representation details · p. 48

G Extended related work

  • Confidence, reward-model, and probe methods score rollouts individually, while RepExp selects for unexplained representations without quality scores; PKTP models outcome dependence across rollouts.Source: G Extended related work · p. 48
  • Unlike post-training methods that optimize pass@k, PKTP leaves the model unchanged and chooses which independently sampled partial rollouts to continue.Source: G Extended related work · p. 48
  • PKTP learns the joint outcome distribution from rollout representations and optimizes a closed-form surrogate that preserves submodularity of the exact objective.Source: G Extended related work · p. 48
  • Related work also uses probability of at least one success in retrieval and Bayesian optimization, while other approaches impose hand-set quality–diversity trade-offs.Source: G Extended related work · p. 48

Related papers

Loading 2610.31667v1…