Choosing which AI tries are worth finishing
Original title: Direct Pass@$k$ Optimization for Thought Pruning: Balancing Quality and Diversity in Test-Time Reasoning
Read by Roger, a synthetic voice
The paper tries to save computer effort when an AI makes many attempt paths. The authors keep only the partial paths that together give the best chance of a correct finished answer. They report this beats older rules that keep only strong paths or only different paths on eight models.
Our take: I’d call the evidence fairly strong: the authors report better results across eight models. For an ordinary person, the payoff is indirect, since it mainly changes how AI builders spend computer effort.
- Running every AI attempt to the end can cost a lot of computer work.
- The method keeps partial attempts that together give the best chance of a correct answer.
- The authors report better results than older quality-only and diversity-only rules on eight models.
14 claims checked against the source Read the paper on arXiv PDF
Checked against the source
Every statement in this brief had to match a word-for-word passage from the source; code checks each match, and a second AI model reviews the whole brief for errors. A brief that fails is not published. These are the 14 passages it rests on.
Show the 14 passages
Reasoning models may generate many candidate solutions to find a correct answer.
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 costly.
Running all candidate solutions to completion can be computationally demanding
The paper's title is Direct Pass@k Optimization for Thought Pruning: Balancing Quality and Diversity in Test-Time Reasoning.
Direct Pass@kOptimization for Thought Pruning: Balancing Quality and Diversity in Test-Time Reasoning
Pass@k is the probability that at least one kept partial solution becomes correct when completed.
it chooses the pool to maximize the probability that at least one of its partial solutions, when continued to completion, yields a correct answer.
Pruning stops after a short prefix and spends further computation only on surviving candidates.
recent work prunes the candidate pool after generating only a short prefix, committing further computation only to the surviving candidates.
Only selected rollouts continue within a budget.
Only k of the N rollouts (ringed) continue past the cutoff T , up to the budgetB.
The objective rewards quality and penalizes redundancy.
Because this objective simultaneously rewards high-quality partial solutions and penalizes redundancy among them
The authors report better pruning results on eight open-weight reasoning models.
On eight open-weight reasoning models, our method outperforms existing quality-based and diversity-based pruning strategies.
Gains are clearest with a moderate number of kept partial solutions.
The gains are particularly pronounced when the pruning procedure keeps a moderate number of partial solutions for continued generation
When many tries are kept, the method's advantage narrows.
At large k, however, the performance of PKTP and RepExp gets closer.
The authors attribute weaker large-k behavior to weak correctness signals in partial rollout embeddings.
We believe this is due to the weak correctness signal in the embeddings of the partial rollouts
Finding the exact best subset is NP-hard.
Minimizing log eL(S)over all S⊆ [N ]with |S| = k is NP-hard.
The authors use an approximate greedy procedure.
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.
Pruning can reduce wasted computation by spending more effort only on survivors.
committing further computation only to the surviving candidates.