Reasoning model pruning method balances solution quality and diversity
Direct Pass@$k$ Optimization for Thought Pruning: Balancing Quality and Diversity in Test-Time Reasoning
Artificial IntelligenceComputation and Language
Summary
Generating many possible answers from a reasoning AI can use a lot of computing power, especially when checking all answers fully. The authors found a way to quickly pick a smaller set of partial answers to finish, aiming to have at least one correct final answer. Their method balances between choosing high-quality partial answers and ensuring the chosen answers are different from each other. Tests on several AI models showed this approach works better than previous pruning methods, especially when keeping a moderate number of partial answers.
What this means in practice
- •For ai developers: Reduce computation during AI reasoning by selecting a balanced set of partial solutions to complete for better end results.
- •For software engineers: Improve efficiency in systems that generate many candidate answers by pruning early with a method balancing solution quality and coverage.
Authors
Mahdi Salmani, Adel Javanmard, Devansh Gupta, Meisam Razaviyayn
Abstract
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.