剪枝注意力头能提升模型推理能力，SPRINT动态选最优剪枝方案。
Structured Pruning for Diverse Best-of-N Reasoning Optimization
- 通过对比学习动态选择推理时要剪枝的注意力头和层
- 在MATH500和GSM8K上显著优于传统Best-of-N和随机剪枝
- 适合追求高效高精度推理的模型部署场景
基于Transformer的语言模型中的模型剪枝，传统上被视为实现计算节省的手段，但本研究发现：有选择地剪枝某些注意力头可提升模型的推理能力，尤其在困难任务上表现突出。受此启发，我们提出SPRINT，一种新颖的对比学习框架，可在推理过程中动态选择最优的注意力头与层进行剪枝。通过将问题嵌入与头嵌入对齐，SPRINT识别出能带来更准确推理的剪枝配置。大量实验表明，该方法在MATH500和GSM8K数据集上显著优于传统的Best-of-N及随机头选择策略。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。