用自蒸馏提升大模型剪枝精度,专注压缩MLP模块
SDMPrune: Self-Distillation MLP Pruning for Efficient Large Language Models
- 剪枝时引入自蒸馏损失,更好利用原模型预测信息
- 在1.2B模型上剪除5倍以上参数,性能几乎不变
- 适合追求高效部署的开源大模型研究者
尽管大语言模型(LLM)表现强劲,但其部署成本过高。现有基于梯度的剪枝方法因使用独热标签计算梯度,忽略了对其他词的潜在预测,从而丢失生成能力的关键信息。为此,我们提出在剪枝阶段而非训练后引入自蒸馏损失,充分挖掘原模型的预测能力,获得更精准的梯度用于剪枝。此外,我们发现相比注意力模块,大模型的预测对多层感知机(MLP)模块不敏感,而其参数占比超过5倍(以LLaMA3.2-1.2B为例)。因此,我们聚焦于剪枝MLP模块,显著压缩模型且性能无明显下降。大量零样本基准测试结果表明,本方法显著优于现有剪枝方法,并在1B规模开源模型中表现竞争力。源代码与训练权重已公开于https://github.com/visresearch/SDMPrune。
原文摘要 · Abstract (English)
In spite of strong performance achieved by LLMs, the costs of their deployment are unaffordable. For the compression of LLMs, gradient-based pruning methods present promising effectiveness. However, in these methods, the gradient computation with one-hot labels ignore the potential predictions on other words, thus missing key information for generative capability of the original model. To address this issue, we introduce a self-distillation loss during the pruning phase (rather than post-training) to fully exploit the predictions of the original model, thereby obtaining more accurate gradient information for pruning. Moreover, we find that, compared to attention modules, the predictions of LLM are less sensitive to multilayer perceptron (MLP) modules, which take up more than $5 \times$ parameters (LLaMA3.2-1.2B). To this end, we focus on the pruning of MLP modules, to significantly compress LLM without obvious performance degradation. Experimental results on extensive zero-shot benchmarks demonstrate that our method significantly outperforms existing pruning methods. Furthermore, our method achieves very competitive performance among 1B-scale open source LLMs. The source code and trained weights are available at https://github.com/visresearch/SDMPrune.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。