PKTP thought pruning — Pass@k-optimal pruning of partial reasoning paths — What does it mean?

The news. On October 22, 2026, researchers at the University of Southern California posted Direct Pass@k Optimization for Thought Pruning. For each question the model samples N rollouts, sees only the first T tokens of each, and must pick k to finish. The paper tests eight open-weight reasoning models on competition math. Read the paper →

Picture a stack of 128 scratch cards. You may scratch only 4 completely, and scratching costs tokens. Each card shows a small corner, the first 4000 tokens of a reasoning attempt. The aim is not to find the four cards that each look best. It is to make it likely that at least one of your four wins. That is the pass@k goal, and it changes which cards you should pick. The quality-only rule picks the four best-looking corners. If those corners look alike, they probably hide the same pattern, so they win or lose together and your four scratches act like one.

PKTP turns the look-alike problem into a model. Each rollout gets a hidden score. The rollout is correct when the score is at least zero. The score has a quality part, learned from the prefix's pooled hidden state, and a shared random part. Two rollouts whose prefix vectors are similar share more of that random part, so their covariance is the dot product of their vectors. In the paper's model, similar prefixes mean correlated outcomes, and correlated paths add little to the chance that at least one succeeds. Quality is how likely a path is to work (see how sampling spreads over many continuations); redundancy is how much a path repeats a path already kept.

The exact probability that every kept rollout fails has no closed form, so the paper uses a second-order expansion as a differentiable surrogate. The authors say the approximation is local, so it is a training signal, not an exact bound. Choosing the best k-subset exactly is NP-hard, but the surrogate is monotone and submodular, so greedy selection keeps at least 1 − 1/e of the best possible surrogate value. Each greedy step adds the rollout that lowers the failure probability most, given the ones already kept. One greedy pass gives nested subsets, so k can be chosen at inference time without retraining.

The reported gains are real but modest, and they are largest at moderate k. At N = 128 rollouts and a 4000-token cutoff, PKTP beat the linear probe by 3.0 points and RepExp by 3.2 points for DeepSeek-R1-Distill-Qwen-1.5B at k = 4, and beat the linear probe by 3.3 points for DeepSeek-R1-Distill-Qwen-7B at k = 4. The cutoff sweep shows why a quality-only rule is fragile: for Qwen3-1.7B at k = 4, the linear probe fell from 64.2% to 62.3% as the cutoff moved from 2000 to 8000 tokens, ending below random choice, while PKTP rose from 63.5% to 65.9%. A larger cutoff gives the probe more text but not more diversity.

Worked example, first the coverage effect, then the cost. Say every kept path succeeds with probability 0.5 and you keep two (illustrative). If they are perfectly correlated, pass@2 stays at 0.50. If they are independent, pass@2 is 1 − 0.5 × 0.5 = 0.75. Now the cost, using the paper's settings: N = 128, prefix T = 4000, full budget B = 32,768 tokens, k = 4 (our own hypothetical arithmetic, not a measured result; the paper does not report compute savings). Finishing every rollout costs 128 × 32,768 = 4,194,304 tokens. Pruning costs 4 × 32,768 + 124 × 4,000 = 131,072 + 496,000 = 627,072 tokens, about 6.7 times fewer. The two calculations are separate: the token count is the same whichever four paths you keep, and the pruning rule only decides how much coverage those tokens buy. See temperature and why samples differ for why 128 attempts are not 128 copies.

Goes deeper in: LLM Internals → Text Generation → Top-K and Top-P Sampling

Related explainers

• Sharpening tax in post-training — why pass@1 and pass@k coverage can move in opposite directions