Structured Pruning for Diverse Best-of-$N$ Reasoning Optimization
Authors: Hieu Trung Nguyen, Bao Nguyen, Viet Anh Nguyen
Accepted at: acl25

Structured Pruning for Diverse Best-of-N Reasoning Optimization

Hieu Trung Nguyen, Bao Nguyen, Viet Anh Nguyen
The Chinese University of Hong Kong
{thnguyen, nbnguyen, nguyen}@se.cuhk.edu.hk

Abstract
Model pruning in transformer-based language models, traditionally viewed as a means of achieving computational savings, can enhance the model’s reasoning capabilities. In this work, we uncover a surprising phenomenon: 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, SPRINT identifies those pruned-head configurations that result in more accurate reasoning. Extensive experiments demonstrate that our method significantly outperforms traditional best-of-N and random head selection strategies on the MATH500 and GSM8K datasets.

1 Introduction
Language models (LMs) have recently made remarkable progress in understanding and reasoning with human language, attracting considerable research attention (Brown et al., 2020; Hoffmann et al., 2022; Chowdhery et al., 2022; Rae et al., 2021; Raffel et al., 2019). This unprecedented evolution opened up a wide range of practical applications, including content generation (Li et al., 2024a; Liu et al., 2023b; Achiam et al., 2023), virtual assistance (Sezgin, 2024; Doe et al., 2023; Garcia et al., 2023), population simulation (Bui et al., 2025), and problem solving (Ahn et al., 2024; Imani et al., 2023; Lewkowycz et al., 2022).

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 (Ji et al., 2023; Xu et al., 2024; Nguyen et al., 2025a; Jiang et al., 2025; Li et al., 2024b). This issue is particularly pronounced in tasks that require advanced reasoning, such as automated theorem proving (Wu et al., 2022), mathematical problem solving (Trinh et al., 2024), or heuristic discovery (Romera-Paredes et al., 2024). 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 (Stienenon et al., 2020) or Monte Carlo Tree Search (Xie et al., 2024). 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 (Stienenon et al., 2020; Xie et al., 2024). 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 (Ma et al., 2023; Wang et al., 2019). 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.

2 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 (Wang et al., 2022). A similar approach, best-of-N sampling, employs a reward model or function to select the response with the highest reward (Stienenon et al., 2020). 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 (Yao et al., 2024), Monte Carlo Tree Search (MCTS) (Feng et al., 2023; Zhang et al., 2024; Guan et al., 2025), and Forest-of-Thought (Bi et al., 2024). Damani et al. (2024) highlights that tree-based search strategies are more effective in discovering correct solutions than parallel sampling, particularly for complex problems.

Beyond reasoning strategies, scaling inference-time computation also plays a crucial role in improving output quality. Beeching et al. (2024) 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, Guan et al. (2025) 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 (Wu et al., 2024; Li et al., 2023; Nguyen et al., 2025b; Liu et al., 2023a; Nguyen et al., 2025c). 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.

3 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 \( \ell \), the outputs of the \( H \) attention heads are denoted as \( \{a_h^{\ell}\}_{h=1}^{H} \), where each \( a_h^{\ell} \in \mathbb{R}^d \). The concatenated vector, \( a^{\ell} = (a_1^{\ell}, \ldots, a_H^{\ell}) \in \mathbb{R}^D \) with \( D = H \times d \), 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^{\ell} \) of a specific attention head before the concatenation step, as illustrated in Figure 1. In modern decoder-only 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.

3.1 Experimental Procedure

To systematically examine the effect of individual head pruning, we follow the procedure below:

1. Layer and Head Selection: We select \( L = 4 \)
Figure 1: Illustration of our attention head pruning technique. Most of the plot depicts the standard multi-head attention mechanism in transformer-based models. Our approach differs by pruning a specific attention head \( a_j \)—that is, we zero out its output vector before concatenating the head outputs and passing the result to the \( o_{proj} \) layer.

layers from the model, yielding a total of \( LH \) attention heads. For each head, we create a variant of the base model with that specific head pruned, resulting in \( LH \) model variants in addition to the base model. We denote this model pool as \( \mathcal{M} = \{\mathcal{M}_{base}\} \cup \{\mathcal{M}_h^j\} \).

2. Measuring the Impact: For each model \( m \in \mathcal{M} \), we generate solutions on the MATH500 dataset, which consists of \( n = 500 \) samples. The predictions are represented by a binary accuracy vector \( z \in \{0,1\}^n \), where \( z_i = 1 \) if sample \( i \) is solved correctly and 0 otherwise. Evaluating all models yields a matrix \( Z \in \{0,1\}^{n \times 4H} \) of accuracy values.

In this experiment, we study two models specialized for mathematical reasoning, Qwen2.5-Math-1.5B-Instruct and Qwen2.5-Math-7B-Instruct, as well as two general-purpose models, Qwen2.5-7B-Instruct, and Meta-Llama-3-8B-Instruct. Our experiments involve evaluating model performance with and without head pruning on 500 samples from the MATH500 dataset (Hendrycks et al., 2021).

3.2 Empirical Observations

Our empirical results, shown in Tables 2, 3, 4, and 1, reveal several important and sometimes surprising findings about attention head pruning. Contrary to the common belief that all attention heads are essential for optimal model performance, we find that pruning individual heads often does not reduce accuracy, and in many cases, actually increases it.

The impact of pruning depends strongly on both the specific head and its layer: while some heads are crucial (with pruning causing sharp drops in accuracy), many others are unnecessary, and removing them can leave performance unchanged or even lead to notable improvements. For instance, in the ‘Counting & Probability’ category, pruning head 0 in layer 15 of Qwen2.5-Math-1.5B-Instruct increased accuracy by 16% over the unpruned baseline.

We further illustrate this trend in Figure 2, which displays violin plots of additive performance gains (the difference between the best pruned head accuracy and the baseline accuracy) across all subcategories for each model. These plots show that head pruning consistently produces positive gains for all evaluated models. The specialized mathematical models (Qwen2.5-Math-1.5B-Instruct and Qwen2.5-Math-7B-Instruct) exhibit higher median gains and more frequent large improvements, particularly in domains such as ‘Counting & Probability’ and ‘Geometry’. In contrast, the general-purpose models (Qwen2.5-7B-Instruct and Meta-Llama-3-8B-Instruct) show more modest but still positive and consistent benefits from pruning.

Overall, these patterns suggest that some attention heads contribute redundant or even harmful information, and that pruning can act as a useful form of regularization. Our findings highlight the significant redundancy present in multi-head attention and indicate promising directions for simplifying models and improving efficiency through targeted
head pruning.

4 Methodology

![Distribution of additive performance gains (best pruned accuracy minus no-prune accuracy) for Qwen2.5-Math-1.5B-Instruct, Qwen2.5-Math-7B-Instruct, Qwen2.5-7B-Instruct, and Meta-Llama-3-8B-Instruct. Each violin illustrates the improvement across all subject-layer pairs.](page_186_320_495_312.png)

Figure 2: Distribution of additive performance gains (best pruned accuracy minus no-prune accuracy) for Qwen2.5-Math-1.5B-Instruct, Qwen2.5-Math-7B-Instruct, Qwen2.5-7B-Instruct, and Meta-Llama-3-8B-Instruct. Each violin illustrates the improvement across all subject-layer pairs.

Motivated by these counterintuitive findings, we propose a framework called Structured Pruning for Diverse Reasoning Optimization (SPRINT) to leverage selective head pruning to improve reasoning performance. At the core of SPRINT is a contrastive learning approach that aligns question embeddings with attention head embeddings to identify the heads most beneficial to prune. Initially, every head is associated with a learnable vector from a set \( \mathcal{V} = \{v_j\}_{j=1}^{LH} \), where each \( v_j \in \mathbb{R}^p \) is randomly initialized. A sentence embedding model, denoted by \( \phi \), processes each question \( x_i \), after which a linear transformation \( \theta \) produces the question embedding \( q_i = \theta(\phi(x_i)) \). For each question, we define the set of pruned-head models as \( \mathcal{M} = \{1, \ldots, LH\} \) and partition this set into a positive subset \( \mathcal{M}_i^+ = \{j : z_{ij} = 1\} \) containing heads that result in correct answers to question \( i \), and a negative subset \( \mathcal{M}_i^- = \{j : z_{ij} = 0\} \).

To further quantify the correspondence between different pruned-head models, we introduce a similarity measure between any two heads \( j \) and \( k \) as follows:

\[
s_{jk} = \frac{1}{n} \sum_{i=1}^n \mathbb{1}(z_{ij} = z_{ik}) \in [0, 1],
\]

which computes the fraction of questions for which both heads produce the same outcome. We jointly optimize the transformation \( \theta \) and the head embeddings \( \mathcal{V} \) with the following loss function:

\[
\mathcal{L} = \frac{1}{n} \sum_{i=1}^n - \log \left( \frac{\sum_{j \in \mathcal{M}_i^+} \exp \left( - \|q_i - v_j\|_2^2 \right)}{\sum_{j=1}^{LH} \exp \left( - \|q_i - v_j\|_2^2 \right)} \right) 
- \lambda \sum_{j<k} s_{jk} \|v_j - v_k\|_2^2.
\]

Our joint optimization loss \( \mathcal{L} \) consists of two components. The first term aims to align the question embedding \( q_i \) with the embeddings \( v_j \) of those pruned-head models that yield the correct answer. The loss employs a softmax formulation for each question, over the negative squared Euclidean distances between \( q_i \) and the head embeddings. This loss encourages the model to produce a lower distance, and thus a higher similarity, for those heads in the positive set \( \mathcal{M}_i^+ \), which results in a correct answer.

The second term of the loss addresses the need for diversity among head embeddings. By summing the pairwise distances weighted by the similarity measure \( s_{jk} \), it acts as a regularizer that encourages distinctiveness between head embeddings, especially for those heads that tend to yield similar outcomes. This term effectively pushes similar heads apart in the embedding space to reduce redundancy and promote a richer, more discriminative representation. Overall, combining these two terms helps align the question representation with high-performing pruned models while maintaining a diverse set of head embeddings.

Inference. At inference time, we obtain the embedding for the new question \( x \) by feedforwarding \( q = \theta(\phi(x)) \). The optimal head to prune for a given test question is identified by selecting the head whose embedding minimizes the squared Euclidean distance to the question embedding, i.e.,

\[
j^* = \arg \min_{j \in \mathcal{M}} \|q - v_j\|_2^2.
\]

We integrate best-of-\( N \) reasoning into our framework to enhance the performance by selecting the top-\( N \) pruned-head configurations based on their proximity to \( q \). We obtain candidate answers via greedy decoding with a temperature of 0.

Overall, our unified methodology demonstrates that pruning the attention heads, guided by a contrastive alignment strategy, can improve model performance. Thus, the SPRINT framework identifies which heads benefit from pruning and supports a
dynamic simulation approach to generate optimal predictions.

5 Numerical Experiments

![Schematic overview of our head pruning framework.](page_186_120_1077_349.png)

Figure 3: Schematic overview of our head pruning framework.

![Average Pass@N on MATH500 and GSM8K dataset.](page_186_624_1077_181.png)

Figure 4: Average Pass@N on MATH500 and GSM8K dataset.

Datasets. For our experiments, we use the MATH and GSM8K datasets. We provide more details in the appendix.

Base Models. As a proof of concept, we select Qwen2.5-Math-1.5B-Instruct as our generation model. This choice is motivated by its strong empirical gains observed in Figure 2. All experiments are performed using the original, unquantized version of this model. We consider five layers {5, 10, 15, 20, 25} for pruning.

Performance metrics. To assess performance, we use the Pass@N metric. Specifically, each method is used to sample N candidate answers, and the accuracy is computed based on the best answer as determined by an oracle.

Baselines. We compare our methods against two baselines: multinomial sampling, commonly used for Best-of-N Reasoning, and a random head selection strategy. We perform Best-of-N evaluations using three temperature values for multinomial sampling: 0.1, 0.6, and 0.9. For the random head selection strategy, we use a greedy approach to identify the heads that solve the most samples in the training set. During testing, we randomly select N heads and measure Pass@N.

The project’s repository is available at https://github.com/HieuNT91/attention_pruning.

5.1 Quantitative Results

Figure 4 summarizes our primary findings. Our method consistently outperforms the baselines for nearly all values of N across the MATH and GSM8K datasets. Notably, the greatest improvements are observed at lower values \( N \in \{1, 2, 3, 4\} \), while the performance gains become more gradual as N increases.

6 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.
Limitations

Due to the space limitations of a short paper, our work does not experiment with a more diverse range of model architectures and sizes. Our work primarily focuses on the mathematical domain, but it can be extended to other scientific domains that require LMs’ reasoning.

References

Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, and 1 others. 2023. Gpt-4 technical report. arXiv preprint arXiv:2303.08774.

Janice Ahn, Rishu Verma, Renze Lou, Di Liu, Rui Zhang, and Wenpeng Yin. 2024. Large language models for mathematical reasoning: Progresses and challenges. arXiv preprint arXiv:2402.00157.

Edward Beeching, Lewis Tunstall, and Sasha Rush. 2024. Scaling test-time compute with open models.

Zhenni Bi, Kai Han, Chuanjian Liu, Yehui Tang, and Yunhe Wang. 2024. Forest-of-thought: Scaling test-time compute for enhancing LLM reasoning. arXiv preprint arXiv:2412.09078.

Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D. Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, and 1 others. 2020. Language models are few-shot learners. Advances in Neural Information Processing Systems, 33:1877–1901.

Ngoc Bui, Hieu Trung Nguyen, Shantanu Kumar, Julian Theodore, Weikang Qiu, Viet Anh Nguyen, and Rex Ying. 2025. Mixture-of-personas language models for population simulation. arXiv preprint arXiv:2504.05019.

Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, and 1 others. 2022. Palm: Scaling language modeling with pathways. arXiv preprint arXiv:2204.02311.

Mehul Damani, Idan Shenfeld, Andi Peng, Andreia Bobu, and Jacob Andreas. 2024. Learning how hard to think: Input-adaptive allocation of LM computation. arXiv preprint arXiv:2410.04707.

John Doe, Jane Smith, and Alan Turing. 2023. Natural language processing in virtual assistants: Current approaches and challenges. In Proceedings of the IEEE Conference on Artificial Intelligence, pages 122–128. IEEE.

Xidong Feng, Ziyu Wan, Muning Wen, Stephen Marcus McAleer, Ying Wen, Weinan Zhang, and Jun Wang. 2023. Alphazero-like tree-search can guide large language model decoding and training. arXiv preprint arXiv:2309.17179.

Maria Garcia, Li Wei, and Ahmed El-Sayed. 2023. Ethical challenges in the development of virtual assistants powered by large language models. Electronics, 12(14):3170.

Xinyu Guan, Li Lyna Zhang, Yifei Liu, Ning Shang, Youran Sun, Yi Zhu, Fan Yang, and Mao Yang. 2025. rStar-Math: Small LLMs can master math reasoning with self-evolved deep thinking. arXiv preprint arXiv:2501.04519.

Dan Hendrycks, Collin Burns, Saurav Kadavath, Akul Arora, Steven Basart, Eric Tang, Dawn Song, and Jacob Steinhardt. 2021. Measuring mathematical problem solving with the math dataset. Advances in Neural Information Processing Systems.

Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, and 1 others. 2022. Training compute-optimal large language models. arXiv preprint arXiv:2203.15556.

Shima Imani, Liang Du, and Harsh Shirvastava. 2023. Mathprompter: Mathematical reasoning using large language models. arXiv preprint arXiv:2303.05398.

Ziwei Ji, Nayeon Lee, Rita Frieske, Tiezhen Yu, Dan Su, Yan Xu, Etsuko Ishii, Ye Jin Bang, Andrea Madotto, and Pascale Fung. 2023. Survey of hallucination in natural language generation. ACM Computing Surveys, 55(12):1–38.

Chonghe Jiang, Bao Nguyen, Anthony Man-Cho So, and Viet Anh Nguyen. 2025. Probe-free low-rank activation intervention. arXiv preprint arXiv:2502.04043.

Aitor Lewkowycz, Anders Andreassen, David Dohan, Ethan Dyer, Henryk Michalewski, Vinay Ramasesh, Ambrose Slone, Cem Anil, Imanol Schlag, Theo Gutman-Solo, and 1 others. 2022. Solving quantitative reasoning problems with language models. Advances in Neural Information Processing Systems, 35:3843–3857.

Junyi Li, Tianyi Tang, Wayne Xin Zhao, Jian-Yun Nie, and Ji-Rong Wen. 2024a. Pre-trained language models for text generation: A survey. ACM Computing Surveys, 56(9):1–39.

Kenneth Li, Oam Patel, Fernanda Viégas, Hanspeter Pfister, and Martin Wattenberg. 2023. Inference-time intervention: Eliciting truthful answers from a language model. Advances in Neural Information Processing Systems, 36:41451–41530.

Kenneth Li, Oam Patel, Fernanda Viégas, Hanspeter Pfister, and Martin Wattenberg. 2024b. Inference-time intervention: Eliciting truthful answers from a language model. Advances in Neural Information Processing Systems, 36.
Sheng Liu, Haotian Ye, Lei Xing, and James Zou. 2023a. In-context vectors: Making in context learning more effective and controllable through latent space steering. arXiv preprint arXiv:2311.06668.

Yiheng Liu, Tianle Han, Siyuan Ma, Jiayue Zhang, Yuanyuan Yang, Jiaming Tian, Hao He, Antong Li, Mengshen He, Zhengliang Liu, and 1 others. 2023b. Summary of chatgpt-related research and perspective towards the future of large language models. Meta-Radiology, page 100017.

Xinyin Ma, Gongfan Fang, and Xinchao Wang. 2023. Lim-pruner: On the structural pruning of large language models. Advances in neural information processing systems, 36:21702–21720.

Bao Nguyen, Binh Nguyen, Duy Nguyen, and Viet Anh Nguyen. 2025a. Risk-aware distributional intervention policies for language models. arXiv preprint arXiv:2501.15758.

Duy Nguyen, Archiki Prasad, Elias Stengel-Eskin, and Mohit Bansal. 2025b. Multi-attribute steering of language models via targeted intervention. arXiv preprint arXiv:2502.12446.

Hieu Trung Nguyen, Bao Nguyen, Binh Nguyen, and Viet Anh Nguyen. 2025c. Task-driven layerwise additive activation intervention. In Proceedings of the 2025 Conference of the Nations of the Americas Chapter of the Association for Computational Linguistics: Human Language Technologies (Volume 2: Short Papers), pages 506–513, Albuquerque, New Mexico. Association for Computational Linguistics.

Jack W. Rae, Sebastian Borgeaud, Trevor Cai, Katie Millican, Jordan Hoffmann, Francis Song, John Aslanides, Sarah Henderson, Roman Ring, Susanah Young, and 1 others. 2021. Scaling language models: Methods, analysis & insights from training gopher. arXiv preprint arXiv:2112.11446.

Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. 2019. Exploring the limits of transfer learning with a unified text-to-text transformer. arXiv preprint arXiv:1910.10683.

Bernardino Romera-Paredes, Mohammadamin Barekatain, Alexander Novikov, Matej Balog, M Pawan Kumar, Emilien Dupont, Francisco JR Ruiz, Jordan S Ellenberg, Pengming Wang, Omar Fawzi, and 1 others. 2024. Mathematical discoveries from program search with large language models. Nature, 625(7995):468–475.

Emre Sezgin. 2024. Redefining virtual assistants in health care: The future with large language models. Journal of Medical Internet Research, 26:e53225.

Nisan Stiennon, Long Ouyang, Jeffrey Wu, Daniel Ziegler, Ryan Lowe, Chelsea Voss, Alec Radford, Dario Amodei, and Paul F Christiano. 2020. Learning to summarize with human feedback. Advances in Neural Information Processing Systems, 33:3008–3021.

Trieu H Trinh, Yuhuai Wu, Quoc V Le, He He, and Thang Luong. 2024. Solving olympiad geometry without human demonstrations. Nature, 625(7995):476–482.

Xuezhi Wang, Jason Wei, Dale Schuurmans, Quoc Le, Ed Chi, Sharan Narang, Aakanksha Chowdhery, and Denny Zhou. 2022. Self-consistency improves chain of thought reasoning in language models. arXiv preprint arXiv:2203.11171.

Ziheng Wang, Jeremy Wohlwend, and Tao Lei. 2019. Structured pruning of large language models. arXiv preprint arXiv:1910.04732.

Yuhuai Wu, Albert Qiaochu Jiang, Wenda Li, Markus Rabe, Charles Staats, Mateja Jamnik, and Christian Szegedy. 2022. Autoformalization with large language models. Advances in Neural Information Processing Systems, 35:32353–32368.

Zhengxuan Wu, Atticus Geiger, Aryaman Arora, Jing Huang, Zheng Wang, Noah D Goodman, Christopher D Manning, and Christopher Potts. 2024. pyvene: A library for understanding and improving torch models via interventions. arXiv preprint arXiv:2403.07809.

Yuxi Xie, Anirudh Goyal, Wenye Zheng, Min-Yen Kan, Timothy P Lillicrap, Kenji Kawaguchi, and Michael Shieh. 2024. Monte Carlo tree search boosts reasoning via iterative preference learning. In The First Workshop on System-2 Reasoning at Scale, NeurIPS 24.

Ziwei Xu, Sanjay Jain, and Mohan Kankanhalli. 2024. Hallucination is inevitable: An innate limitation of large language models. arXiv preprint arXiv:2401.11817.

Shunyu Yao, Dian Yu, Jeffrey Zhao, Izhak Shafran, Tom Griffiths, Yuan Cao, and Karthik Narasimhan. 2024. Tree of thoughts: Deliberate problem solving with large language models. Advances in Neural Information Processing Systems, 36.

Dan Zhang, Sining Zhoubian, Ziniu Hu, Yisong Yue, Yuxiao Dong, and Jie Tang. 2024. Rest-MCTS*: LLM self-training via process reward guided tree search. arXiv preprint arXiv:2406.03816.

A Implementation details.

Datasets. The MATH dataset contains 12,500 math problems, with 12,000 problems designated for training and 500 for testing. Similarly, the GSM8K dataset comprises 8,792 math problems, with 7,473 used for training and 1,319 for testing. In our experiments, we subsample 1,500 training examples from the respective training sets. The head embeddings \( \mathcal{V} \) are then trained using the subsampled training set. Finally, we evaluate performance on the test splits of the aforementioned datasets.
B  Empirical Observation on the Impact of Attention Head Pruning.

Figure 5 shows that pruning an attention head in Qwen-2.5-1.5B Math Instruct model mostly leads to improvement in accuracy@1 on the MATH500 dataset.

![Accuracy Matrix per Subject and Attention Heads at Layer 23 (Qwen-2.5-1.5B Math Instruct)](page_495_682_474_180.png)

Figure 5: Impact of individual head pruning on model performance.
Table 1: We report Pass@1 scores for single head pruning across different layers of Meta-Llama-3-8B-Instruct, organized by question statement subjects.

<table>
  <tr>
    <th rowspan="2">Layer</th>
    <th colspan="5">Algebra</th>
    <th colspan="5">Counting & Probability</th>
    <th colspan="5">Geometry</th>
    <th colspan="5">Intermediate Algebra</th>
    <th colspan="5">Number Theory</th>
    <th colspan="5">Prealgebra</th>
    <th colspan="5">Precalculus</th>
  </tr>
  <tr>
    <th>5</th><th>10</th><th>15</th><th>31</th>
    <th>5</th><th>10</th><th>15</th><th>31</th>
    <th>5</th><th>10</th><th>15</th><th>31</th>
    <th>5</th><th>10</th><th>15</th><th>31</th>
    <th>5</th><th>10</th><th>15</th><th>31</th>
    <th>5</th><th>10</th><th>15</th><th>31</th>
    <th>5</th><th>10</th><th>15</th><th>31</th>
    <th>5</th><th>10</th><th>15</th><th>31</th>
  </tr>
  <tr>
    <td>Prune Head 0</td><td>0.35</td><td>0.37</td><td>0.32</td><td>0.32</td><td>0.11</td><td>0.11</td><td>0.13</td><td>0.16</td><td>0.2</td><td>0.24</td><td>0.2</td><td>0.2</td><td>0.22</td><td>0.13</td><td>0.14</td><td>0.14</td><td>0.16</td><td>0.19</td><td>0.19</td><td>0.21</td><td>0.18</td><td>0.4</td><td>0.38</td><td>0.41</td><td>0.39</td><td>0.41</td><td>0.12</td><td>0.14</td><td>0.12</td><td>0.12</td>
  </tr>
  <tr>
    <td>Prune Head 1</td><td>0.31</td><td>0.37</td><td>0.32</td><td>0.34</td><td>0.11</td><td>0.11</td><td>0.08</td><td>0.08</td><td>0.08</td><td>0.2</td><td>0.22</td><td>0.22</td><td>0.4</td><td>0.14</td><td>0.13</td><td>0.14</td><td>0.19</td><td>0.24</td><td>0.19</td><td>0.24</td><td>0.41</td><td>0.37</td><td>0.37</td><td>0.4</td><td>0.11</td><td>0.14</td><td>0.11</td><td>0.11</td><td>0.11</td>
  </tr>
  <!-- Table continued in rows below -->
</table>
Table 2: We report Pass@1 scores for single head pruning across different layers of Qwen2.5-Math-1.5B-Instruct, organized by question statement subjects.

<table>
  <tr>
    <th rowspan="2">Layer</th>
    <th colspan="4">Algebra</th>
    <th colspan="4">Counting & Probability</th>
    <th colspan="4">Geometry</th>
    <th colspan="4">Intermediate Algebra</th>
    <th colspan="4">Number Theory</th>
    <th colspan="4">Prealgebra</th>
    <th colspan="4">Precalculus</th>
  </tr>
  <tr>
    <th>0</th><th>5</th><th>15</th><th>27</th>
    <th>0</th><th>5</th><th>15</th><th>27</th>
    <th>0</th><th>5</th><th>15</th><th>27</th>
    <th>0</th><th>5</th><th>15</th><th>27</th>
    <th>0</th><th>5</th><th>15</th><th>27</th>
    <th>0</th><th>5</th><th>15</th><th>27</th>
    <th>0</th><th>5</th><th>15</th><th>27</th>
  </tr>
  <tr>
    <td>Prune Head 0</td><td>0.92</td><td>0.9</td><td>0.91</td><td>0.91</td><td>0.63</td><td>0.63</td><td>0.79</td><td>0.66</td><td>0.56</td><td>0.49</td><td>0.54</td><td>0.51</td><td>0.59</td><td>0.62</td><td>0.55</td><td>0.62</td><td>0.82</td><td>0.87</td><td>0.89</td><td>0.87</td><td>0.82</td><td>0.82</td><td>0.79</td><td>0.78</td><td>0.62</td><td>0.59</td><td>0.57</td><td>0.68</td>
  </tr>
  <tr>
    <td>Prune Head 1</td><td>0.92</td><td>0.92</td><td>0.92</td><td>0.92</td><td>0.71</td><td>0.68</td><td>0.63</td><td>0.63</td><td>0.51</td><td>0.56</td><td>0.59</td><td>0.51</td><td>0.6</td><td>0.52</td><td>0.53</td><td>0.58</td><td>0.84</td><td>0.85</td><td>0.85</td><td>0.90</td><td>0.82</td><td>0.77</td><td>0.78</td><td>0.79</td><td>0.62</td><td>0.62</td><td>0.64</td><td>0.66</td>
  </tr>
  <tr>
    <td>Prune Head 2</td><td>0.93</td><td>0.91</td><td>0.91</td><td>0.91</td><td>0.74</td><td>0.66</td><td>0.66</td><td>0.71</td><td>0.59</td><td>0.54</td><td>0.63</td><td>0.54</td><td>0.61</td><td>0.58</td><td>0.54</td><td>0.59</td><td>0.92</td><td>0.89</td><td>0.84</td><td>0.89</td><td>0.82</td><td>0.83</td><td>0.77</td><td>0.79</td><td>0.61</td><td>0.62</td><td>0.61</td><td>0.68</td>
  </tr>
  <tr>
    <td>Prune Head 3</td><td>0.06</td><td>0.91</td><td>0.92</td><td>0.9</td><td>0.08</td><td>0.71</td><td>0.61</td><td>0.71</td><td>0.07</td><td>0.56</td><td>0.61</td><td>0.49</td><td>0.07</td><td>0.6</td><td>0.59</td><td>0.6</td><td>0.1</td><td>0.85</td><td>0.89</td><td>0.90</td><td>0.12</td><td>0.79</td><td>0.83</td><td>0.79</td><td>0.09</td><td>0.62</td><td>0.66</td><td>0.68</td>
  </tr>
  <tr>
    <td>Prune Head 4</td><td>0.89</td><td>0.92</td><td>0.93</td><td>0.9</td><td>0.71</td><td>0.66</td><td>0.66</td><td>0.68</td><td>0.59</td><td>0.54</td><td>0.54</td><td>0.51</td><td>0.59</td><td>0.6</td><td>0.55</td><td>0.59</td><td>0.85</td><td>0.82</td><td>0.87</td><td>0.89</td><td>0.82</td><td>0.78</td><td>0.77</td><td>0.78</td><td>0.62</td><td>0.62</td><td>0.62</td><td>0.68</td>
  </tr>
  <tr>
    <td>Prune Head 5</td><td>0.56</td><td>0.91</td><td>0.9</td><td>0.91</td><td>0.5</td><td>0.61</td><td>0.61</td><td>0.66</td><td>0.46</td><td>0.59</td><td>0.54</td><td>0.56</td><td>0.37</td><td>0.58</td><td>0.53</td><td>0.59</td><td>0.47</td><td>0.89</td><td>0.81</td><td>0.84</td><td>0.54</td><td>0.83</td><td>0.76</td><td>0.77</td><td>0.41</td><td>0.64</td><td>0.68</td><td>0.66</td>
  </tr>
  <tr>
    <td>Prune Head 6</td><td>0.81</td><td>0.92</td><td>0.9</td><td>0.9</td><td>0.61</td><td>0.66</td><td>0.63</td><td>0.63</td><td>0.56</td><td>0.56</td><td>0.59</td><td>0.51</td><td>0.51</td><td>0.59</td><td>0.59</td><td>0.57</td><td>0.71</td><td>0.82</td><td>0.82</td><td>0.89</td><td>0.73</td><td>0.82</td><td>0.74</td><td>0.76</td><td>0.54</td><td>0.66</td><td>0.66</td><td>0.64</td>
  </tr>
  <tr>
    <td>Prune Head 7</td><td>0.91</td><td>0.92</td><td>0.9</td><td>0.9</td><td>0.63</td><td>0.61</td><td>0.63</td><td>0.66</td><td>0.56</td><td>0.56</td><td>0.54</td><td>0.54</td><td>0.58</td><td>0.6</td><td>0.56</td><td>0.59</td><td>0.85</td><td>0.87</td><td>0.89</td><td>0.87</td><td>0.78</td><td>0.79</td><td>0.8</td><td>0.78</td><td>0.64</td><td>0.68</td><td>0.55</td><td>0.62</td>
  </tr>
  <tr>
    <td>Prune Head 8</td><td>0.86</td><td>0.93</td><td>0.9</td><td>0.93</td><td>0.63</td><td>0.74</td><td>0.68</td><td>0.68</td><td>0.59</td><td>0.49</td><td>0.56</td><td>0.54</td><td>0.53</td><td>0.56</td><td>0.59</td><td>0.58</td><td>0.82</td><td>0.85</td><td>0.87</td><td>0.85</td><td>0.79</td><td>0.8</td><td>0.82</td><td>0.80</td><td>0.62</td><td>0.61</td><td>0.62</td><td>0.64</td>
  </tr>
  <tr>
    <td>Prune Head 9</td><td>0.93</td><td>0.91</td><td>0.91</td><td>0.91</td><td>0.66</td><td>0.66</td><td>0.68</td><td>0.66</td><td>0.56</td><td>0.49</td><td>0.54</td><td>0.49</td><td>0.57</td><td>0.54</td><td>0.52</td><td>0.57</td><td>0.9</td><td>0.84</td><td>0.85</td><td>0.89</td><td>0.78</td><td>0.82</td><td>0.8</td><td>0.79</td><td>0.64</td><td>0.66</td><td>0.66</td><td>0.66</td>
  </tr>
  <tr>
    <td>Prune Head 10</td><td>0.92</td><td>0.92</td><td>0.92</td><td>0.89</td><td>0.68</td><td>0.66</td><td>0.68</td><td>0.66</td><td>0.54</td><td>0.51</td><td>0.56</td><td>0.51</td><td>0.55</td><td>0.57</td><td>0.56</td><td>0.57</td><td>0.85</td><td>0.84</td><td>0.87</td><td>0.8</td><td>0.8</td><td>0.77</td><td>0.78</td><td>0.66</td><td>0.62</td><td>0.62</td><td>0.70</td>
  </tr>
  <tr>
    <td>Prune Head 11</td><td>0.9</td><td>0.92</td><td>0.91</td><td>0.9</td><td>0.68</td><td>0.66</td><td>0.71</td><td>0.66</td><td>0.51</td><td>0.56</td><td>0.56</td><td>0.49</td><td>0.6</td><td>0.61</td><td>0.55</td><td>0.6</td><td>0.87</td><td>0.84</td><td>0.87</td><td>0.78</td><td>0.8</td><td>0.77</td><td>0.80</td><td>0.66</td><td>0.62</td><td>0.61</td><td>0.66</td>
  </tr>
  <tr>
    <td>No Prune</td><td>0.91</td><td>0.91</td><td>0.91</td><td>0.91</td><td>0.63</td><td>0.63</td><td>0.63</td><td>0.63</td><td>0.49</td><td>0.49</td><td>0.49</td><td>0.49</td><td>0.58</td><td>0.58</td><td>0.57</td><td>0.89</td><td>0.89</td><td>0.87</td><td>0.79</td><td>0.79</td><td>0.79</td><td>0.79</td><td>0.79</td><td>0.66</td><td>0.66</td><td>0.66</td><td>0.68</td>
  </tr>
</table>
Table 3: We report Pass@1 scores for single head pruning across different layers of Qwen2.5-Math-7B-Instruct, organized by question statement subjects.

<table>
  <tr>
    <th rowspan="2">Layer</th>
    <th colspan="3">Algebra</th>
    <th colspan="3">Counting & Probability</th>
    <th colspan="3">Geometry</th>
    <th colspan="3">Intermediate Algebra</th>
    <th colspan="3">Number Theory</th>
    <th colspan="3">Prealgebra</th>
    <th colspan="3">Precalculus</th>
  </tr>
  <tr>
    <th>0</th><th>5</th><th>15</th><th>27</th><th>0</th><th>5</th><th>15</th><th>27</th><th>0</th><th>5</th><th>15</th><th>27</th><th>0</th><th>5</th><th>15</th><th>27</th><th>0</th><th>5</th><th>15</th><th>27</th><th>0</th><th>5</th><th>15</th><th>27</th>
  </tr>
  <tr>
    <td>Prune Head 0</td><td>0.96</td><td>0.96</td><td>0.97</td><td>0.97</td><td>0.84</td><td>0.82</td><td>0.84</td><td>0.84</td><td>0.68</td><td>0.68</td><td>0.59</td><td>0.66</td><td><b>0.68</b></td><td>0.67</td><td>0.64</td><td>0.62</td><td>0.9</td><td>0.95</td><td>0.95</td><td><b>0.97</b></td><td>0.9</td><td>0.88</td><td>0.9</td><td>0.88</td><td>0.71</td><td>0.77</td><td>0.7</td><td>0.80</td>
  </tr>
  <tr>
    <td>Prune Head 1</td><td><b>0.97</b></td><td>0.96</td><td>0.97</td><td>0.97</td><td>0.84</td><td>0.79</td><td>0.84</td><td>0.82</td><td>0.68</td><td>0.66</td><td>0.66</td><td>0.63</td><td>0.65</td><td>0.63</td><td>0.62</td><td>0.95</td><td>0.95</td><td>0.97</td><td>0.95</td><td>0.89</td><td>0.89</td><td>0.89</td><td>0.88</td><td>0.73</td><td>0.75</td><td>0.75</td><td>0.77</td>
  </tr>
  <tr>
    <td>Prune Head 2</td><td><b>0.97</b></td><td>0.96</td><td><b>0.98</b></td><td>0.97</td><td>0.84</td><td><b>0.84</b></td><td><b>0.87</b></td><td><b>0.84</b></td><td>0.66</td><td>0.66</td><td>0.66</td><td>0.65</td><td>0.64</td><td>0.65</td><td>0.64</td><td>0.92</td><td>0.95</td><td>0.92</td><td>0.95</td><td>0.89</td><td>0.89</td><td>0.9</td><td>0.87</td><td>0.75</td><td>0.71</td><td>0.77</td><td>0.77</td>
  </tr>
  <!-- Table continues with similar row structure and data for Prune Heads 3-27, No Prune -->
</table>
Table 4: We report Pass@1 scores for single head pruning across different layers of Qwen2.5-7B-Instruct, organized by question statement subjects.

<table>
  <tr>
    <th rowspan="2">Layer</th>
    <th colspan="3">Algebra</th>
    <th colspan="3">Counting & Probability</th>
    <th colspan="3">Geometry</th>
    <th colspan="3">Intermediate Algebra</th>
    <th colspan="3">Number Theory</th>
    <th colspan="3">Prealgebra</th>
    <th colspan="3">Precalculus</th>
  </tr>
  <tr>
    <th>0</th><th>5</th><th>15</th><th>27</th>
    <th>0</th><th>5</th><th>15</th><th>27</th>
    <th>0</th><th>5</th><th>15</th><th>27</th>
    <th>0</th><th>5</th><th>15</th><th>27</th>
    <th>0</th><th>5</th><th>15</th><th>27</th>
    <th>0</th><th>5</th><th>15</th><th>27</th>
  </tr>
  <tr><td>Prune Head 0</td><td>0.55</td><td>0.94</td><td>0.95</td><td>0.95</td><td>0.5</td><td>0.74</td><td>0.76</td><td>0.71</td><td>0.44</td><td>0.56</td><td>0.56</td><td>0.59</td><td>0.38</td><td>0.55</td><td>0.54</td><td>0.57</td><td>0.34</td><td>0.85</td><td>0.85</td><td>0.87</td><td>0.49</td><td>0.83</td><td>0.84</td><td>0.87</td><td>0.36</td><td>0.59</td><td>0.61</td><td>0.59</td></tr>
  <tr><td>Prune Head 1</td><td>0.85</td><td>0.94</td><td>0.93</td><td>0.94</td><td>0.61</td><td>0.71</td><td>0.74</td><td>0.71</td><td>0.49</td><td>0.63</td><td>0.56</td><td>0.61</td><td>0.51</td><td>0.57</td><td>0.53</td><td>0.59</td><td>0.82</td><td>0.87</td><td>0.87</td><td>0.73</td><td>0.84</td><td>0.83</td><td>0.87</td><td>0.70</td><td>0.62</td><td>0.59</td><td>0.61</td></tr>
  <tr><td>Prune Head 2</td><td>0.94</td><td>0.94</td><td>0.95</td><td>0.94</td><td>0.68</td><td>0.71</td><td>0.71</td><td>0.71</td><td>0.63</td><td>0.61</td><td>0.71</td><td>0.66</td><td>0.58</td><td>0.56</td><td>0.54</td><td>0.58</td><td>0.82</td><td>0.84</td><td>0.81</td><td>0.85</td><td>0.87</td><td>0.84</td><td>0.83</td><td>0.87</td><td>0.61</td><td>0.61</td><td>0.61</td><td>0.64</td></tr>
  <tr><td>Prune Head 3</td><td>0.02</td><td>0.94</td><td>0.92</td><td>0.94</td><td>0</td><td>0.76</td><td>0.76</td><td>0.74</td><td>0.05</td><td>0.61</td><td>0.68</td><td>0.61</td><td>0.04</td><td>0.58</td><td>0.60</td><td>0.6</td><td>0</td><td>0.85</td><td>0.84</td><td>0.85</td><td>0.01</td><td>0.84</td><td>0.84</td><td>0.87</td><td>0.02</td><td>0.61</td><td>0.62</td><td>0.61</td></tr>
  <tr><td>Prune Head 4</td><td>0.94</td><td>0.94</td><td>0.94</td><td>0.94</td><td>0.71</td><td>0.68</td><td>0.74</td><td>0.71</td><td>0.61</td><td>0.59</td><td>0.56</td><td>0.61</td><td>0.59</td><td>0.58</td><td>0.54</td><td>0.59</td><td>0.85</td><td>0.82</td><td>0.79</td><td>0.85</td><td>0.87</td><td>0.84</td><td>0.83</td><td>0.85</td><td>0.62</td><td>0.57</td><td>0.64</td><td>0.61</td></tr>
  <tr><td>Prune Head 5</td><td>0.95</td><td>0.95</td><td>0.93</td><td>0.94</td><td>0.71</td><td>0.74</td><td>0.74</td><td>0.71</td><td>0.59</td><td>0.61</td><td>0.61</td><td>0.56</td><td>0.58</td><td>0.54</td><td>0.57</td><td>0.6</td><td>0.85</td><td>0.84</td><td>0.82</td><td>0.87</td><td>0.82</td><td>0.84</td><td>0.85</td><td>0.62</td><td>0.62</td><td>0.62</td><td>0.61</td></tr>
  <tr><td>Prune Head 6</td><td>0.63</td><td>0.95</td><td>0.91</td><td>0.95</td><td>0.53</td><td>0.76</td><td>0.66</td><td>0.71</td><td>0.54</td><td>0.59</td><td>0.56</td><td>0.61</td><td>0.44</td><td>0.57</td><td>0.53</td><td>0.6</td><td>0.45</td><td>0.89</td><td>0.84</td><td>0.87</td><td>0.65</td><td>0.85</td><td>0.83</td><td>0.85</td><td>0.39</td><td>0.62</td><td>0.64</td><td>0.62</td></tr>
  <tr><td>Prune Head 7</td><td>0.96</td><td>0.93</td><td>0.91</td><td>0.94</td><td>0.76</td><td>0.74</td><td>0.76</td><td>0.74</td><td>0.68</td><td>0.61</td><td>0.54</td><td>0.59</td><td>0.55</td><td>0.57</td><td>0.59</td><td>0.58</td><td>0.84</td><td>0.87</td><td>0.82</td><td>0.87</td><td>0.84</td><td>0.85</td><td>0.85</td><td>0.84</td><td>0.64</td><td>0.64</td><td>0.61</td><td>0.64</td></tr>
  <tr><td>Prune Head 8</td><td>0.94</td><td>0.95</td><td>0.94</td><td>0.94</td><td>0.74</td><td>0.74</td><td>0.74</td><td>0.74</td><td>0.63</td><td>0.63</td><td>0.59</td><td>0.61</td><td>0.58</td><td>0.58</td><td>0.58</td><td>0.82</td><td>0.87</td><td>0.84</td><td>0.84</td><td>0.85</td><td>0.85</td><td>0.85</td><td>0.55</td><td>0.57</td><td>0.59</td><td>0.59</td></tr>
  <tr><td>Prune Head 9</td><td>0.95</td><td>0.95</td><td>0.94</td><td>0.94</td><td>0.97</td><td>0.68</td><td>0.74</td><td>0.68</td><td>0.71</td><td>0.56</td><td>0.66</td><td>0.56</td><td>0.59</td><td>0.58</td><td>0.55</td><td>0.57</td><td>0.85</td><td>0.85</td><td>0.92</td><td>0.87</td><td>0.87</td><td>0.87</td><td>0.87</td><td>0.85</td><td>0.61</td><td>0.57</td><td>0.62</td><td>0.64</td></tr>
  <tr><td>Prune Head 10</td><td>0.91</td><td>0.94</td><td>0.94</td><td>0.94</td><td>0.61</td><td>0.74</td><td>0.68</td><td>0.76</td><td>0.54</td><td>0.61</td><td>0.63</td><td>0.61</td><td>0.60</td><td>0.6</td><td>0.55</td><td>0.56</td><td>0.82</td><td>0.85</td><td>0.77</td><td>0.85</td><td>0.82</td><td>0.85</td><td>0.8</td><td>0.8</td><td>0.61</td><td>0.61</td><td>0.64</td><td>0.55</td></tr>
  <tr><td>Prune Head 11</td><td>0.87</td><td>0.93</td><td>0.95</td><td>0.94</td><td>0.71</td><td>0.74</td><td>0.76</td><td>0.74</td><td>0.59</td><td>0.61</td><td>0.56</td><td>0.61</td><td>0.57</td><td>0.56</td><td>0.55</td><td>0.57</td><td>0.77</td><td>0.87</td><td>0.89</td><td>0.82</td><td>0.83</td><td>0.84</td><td>0.84</td><td>0.52</td><td>0.62</td><td>0.68</td><td>0.61</td></tr>
  <tr><td>Prune Head 12</td><td>0.94</td><td>0.94</td><td>0.94</td><td>0.97</td><td>0.74</td><td>0.71</td><td>0.74</td><td>0.74</td><td>0.63</td><td>0.59</td><td>0.59</td><td>0.61</td><td>0.56</td><td>0.59</td><td>0.57</td><td>0.57</td><td>0.84</td><td>0.85</td><td>0.89</td><td>0.82</td><td>0.82</td><td>0.82</td><td>0.82</td><td>0.68</td><td>0.61</td><td>0.57</td><td>0.57</td></tr>
  <tr><td>Prune Head 13</td><td>0.95</td><td>0.95</td><td>0.93</td><td>0.94</td><td>0.71</td><td>0.71</td><td>0.71</td><td>0.79</td><td>0.61</td><td>0.63</td><td>0.61</td><td>0.59</td><td>0.57</td><td>0.6</td><td>0.60</td><td>0.54</td><td>0.85</td><td>0.89</td><td>0.87</td><td>0.85</td><td>0.87</td><td>0.88</td><td>0.83</td><td>0.83</td><td>0.62</td><td>0.59</td><td>0.55</td><td>0.55</td></tr>
  <tr><td>Prune Head 14</td><td>0.94</td><td>0.93</td><td>0.94</td><td>0.94</td><td>0.71</td><td>0.71</td><td>0.68</td><td>0.76</td><td>0.63</td><td>0.56</td><td>0.61</td><td>0.59</td><td>0.54</td><td>0.55</td><td>0.58</td><td>0.77</td><td>0.84</td><td>0.85</td><td>0.84</td><td>0.82</td><td>0.85</td><td>0.84</td><td>0.87</td><td>0.70</td><td>0.61</td><td>0.64</td><td>0.61</td></tr>
  <tr><td>Prune Head 15</td><td>0.81</td><td>0.94</td><td>0.94</td><td>0.95</td><td>0.68</td><td>0.74</td><td>0.76</td><td>0.71</td><td>0.54</td><td>0.56</td><td>0.61</td><td>0.61</td><td>0.52</td><td>0.56</td><td>0.58</td><td>0.59</td><td>0.77</td><td>0.84</td><td>0.85</td><td>0.82</td><td>0.82</td><td>0.88</td><td>0.85</td><td>0.59</td><td>0.57</td><td>0.64</td><td>0.59</td></tr>
  <tr><td>Prune Head 16</td><td>0.96</td><td>0.94</td><td>0.94</td><td>0.94</td><td>0.68</td><td>0.68</td><td>0.76</td><td>0.74</td><td>0.66</td><td>0.61</td><td>0.63</td><td>0.59</td><td>0.58</td><td>0.61</td><td>0.54</td><td>0.59</td><td>0.84</td><td>0.87</td><td>0.84</td><td>0.88</td><td>0.87</td><td>0.82</td><td>0.88</td><td>0.59</td><td>0.57</td><td>0.62</td><td>0.62</td></tr>
  <tr><td>Prune Head 17</td><td>0.93</td><td>0.95</td><td>0.94</td><td>0.94</td><td>0.71</td><td>0.71</td><td>0.74</td><td>0.76</td><td>0.59</td><td>0.59</td><td>0.63</td><td>0.59</td><td>0.54</td><td>0.55</td><td>0.60</td><td>0.61</td><td>0.81</td><td>0.84</td><td>0.87</td><td>0.85</td><td>0.89</td><td>0.87</td><td>0.82</td><td>0.82</td><td>0.62</td><td>0.59</td><td>0.61</td><td>0.61</td></tr>
  <tr><td>Prune Head 18</td><td>0.94</td><td>0.94</td><td>0.94</td><td>0.94</td><td>0.74</td><td>0.71</td><td>0.74</td><td>0.71</td><td>0.66</td><td>0.61</td><td>0.56</td><td>0.61</td><td>0.56</td><td>0.56</td><td>0.58</td><td>0.56</td><td>0.85</td><td>0.84</td><td>0.85</td><td>0.87</td><td>0.85</td><td>0.83</td><td>0.84</td><td>0.64</td><td>0.59</td><td>0.66</td><td>0.61</td></tr>
  <tr><td>Prune Head 19</td><td>0.94</td><td>0.96</td><td>0.92</td><td>0.94</td><td>0.71</td><td>0.76</td><td>0.76</td><td>0.71</td><td>0.63</td><td>0.56</td><td>0.66</td><td>0.59</td><td>0.56</td><td>0.53</td><td>0.58</td><td>0.82</td><td>0.87</td><td>0.84</td><td>0.85</td><td>0.87</td><td>0.79</td><td>0.87</td><td>0.85</td><td>0.66</td><td>0.57</td><td>0.70</td><td>0.61</td></tr>
  <tr><td>Prune Head 20</td><td>0.95</td><td>0.94</td><td>0.94</td><td>0.94</td><td>0.74</td><td>0.74</td><td>0.74</td><td>0.71</td><td>0.66</td><td>0.59</td><td>0.61</td><td>0.60</td><td>0.57</td><td>0.57</td><td>0.58</td><td>0.81</td><td>0.87</td><td>0.82</td><td>0.84</td><td>0.84</td><td>0.88</td><td>0.61</td><td>0.62</td><td>0.66</td><td>0.61</td></tr>
  <tr><td>Prune Head 21</td><td>0.95</td><td>0.96</td><td>0.96</td><td>0.96</td><td>0.76</td><td>0.74</td><td>0.74</td><td>0.71</td><td>0.61</td><td>0.61</td><td>0.61</td><td>0.56</td><td>0.58</td><td>0.6</td><td>0.85</td><td>0.85</td><td>0.87</td><td>0.85</td><td>0.84</td><td>0.87</td><td>0.61</td><td>0.62</td><td>0.62</td><td>0.61</td></tr>
  <tr><td>Prune Head 22</td><td>0.44</td><td>0.96</td><td>0.93</td><td>0.94</td><td>0.24</td><td>0.76</td><td>0.74</td><td>0.71</td><td>0.22</td><td>0.59</td><td>0.66</td><td>0.59</td><td>0.12</td><td>0.58</td><td>0.55</td><td>0.61</td><td>0.26</td><td>0.85</td><td>0.87</td><td>0.84</td><td>0.52</td><td>0.89</td><td>0.83</td><td>0.84</td><td>0.07</td><td>0.62</td><td>0.64</td><td>0.61</td></tr>
  <tr><td>Prune Head 23</td><td>0.11</td><td>0.94</td><td>0.92</td><td>0.94</td><td>0.21</td><td>0.74</td><td>0.63</td><td>0.71</td><td>0.17</td><td>0.66</td><td>0.68</td><td>0.59</td><td>0.04</td><td>0.62</td><td>0.49</td><td>0.58</td><td>0.06</td><td>0.84</td><td>0.79</td><td>0.87</td><td>0.27</td><td>0.85</td><td>0.82</td><td>0.87</td><td>0.04</td><td>0.61</td><td>0.59</td><td>0.59</td></tr>
  <tr><td>Prune Head 24</td><td>0.83</td><td>0.94</td><td>0.93</td><td>0.96</td><td>0.61</td><td>0.74</td><td>0.71</td><td>0.71</td><td>0.56</td><td>0.66</td><td>0.63</td><td>0.66</td><td>0.51</td><td>0.6</td><td>0.56</td><td>0.6</td><td>0.77</td><td>0.85</td><td>0.84</td><td>0.85</td><td>0.68</td><td>0.87</td><td>0.84</td><td>0.84</td><td>0.57</td><td>0.62</td><td>0.55</td><td>0.62</td></tr>
  <tr><td>Prune Head 25</td><td>0.02</td><td>0.95</td><td>0.95</td><td>0.95</td><td>0</td><td>0.76</td><td>0.74</td><td>0.66</td><td>0</td><td>0.63</td><td>0.66</td><td>0.63</td><td>0</td><td>0.55</td><td>0.60</td><td>0.59</td><td>0</td><td>0.89</td><td>0.81</td><td>0.85</td><td>0.02</td><td>0.84</td><td>0.83</td><td>0.85</td><td>0.02</td><td>0.70</td><td>0.59</td><td>0.66</td></tr>
  <tr><td>Prune Head 26</td><td>0.5</td><td>0.95</td><td>0.93</td><td>0.95</td><td>0.47</td><td>0.71</td><td>0.71</td><td>0.71</td><td>0.37</td><td>0.54</td><td>0.54</td><td>0.56</td><td>0.25</td><td>0.58</td><td>0.60</td><td>0.56</td><td>0.85</td><td>0.89</td><td>0.89</td><td>0.63</td><td>0.88</td><td>0.85</td><td>0.85</td><td>0.21</td><td>0.55</td><td>0.64</td><td>0.66</td></tr>
  <tr><td>Prune Head 27</td><td>0.83</td><td>0.96</td><td>0.94</td><td>0.95</td><td>0.68</td><td>0.71</td><td>0.71</td><td>0.63</td><td>0.46</td><td>0.68</td><td>0.59</td><td>0.56</td><td>0.49</td><td>0.6</td><td>0.52</td><td>0.59</td><td>0.69</td><td>0.85</td><td>0.82</td><td>0.82</td><td>0.63</td><td>0.84</td><td>0.84</td><td>0.85</td><td>0.61</td><td>0.59</td><td>0.59</td><td>0.64</td></tr>
  <tr><td>No Prune</td><td>0.94</td><td>0.94</td><td>0.94</td><td>0.94</td><td>0.74</td><td>0.74</td><td>0.74</td><td>0.74</td><td>0.59</td><td>0.59</td><td>0.59</td><td>0.59</td><td>0.59</td><td>0.59</td><td>0.59</td><td>0.59</td><td>0.84</td><td>0.84</td><td>0.84</td><td>0.87</td><td>0.87</td><td>0.87</td><td>0.87</td><td>0.59</td><td>0.59</td><td>0.59</td><td>0.59</td></tr>
</table>

Reviewer paper summaries: This paper introduces **SPRINT**, a novel contrastive learning framework aimed at improving reasoning performance in transformer-based language models through selective attention head pruning. It is motivated by an observation that pruning certain heads can increase model performance on certain categories of questions. In order to decide which attention heads to prune at inference, the authors introduce **Diversity Promoted Contrastive Loss**, which aligns question embedding with the beneficial heads to prune to yield the correct answer. The approach generates N candidate solutions by pruning N different heads. Experiments on the MATH500 and GSM8K datasets show that SPRINT outperforms traditional best-of-N and random head selection methods.

This paper first provides an observation: removing certain attention heads in a transformer doesn’t always hurt performance—sometimes it even improves accuracy. Based on this, they propose SPRINT, a method that learns a vector representation for each attention head at each layer. This trained vector helps decide which heads to prune during inference, making the model more efficient without extra computation.

This paper introduces SPRINT, a novel framework that leverages structured attention head pruning to improve reasoning performance in transformer-based language models. Unlike conventional approaches that treat pruning primarily as a compression technique, SPRINT repurposes it to induce diverse reasoning behaviors by selectively pruning attention heads that may negatively impact performance. A contrastive learning objective is employed to align input questions with effective head-pruning configurations. The method is evaluated on the MATH500 and GSM8K datasets, demonstrating consistent improvements over traditional sampling-based strategies.
Summary of strengths: * The finding that pruning attention heads in language models can improve the reasoning performance of models is interesting. 
* The proposed **Diversity Promoted Contrastive Loss** is intuitive and promising, as the experiment results are positive.
* Paper is well-written overall.

1. The work aims to reduce computational costs in the best-of-N strategy while achieving better accuracy.

1. This paper presents a creative perspective by using attention head pruning not as a compression tool but as a means to enhance answer diversity during generation. This approach addresses the issue of reasoning diversity without incurring significant additional computational overhead.

2. The use of contrastive learning to guide head pruning is well-motivated and shows clear performance gains in the experimental results on MATH500 and GSM8K.

3. The paper is well-organized, with a logical flow from motivation to methodology and experimental validation, making it easy to follow and understand the key insights.
Summary of weaknesses: * The choice of the base model is concerning. The only model used in this paper is Qwen2.5-Math-1.5B Instruct, which is pre-trained on mathematical corpora and optimized for math-related tasks. This specialization may cause different attention heads to capture distinct mathematical concepts, making head pruning effective by removing noisy heads for tasks requiring diverse mathematical understanding. However, this observation may not generalize to other LLMs that are not specifically trained for the math domain. Evaluating the framework on more diverse, non-math-specialized models would strengthen the claims, which contradicts traditional view that pruning attention heads is harmful to model performance.

1. The paper is unfinished and unclear and lacks details. For example, in line 190, they only present the result of pruning head 9 in layer 21 on Qwen2.5-Math-1.5B, yet they generalize the conclusion that "removing certain attention heads does not degrade performance; in fact, it sometimes improves accuracy." without providing sufficient supporting evidence across multiple heads, layers, or models. In line 200-250, it does not illustrate the subject of training and is hard to follow the training process, whether they train the base model or a learnable vector. 

2. The paper is not technical, as it introduces many variables without explaining their usage, making it difficult to follow. Additionally, it inconsistently refers to the MATH and MATH500 datasets, causing further confusion. The MATH dataset was originally released by DeepMind, while OpenAI later released PRM800K, which is derived from the MATH dataset. MATH500, on the other hand, consists of 500 questions sampled from PRM800K. They also don't cite the source of dataset, which makes it hard to pinpoint the training data. Until the end, I still don't know what's the training dataset. 

3. They only evaluate their method with Qwen2.5-Math-1.5B, doesn't evaluate the weaker or stronger policy model, like Qwen-Math-7B, or deepseek-Math, or general models, etc. It also lacks any analysis and only provides a table and a graph, such as heatmap or gradient accumulation graph when pruning one specific head in one layer. 

4. The paper fails to compare its pruned model with the original model, raising concerns about experimental reliability, especially since its reported performance significantly exceeds the official Qwen2.5-Math-1.5B benchmarks (76.8% on GSM8K, 49.8% on MATH in few-shot, 84.8% on GSM8K, 75.8% on MATH in best-of-8), making the results highly questionable.

5. The paper lacks novelty and practical value, as modern large models already have a high number of attention heads (e.g., Qwen-1.5B has 12 heads, Qwen-7B has 28 heads, and Qwen-72B has 64 heads). Even if it optimizes best-of-N, the computational reduction is only 1/N, which becomes negligible as model size increases, while simultaneously introducing performance uncertainty.

1. Pruning can potentially degrade LLMs. As illustrated in Figure 2, randomly pruning heads tends to perform worse than not pruning at all. While the contrastive learning approach helps mitigate this issue, it introduces additional training steps and requires training data. Compared to other methods that do not require any training or training data, this may present an unfair baseline comparison.

2. The effectiveness of the proposed method is only demonstrated on Qwen2.5-Math-1.5B-Instruct. It remains unclear whether similar benefits hold for other models such as LLaMA-3 or across different model sizes (e.g., 7B or 72B). Broader experimentation would strengthen the paper and demonstrate the robustness of the approach.
Comments, suggestions and typos: * What is the exact value of l, the number of layers pruned during the experiments? Knowing this value and its ratio relative to the total number of layers in the model would provide better insight into the extent of pruning.
* If accepted, I suggest the authors move Figure 3 from the appendix to the main paper, as the current text draws key conclusions based on this figure.
* Lines 242–244 and 260–263 are confusing. The metric used is Accuracy@k, but these sections mention that the final answer is aggregated over N generations and that accuracy is based on the best answer. I suggest the authors clarify these descriptions to avoid ambiguity.

See weakness

See above.
Overall assessment scores: 3

2

3.5