Structured Pruning for Diverse Best-of- N Reasoning Optimization (2025)

Abstract: Model pruning in transformer-based language models, traditionally seen as a means of computational savings, can enhance the model’s reasoning capabilities. In this work, we uncover the surprising phenomenon that the selective pruning of certain attention heads leads to improvements in reasoning performance, particularly on challenging tasks. Motivated by this observation, we propose SPRINT, a novel contrastive learning framework that dynamically selects the optimal head and layer to prune during inference. By aligning question embeddings with head embeddings, our approach identifies those pruned-head configurations that result in more accurate reasoning. Extensive experiments on the MATH dataset demonstrate that our method significantly outperforms traditional best-of- N and random head selection strategies on the MATH500 and GSM8K datasets.

## introduction language models (lms) have recently made remarkable progress in understanding and reasoning with human language, attracting considerable research attention @xcite @xcite @xcite . this unprecedented evolution opened up a wide range of practical applications, including content generation @xcite @xcite , virtual assistance @xcite @xcite , population simulation @xcite , and problem solving @xcite @xcite . despite these advances, the reliability and correctness of lm-generated answers remain a significant concern. analyses of their output frequently reveal faulty reasoning and factual inaccuracies @xcite @xcite @xcite . this issue is particularly pronounced in tasks that require advanced reasoning, such as automated theorem proving @xcite , mathematical problem solving @xcite , or heuristic discovery @xcite . llms often struggle to produce accurate responses in these scenarios in a single pass. to address this challenge, iterative generation strategies have been employed to refine and select the most appropriate response, often by combining language generation with aggregation or search techniques such as best-of-n sampling @xcite or monte carlo tree search @xcite . throughout this paper, we refer to these iterative approaches as reasoning methods. current reasoning methods typically leverage the diversity induced by the llm's decoding process through parameters such as temperature, top-k, and top-p to generate multiple candidate answers @xcite . although this technique introduces textual diversity and paraphrasing, it does not capture fundamentally different viewpoints, since all candidates are produced by the same underlying model with similar reasoning capabilities. one potential remedy is to employ a variety of lms to generate answers; however, this approach is often computationally expensive and memory inefficient. at the same time, recent research on model pruning has shown that comparable performance can be achieved by selectively disabling certain heads in the transformer architecture @xcite . removing an attention head does not necessarily weaken the overall ability of the model because a head may introduce noise into the decision-making process, leading to incorrect answers. in fact, our empirical observations suggest that pruning certain heads can improve the reasoning methods' performance on specific categories of questions; section 3.2 discusses this phenomenon in more detail. this finding motivates our approach: we aim to construct a pool of pruned models from one base model with diverse reasoning capabilities without incurring additional computational or memory overhead. contributions. we begin by analyzing and providing new insights into the surprising phenomenon that pruning attention heads can improve the performance of the language model on the math benchmark. motivated by these insights, we introduce the structured pruning for diverse reasoning optimization (sprint) framework. sprint first learns a set of embeddings for each layer-attention head pair using contrastive learning, guided by a novel diversity-promoted contrastive loss. this loss function encourages the question embedding to move closer to the embeddings of heads whose pruning leads to correct answers for that question. finally, we leverage sprint to enhance best-of-n sampling by pruning the top n nearest layer-model configurations, determined by the distance between the learned embeddings and the question embedding. experiments on the math500 and gsm8k datasets show that our method outperforms traditional best-of-n sampling approaches, achieving higher accuracy without introducing inference-time overhead. our paper unfolds as follows: section 2 discusses related work on reasoning in lms. section 3 investigates the impact of pruning individual attention heads in transformer models, revealing that certain heads can be removed without performance loss and, in some cases, even improve accuracy. section 4 delineates our sprint framework for matching heads for pruning with input questions, and section 5 presents the extensive numerical results of the mathematical reasoning task. ## related works a straightforward decoding process may struggle to produce accurate solutions for complex reasoning tasks. to address this, self-consistency (sc) improves reliability by generating multiple outputs and selecting the final response by majority voting @xcite . a similar approach, bestof-n sampling, employs a reward model or function to select the response with the highest reward @xcite . while both techniques enhance output quality, they also increase computational cost proportional to the number of samples. to explore reasoning pathways more effectively, tree-search-based methods have been introduced, such as tree-of-thought @xcite , monte carlo tree search (mcts) @xcite @xcite , and forestof-thought @xcite . @xcite highlights that tree-based search strategies are more effective in discovering correct solutions than parallel sampling, particularly for complex problems. beyond reasoning strategies, scaling inferencetime computation also plays a crucial role in improving output quality. @xcite shows that the accuracy on the math-500 benchmark improves as the test-time computation (that is, the number of generations per problem) increases for methods such as best-of-n, beam search, and diverse verifier tree search (dvts). similarly, @xcite achieved an average accuracy of 53.3% on the usa math olympiad (aime) by conducting extensive mcts rollouts. our work is related to recent research on interventions in the internal state of language models @xcite @xcite @xcite . although prior approaches often focus on learning intervention vectors to add to model activations, instead we investigate the effect of excluding specific components of the internal state of the model. ## does pruning an attention head improve lm's performance? we consider a transformer-based model, where each layer contains h attention heads. in each layer ℓ, the outputs of the h attention heads are denoted as {a ℓ h } h @xmath0, where each a ℓ @xmath1. the concatenated vector, a ℓ = (a ℓ 1 , . . . , a ℓ h ) ∈ r d with @xmath2, is then passed to a linear layer to produce the output of a multi-head attention block. attention head pruning. in our study, we define head pruning as zeroing out the output vector a ℓ h of a specific attention head before the concatenation step, as illustrated in figure 1 . in modern decoderonly architectures (), this is equivalent to pruning the input of o proj layer. in this section, we detail our experimental procedure for attention head pruning and present a surprising empirical phenomenon that motivates our subsequent work. ## conclusion in this work, we demonstrated that pruning attention heads, a technique traditionally viewed as harmful, can enhance the reasoning capabilities of transformer-based language models. our study revealed that the selective removal of specific heads mitigates the propagation of redundant or noisy signals and enhances overall model performance. we developed a contrastive learning framework to harness this phenomenon that dynamically selects the optimal head and layer to prune at inference time. experimental results confirm that our approach significantly outperforms traditional best-of-n methods and random head selection strategies.