MINI-LLM: Memory-Efficient Structured Pruning for Large Language Models

Gradient-based Pruning. To evaluate the significance of a specific weight $W_{l}^{k}$, one common way is using its sensitive effect on the loss function, where $k$ denotes the $k$-th weight in the $l$-th layer. Specifically, one can compare the difference in the loss function when $W_{l}^{k}$ is included versus when it is excluded from the model (i.e., LLaMA-7B). Thus, the loss change can be formulated as LeCun et al. (1989):

where $\Delta W=W_{l}^{k}$ and the Hessian matrix $H=\nabla^{2}_{W_{l}^{k}}\mathcal{L}(W_{0})$. Unlike previous work Liu et al. (2021); Kurtic et al. (2022), the pre-trained datasets of a large language model are inconsistent with the datasets of the downstream tasks, hence $\frac{\partial\mathcal{L}(\mathcal{D})}{\partial W_{l}^{k}}\not\approx 0$. This characteristic is advantageous for assessing the importance of weights through the gradient term in the context of LLMs, as calculating the Hessian matrix in the second term is impractical on LLMs with $\mathcal{O}(N^{2})$ complexity. Therefore, $\Delta\mathcal{L}$ can be approximated as:

Activation-based Pruning. As recently observed in LLMs larger than 6.7B (Dettmers et al. (2022)), a small set of hidden state features emerges with significantly larger magnitudes (outliers) than the remainders and zeroing out these features causes a significant degradation of performance. The vanilla scoring function Eq. (3) does not highlight the unique characteristics of LLMs compared with smaller models. Given the $l$th layer’s input activation $X_{l}$ (i.e., the output from the ($l-1$)th layer of the network) and the $l$th layer’s weights $W_{l}$, the pruning problem is also commonly treated as finding the subset $\hat{W}_{l}$ of $W_{l}$ respecting a compression constraint $\mathcal{C}$, which most closely approximates the initial output as determined by the squared error metric. Assuming that the activation and weight matrices possess a suitable rectangular shape, the neural network pruning is defined as the following optimization problem Frantar and Alistarh (2023):

To evaluate the significance of a specific weight $W_{l}^{k}$, one can compare the difference in the layer-wise output when $W_{l}^{k}$ is preserved versus when it is excluded from the model and write the formulation as:

where $\Delta W_{l}=W_{l}^{k}$.

Feature Map Sensitivity (FMS). Eq. (5) only considers the output changes in a single layer and does not take into account the global loss change across the entire network. We notice that, with the weight gradients, the global loss change can be quantified with weight change as shown in Eq. (3). To calculate the salience of each weight relative to the change in global loss and layer-wise output, we measure $\Delta W_{l}$ in Eq. (5) through its global sensitivity $\Delta\mathcal{L}$ from Eq. (3) and propose a heuristic hybrid sensitivity scoring function called Feature Map Sensitivity (FMS) as follows:

Compared to Eq. (3) and Eq. (5), our criterion Eq. (6) integrates magnitude, activation, and gradient to optimally utilize the pivotal information from the three critical aspects, so as calculate the feature map sensitivity along with the loss changes.

where $\nabla\mathcal{L}(W;\mathcal{B})$ is the gradient with backpropagation, $z\in\mathbb{R}^{d}$ with $z\sim\mathcal{N}(0,I_{d})$ and $\epsilon$ is the perturbation scale. The $n$-ZO gradient estimate averages $\hat{\nabla}\mathcal{L}(W;\mathcal{B})$ over $n$ randomly sampled $z$. Malladi et al. Malladi et al. (2023) found that $n=1$ is the most efficient. Therefore, we choose $n=1$ as the default. For each weight $W_{l}^{k}$ in the model, the estimation of its gradient $\frac{\partial\mathcal{L}(\mathcal{D})}{\partial W_{l}^{k}}$ (defined as $\hat{g}_{l}^{k}$) is then

where $z_{l}^{k}\in z$ is the random corresponding to $W_{l}^{k}$. In this way, the practical pruning score used in our MINI-LLM is defined as:

Dependency-aware structured LLM pruning. To maintain structural integrity, it is crucial for structured pruning to identify groups of interdependent structures within LLMs. Following Ma et al. Ma et al. (2023), we prune heads for Multi-Head Attention (MHA) and channels for Feed-Forward Network (FFN), respectively. We arrange the interconnected weights into groups and determine the sensitivity of each group (a set of coupled structures) defined as $G=\{W_{i}\}_{i=1}^{M}$ by choosing the maximum sensitivity score of the structures in it, i.e., $\hat{S}({G})=\text{max}_{i=1}^{M}\sum_{k}\hat{S}(W_{i}^{k})$, where $M$ is the number of interdependent structures in the group. Our structured pruning approach MINI-LLM is outlined in Algorithm 1.

Pruning and fine-tuning settings. Our MINI-LLM conducts in a one-shot pruning framework. That is scoring only once and then pruning the network to the target prune ratio Cheng et al. (2023). In the model pruning process, we use 10 randomly selected samples from Bookcorpus Zhu et al. (2015) as the calibration data for evaluating the weight gradients and 128 samples for computing each layer’s input (i.e., activation). Due to the varying sensitivity of each layer to pruning Ma et al. (2023), the first four layers and the last three layers are retained. During the recovery phase, we utilize the Alpaca-cleaned Taori et al. (2023) as the training dataset, which contains approximately 50k samples, to fine-tune the pruned models with a batch size of 64. Following [Ma et al. (2023)], the learning rate is set to 1e-4 and a total of 2 epochs. Each pruned model is recovered by an Adam optimizer Kingma and Ba (2015) paired with a cosine decay schedule for the learning rate. We set LoRA $r=8$, $\alpha=16$, and attach LoRA modules on all linear layers of the base model. In the inference stage, all the evaluations are implemented with a context length of 128.

Baselines. We compare MINI-LLM with four one-shot structured pruning methods for LLMs. Magnitude-l1/l2: pruning based on the absolute values or the l2-norm of weights, respectively. LLM-Pruner Ma et al. (2023): pruning using criterion Eq. (3) with backpropagation gradients. Wanda Sun et al. (2024): pruning based on the product of the magnitude of weights and their corresponding activations. Given that vanilla SparseGPT and Wanda are retraining-free unstructured methods, we adapt them for structured pruning with pruning and fine-tuning stages for a fair comparison while maintaining the same criterion. Except for LLM-pruner, which is a gradient-based method, the other methods are all gradient-free methods.

Zero-shot performance on BLOOM-7B and OPT-6.7B. To validate MINI-LLM on other LLMs broadly, we prune both BLOOM-7B and OPT-6.7B with two prune ratios: 10% and 30%, and fine-tune the pruned models to restore model accuracy. The results in Table 2 illustrate that our MINI-LLM steadily outperforms all gradient-free methods and exhibits performance comparable to, even surpasses at times, that of LLM-Pruner. For instance, at a 30% compression rate on BLOOM-7B, MINI-LLM achieves a perplexity of 54.07 on the WikiText2 dataset, obviously outperforming LLM-Pruner’s perplexity of 58.11. Similarly, at a 30% compression rate on OPT-6.7B, MINI-LLM achieves a perplexity of 40.89 on the WikiText2 dataset and 57.44 on the PTB dataset, outperforming LLM-Pruner’s perplexity of 42.94 and 65.09, respectively. In addition, at a 10% prune ratio on OPT-6.7B, MINI-LLM achieves an average classification accuracy of 67.81% across four datasets and obtains 98.60% of the accuracy achieved by the original model, which is even better than LLM-Pruner’ s 67.50% and 98.15%. This demonstration validates the effectiveness of MINI-LLM in efficiently compressing models of various structures to a specified size, while optimizing memory usage.

Zero-shot Performance on LLaMA-13B Due to the efficient approximation for the gradients of the pre-trained weights, MINI-LLM enables pruning on larger-scale LLMs, such as LLaMA-13B ⁴⁴4https://huggingface.co/huggyllama/llama-13b/tree/main. We prune LLaMA-13B with five pruning ratios: from 10% to 50% and present the zero-shot performance of the pruned LLaMA-13B without fine-tuning in Table 3 and with fine-tuning in Figure 3, respectively. Except for the model pruned with a ratio of 50%, for which fine-tuning is conducted for two epochs, the compressed models obtained from all other pruning ratios are fine-tuned for just one epoch. The other fine-tuning settings, such as the learning rate and batch size, are the same as those for recovering LLaMA-7B. We follow Wanda Sun et al. (2024) to conduct inference with 2048 tokens.

Few-shot performance on LLaMA-7B. In Table 4, we report the mean accuracies for both dense LLMs and sparse LLMs with 20% to 50% sparsity. In the few-shot setting, MINI-LLM performs competitively with other methods, including backpropagation gradient-based LLM-Pruner. Specifically, at a 20% prune ratio, MINI-LLM achieves an average accuracy of 26.60%, which surpasses SparseGPT’s 25.80% and LLM-Pruner’s 25.30%. Notably, estimated gradient-based MINI-LLM consistently surpasses backpropagation gradient-based LLM-Pruner. This performance is not observed in zero-shot setting.

Model size, complexity, and inference time. Table 5 shows the number of parameters, MACs, GPU memory requirements, and total inference time for running the original model and the pruned LLaMA-7B models at different prune ratios. The results indicate that when the model is pruned by 50%, the total inference time is reduced to 58% and the GPU memory usage concurrently drops to 50% of its original values, respectively. The evaluation is conducted in the inference mode and the sequence length is set to 64. The inference time is tested under the test dataset of WikiText2 on a single NVIDIA GeForce RTX 3090Ti (24GB).

Efficacy of Estimated Gradients on BLOOM-7B and OPT-6.7B In the main body, we explored the effectiveness of the estimated gradients for guiding the pruning of LLaMA-7B. Here, we delve further into the effectiveness of the estimated gradients in guiding the pruning process for BLOOM-7B ⁵⁵5https://huggingface.co/bigscience/bloom-7b1/tree/main and OPT-6.7B ⁶⁶6https://huggingface.co/facebook/opt-6.7b/tree/main. The experimental results presented in Table 6 indicate that our score function FMS, $|W\hat{\nabla}\mathcal{L}(W)X|$, consistently yields better performance compared to Wanda’s pruning criterion, $|WX|$, for pruning BLOOM-7B and OPT-6.7B on the WikiText2 and PTB datasets. For example, at a 30% pruning ratio, $|W\hat{\nabla}\mathcal{L}(W)X|$ achieves a perplexity of 54.07 on WikiText2 for pruning BLOOM-7B. This result shows a 30.82 improvement over $|WX|$. Similarly, at a 30% pruning ratio, $|W\hat{\nabla}\mathcal{L}(W)X|$ achieves a perplexity of 40.89 on WikiText2 for pruning OPT-6.7B, which surpasses the perplexity of 82.93 obtained by using $|WX|$.

Efficacy of activation on LLaMA-7B. Dettmers et al. (2022); Kovaleva et al. (2021) identified a distinct property of LLMs that a few hidden state features possess notably high magnitudes. Eliminating these features results in a considerable decline in performance. As argued in Section 3.1, the vanilla pruning criterion $\left|W\nabla\mathcal{L}(W)\right|$ ( Eq. (3)) does not highlight this characteristic of LLMs. To validate that activation in FMS, i.e., $\left|W\nabla\mathcal{L}(W)X\right|$ (Eq. (6)), can bring improved performance, we compare the performance of the pruned models obtained by using the criteria Eq. (3) and Eq. (6) over four prune ratios on the LLM. The difference between these two criteria is that Eq. (6) includes an additional activation term $X$ compared to Eq. (3). From the results shown in Table 7, we can see that the activations enable an effective increase in performance on both the WikiText2 and PTB datasets. For example, at a 20% prune ratio, the model pruned with $\left|W\nabla\mathcal{L}(W)X\right|$ achieves a perplexity of 17.45 on the WikiText2 dataset. This result surpasses the 17.79 perplexity achieved by the model compressed by using $\left|W\nabla\mathcal{L}(W)\right|$. In contrast, on the PTB dataset, the performance enhancement provided by activation is generally more noticeable than on WikiText2. For instance, with 50% parameters pruned, the model pruned with $\left|W\nabla\mathcal{L}(W)X\right|$ achieves a perplexity of 62.84 on the PTB dataset. Compared to using $\left|W\nabla\mathcal{L}(W)\right|$, the perplexity result decreased by 3.53. It is worth noting that activations can be estimated using a small set of calibration data and executed in a single forward pass. Therefore, computing $\left|W\nabla\mathcal{L}(W)X\right|$, as opposed to $\left|W\nabla\mathcal{L}(W)\right|$, almost does not require additional GPU memory overhead.

Efficacy of activations on BLOOM-7B and OPT-6.7B Table 9, Table 10 and Table 11 show the zero-shot results of the pruned models with/without fine-tuning to discover the effectiveness of activations for guiding pruning LLMs. From the results in Table 9, it is evident that incorporating activations $|W\nabla\mathcal{L}(W)X|$ results in an overall enhancement of performance compared to the standard pruning criterion $|W\nabla\mathcal{L}(W)|$. For example, at a 10% pruning ratio, the model pruned with $|W\nabla\mathcal{L}(W)X|$ achieves a perplexity of 35.53 on the WikiText2 dataset for pruning BLOOM-7B. This result surpasses the 38.12 perplexity achieved by the model compressed by using $|W\nabla\mathcal{L}(W)|$. For OPT-6.7B, the performance enhancement provided by activation is also noticeable. For instance, with 30% parameters pruned, the model pruned with $|W\nabla\mathcal{L}(W)X|$ achieves a perplexity of 64.58 on PTB. Compared to using $|W\nabla\mathcal{L}(W)|$, the perplexity result decreased by 0.51.