Structured Pruning for Diverse Best-of-N Reasoning Optimization

2025/06/04 by Hieu Trung Nguyen, Bao D. Nguyen, Nguyen, Hieu Trung +3 · 1 citation
Computer Science · #Advanced Graph Neural Networks #Computation and Language (cs.CL) #FOS: Computer and information sciences #Machine Learning (cs.LG) #Multimodal Machine Learning Applications #Topic Modeling

paper · pdf · doi:10.48550/arxiv.2506.03978

openalex publication_date 2025/06/04 · openalex created_date 2025/10/10 · openalex updated_date 2026/07/28

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.

Citations

• Mixture-of-Personas Language Models for Population Simulation
• Multi-Attribute Steering of Language Models via Targeted Intervention
• Task-driven Layerwise Additive Activation Intervention
• Probe-Free Low-Rank Activation Intervention
• Distributional Surgery for Language Model Activations
• rStar-Math: Small LLMs Can Master Math Reasoning with Self-Evolved Deep Thinking
• Forest-of-Thought: Scaling Test-Time Compute for Enhancing LLM Reasoning
• Learning How Hard to Think: Input-Adaptive Allocation of LM Computation
• Monte Carlo Tree Search Boosts Reasoning via Iterative Preference Learning
• pyvene: A Library for Understanding and Improving PyTorch Models via Interventions
• In-context Vectors: Making In Context Learning More Effective and Controllable Through Latent Space Steering
• Alphazero-like Tree-Search can Guide Large Language Model Decoding and Training
• Inference-Time Intervention: Eliciting Truthful Answers from a Language Model
• Let's Verify Step by Step
• Solving Quantitative Reasoning Problems with Language Models
• Self-Consistency Improves Chain of Thought Reasoning in Language Models
• Pretrained Language Models for Text Generation: A Survey
• Scaling Language Models: Methods, Analysis & Insights from Training Gopher
• Training Verifiers to Solve Math Word Problems
• Measuring Mathematical Problem Solving With the MATH Dataset
• Learning to summarize from human feedback
• Language Models are Few-Shot Learners
• Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer

Cited by

• Reasoning Planning for Language Models

Related

• Inference-Time Chain-of-Thought Pruning with Latent Informativeness Signals
• Think Before You Prune: Selective Self-Generated Calibration for Pruning Large Reasoning Models
• Think, Prune, Train, Improve: Scaling Reasoning without Scaling Models
• Efficient LLMs with AMP: Attention Heads and MLP Pruning
• Gradient-based Intra-attention Pruning on Pre-trained Language Models