SDMPrune: Self-Distillation MLP Pruning for Efficient Large Language Models
This work addresses the problem of efficient LLM deployment for users needing reduced computational costs, though it is incremental as it builds on existing gradient-based pruning techniques.
The paper tackles the high deployment costs of large language models (LLMs) by introducing a pruning method that uses self-distillation during pruning to better preserve generative capabilities, focusing on MLP modules which are less sensitive but parameter-heavy, achieving significant compression without obvious performance degradation and outperforming existing methods on zero-shot benchmarks.
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.