Structured Optimal Brain Pruning for Large Language Models (2024)

Abstract: The massive parameters and computational demands hinder the widespread application of Large Language Models (LLMs). Network pruning provides a practical solution to this problem. However, existing pruning works for LLMs mainly focus on unstructured pruning or necessitate post-pruning fine-tuning. The former relies on special hardware to accelerate computation, while the latter may need substantial computational resources. In this paper, we introduce a retraining-free structured pruning method called SoBP (Structured Optimal Brain Pruning). It leverages global first-order information to select pruning structures, then refines them with a local greedy approach, and finally adopts module-wise reconstruction to mitigate information loss. We assess the effectiveness of SoBP across 14 models from 3 LLM families on 8 distinct datasets. Experimental results demonstrate that SoBP outperforms current state-of-the-art methods.

## introduction extensive research has shown that large language models (llms) @xcite @xcite possess outstanding language comprehension and text generation capabilities, which provide them with significant potential and value across diverse industries @xcite . however, such high performance often comes with enormous model size, high computational overhead, and significant memory consumption @xcite , making it impractical for many scenarios. many studies in model compression @xcite @xcite have been investigated to tackle this problem. among them, network pruning is an effective technique to compress the model size while maintaining the performance. for llms' pruning, most methods adopt unstructured pruning. they analyze the importance of individual weight, zeroing out unimportant weights to obtain sparse weight matrices, as shown in fig. 1(a ). however, unstructured sparse weights require specialized data structure to store and specific hardware to accelerate computation @xcite . for general computing devices, unstructured pruning does not bring any reduction in storage or computation. another branch of network pruning focuses on structured pruning. structured pruning eliminates several rows or columns of the weight matrix or deletes entire layers from the network, which makes it more general for various computing devices @xcite . however, these coarser pruning granularities suffer from non-negligible performance degradation. therefore, most structured pruning methods require fine-tuning to restore network performance @xcite . given the large number of parameters in llms, it might be incredibly resource-intensive and timeconsuming. to address the above issues, we propose sobp (structured optimal brain pruning), a structured pruning method for llms that does not require fine-tuning. it takes into account both the hardware-friendliness of the pruned network and the efficiency of the pruning process. when considering this problem of retraining-free structured pruning, two principal aspects need to be critically addressed: ❶ finding the appropriate pruning structures. ❷ applying a suitable strategy to recover the performance of the pruned network. regarding the first aspect, considering the setting of structured pruning, we prune the heads in multi-head attention (mha) and the intermediate dimension of feed forward network (ffn), similar to previous studies @xcite . since the importance level varies among different layers in llms @xcite , it is necessary to select pruning structures according to their importance within the network. we use global importance scores, which are based on first-order information, to select the pruning structures of each mha and ffn module. however, first-order information ignores the correlation between different structures, which means that the structures selected by it might fail to minimize the change in local output. so we further refine the pruning structures of each module by a local greedy approach. for the second aspect, since fine-tuning is omitted, we reconstruct the output weights of each module after pruning. our method inherits the legacy of the previous optimal brain series @xcite @xcite , extending them into the field of structured pruning and implementing some transformations to address numerical stability issues. meanwhile, we also introduce a method to compensate for the cumulative errors at a high compression rate. the overview of sobp is shown in fig. 2 . to summarize the contribution of our work: ## preliminary sobp prunes heads of mha and the intermediate dimension of ffn. we refer to them as pruning structures or pruning units and use neurons to denote the intermediate dimension of ffn in the following text. ## method the framework of sobp has three key components: global importance-aware selection, local greedy refinement, and module-wise reconstruction. ## experiments we compare sobp with different methods on generative and zero-shot tasks. meanwhile, we perform ablation studies to evaluate the contribution of each component in sobp, as well as the impact of calibration data and the hyperparameter on the pruned models' performance. finally, we measure the inference time and throughput after pruning. ## conclusion in this paper, we investigate retraining-free structured pruning for large language models. we propose a three-stage pruning algorithm named sobp. it determines the appropriate pruning structures in a global-local manner. by reconstructing the output feature of each module, we preserve the pruned model's performance while avoiding retraining. the effectiveness of this approach is confirmed through extensive experiments across various models.