arXiv · 2506.03978
Structured Pruning for Diverse Best-of-N Reasoning Optimization
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.
Explore related subjects
Keep this discovery
Explore connections, maps & timelines
Hieu Trung Nguyen, Bao Nguyen, Viet Anh Nguyen. 2025-06-04. Structured Pruning for Diverse Best-of-N Reasoning Optimization. https://arxiv.org/abs/2506.03978
Cite the original work for its findings. Save a collection to share your selection of sources.