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.
@article{arxiv.2506.03978,
title = {Structured Pruning for Diverse Best-of-N Reasoning Optimization},
author = {Hieu Trung Nguyen and Bao Nguyen and Viet Anh Nguyen},
journal= {arXiv preprint arXiv:2506.03978},
year = {2025}
}