Year: 2025
Type: conference-paper
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.
Authors: Hieu Trung Nguyen, Bao Nguyen, Viet Anh Nguyen
Institution: Chinese University of Hong Kong
Language: English
Topic: Topic Modeling
Subfield: Artificial Intelligence
Field: Computer Science
Domain: Physical Sciences
Open Access status: gold
OpenAlex ID: W4412944633