Pruning as a Domain-specific LLM Extractor (2024)

Abstract: Large Language Models (LLMs) have exhibited remarkable proficiency across a wide array of NLP tasks. However, the escalation in model size also engenders substantial deployment costs. While few efforts have explored model pruning techniques to reduce the size of LLMs, they mainly center on general or task-specific weights. This leads to suboptimal performance due to lacking specificity on the target domain or generality on different tasks when applied to domain-specific challenges. This work introduces an innovative unstructured dual-pruning methodology, D-Pruner, for domain-specific compression on LLM. It extracts a compressed, domain-specific, and task- agnostic LLM by identifying LLM weights that are pivotal for general capabilities, like linguistic capability and multi-task solving, and domain-specific knowledge. More specifically, we first assess general weight importance by quantifying the error incurred upon their removal with the help of an open-domain calibration dataset. Then, we utilize this general weight importance to refine the training loss, so that it preserves generality when fitting into a specific domain. Moreover, by efficiently approximating weight importance with the refined training loss on a domain-specific calibration dataset, we obtain a pruned model emphasizing generality and specificity. Our comprehensive experiments across various tasks in healthcare and legal domains show the effectiveness of D-Pruner in domain-specific compression. Our code is available at https://github.com/psunlpgroup/D-Pruner.

## introduction large language models (llms) such as the gpt family @xcite and the llama family @xcite have exhibited remarkable advancements across a diverse spectrum of nlp † work done as a research intern at nec labs america. here generality refers to the general capabilities of an llm such as language understanding and generation, and multi-task solving, and specificity refers to the capability of an llm to understand domain-specific knowledge. as shown in figure 1 , the weights in an llm work together to support its general capabilities and to store various domain knowledge. the domain-shared weights (or general weights) empower the llm with linguistic and multi-task solving prowess akin to human language usage and thinking. the domain-specific weights (or domain weights) are pivotal for endowing the llm with domain-specific expertise mirroring that of domain experts. however, the current pruning methods mainly focus on preserving general or taskspecific weights, which may not be enough to deal with domain-specific problems. for example, posttraining pruning methods @xcite assume the model is optimized and prune unimportant weights based on an open-domain calibration dataset. this leads to a pruned model that focuses on model generality with domain-specific weights not considered. on the other hand, pruning with fine-tuning methods @xcite utilizes gradients during fine-tuning on a specific task to estimate the importance of parameters. as a result, the pruned model focuses on the model specificity while decreasing the linguistic and multi-task solving capabilities, compromising the llm's capacity as a versatile task-agnostic solver. to this end, this study introduces a novel dualpruning approach, d-pruner, for domain-specific unstructured pruning on llms, which aims to extract a domain-specific llm from the foundation llm. this extracted model is able to solve different tasks in the target domain and facilitates further domain-specific fine-tuning. d-pruner is designed to harness calibration data for guiding llm pruning processes while preserving generality and specificity for multi-task solving and domain challenges. the resulting compressed llm can be seamlessly adapted to the target domain, enabling deployment with limited computing resources. specifically, d-pruner adeptly captures and retains both general and domain parameters while selectively eliminating insignificant model parameters. this mechanism comprises the following steps: firstly, a general weight importance module operates to assess the significance of model parameters for general capabilities. subsequently, we propose an updated training loss function based on the autoregressive training objective for the next token prediction by integrating the general importance as a regularization term. this way, we iden-tify weights contributing to both generality and domain specificity when training on a domain calibration dataset. then, with the updated loss function, we compute the weight importance leveraging gradients without updating the model. moreover, an approximation algorithm, empirical fisher @xcite , is utilized to compute the weight importance efficiently for pruning. we evaluate the performance of d-pruner on llama2 @xcite , a widely adopted open-source llm. our experimental findings demonstrate that d-pruner exhibits remarkable efficiency in the extraction of sparse domain networks from pre-trained llms, with a limited amount of calibration data provided. remarkably, d-pruner achieves comparable results to the full dense model while achieving 50% sparsity, surpassing the performance of alternative pruning techniques across diverse domain-specific datasets in healthcare and legal domains encompassing language comprehension, question answering, and summarization tasks. ## related work model compression involves transforming a large, resource-intensive model into a compact version suitable for low-resource deployment @xcite . there are mainly three techniques for model compression, which are pruning, knowledge distillation, and quantization. pruning. pruning techniques in neural networks can be broadly classified into structured pruning and unstructured pruning @xcite @xcite . structured pruning entails the removal of entire network components, such as channels or layers, guided by specific criteria, while maintaining the overall network architecture. in contrast, unstructured pruning targets individual weights, leading to an irregular sparse structure. while numerous attempts have been made to prune language models of relatively small scales, such as bert (kenton and toutanova, 2019), scant attention has been devoted to pruning llms containing billions of parameters. these larger models possess 100-1000 times more weights, rendering the pruning task significantly more challenging. sparsegpt @xcite , a post-training method for large language models (llms), lacks the capability to identify crucial weights tailored to specific domains or tasks as it refrains from fine-tuning. on the other hand, llm-pruner @xcite employs gradient-based techniques for pruning. however, it falls short in identifying pivotal weights essential for domainshared knowledge, resulting in pruned models that lack the desired level of generality. the existing pruning methods either focus on general or domain-specific weights, yet none of them consider preserving both at the same time. to the best of our knowledge, we are the first to work on pruning llms while preserving weights important to both generality and specificity. knowledge distillation. knowledge distillation (kd) has emerged as a powerful technique, drawing considerable interest for its ability to augment model performance and enhance generalization capacities @xcite . at its core, kd revolves around the transfer of expertise from a complex model, referred to as the "teacher model", to a simplified counterpart known as the "student model". this intricate process of knowledge transfer aims to distill the profound insights encapsulated within the teacher models, condensing them into a more concise and efficient representation within the student models. while kd has been proven a powerful tool for model compression, it needs specific downstream tasks and a large amount of data for the student models to learn from the teacher models. thus, the output that student models produce mainly focuses on a specific task and loses the generality capability. kd generally sets higher requirements on data availability and computation budgets (e.g., gpu memory) than pruning. ## methodology to preserve both generality and specificity on the pruned model, our dual-pruning method d-pruner considers weights important to both generality and specificity during training on a calibration dataset. note we only use the weight gradient generated from the training process but do not update the model weights. our model is pruned in a task-agnostic fashion (e.g., we adopted a pretraining objective, next token prediction, as a part of training loss) so that the pruned model can solve different tasks in the target domain. d-pruner comprises the following steps: firstly, a general weight locating module operates to assess the significance of model parameters for general understanding (section 3.1). subsequently, an updated loss function for the training process is proposed by integrating the general weight importance as a regularization term. this way, we identify weights contributing to both general and domain knowledge (section 3.2). finally, with the updated loss function, we compute the weight gradients on a small domain calibration dataset without updating the model and approximate our dual-pruning weight importance by utilizing the empirical fisher index @xcite for pruning (section 3.3). our method concentrates on unstructured pruning in a layer-by-layer manner for the transformers model. we consider query, key, value, and output projections of all self-attention layers and gate @xcite , down, and up projections of all mlp (multilayer perceptron) layers for pruning. ## experiment setup we evaluate d-pruner on two knowledgeintensive domains, which are healthcare and legal. for model generality under domain-specific challenges, we evaluate the linguistic capability using domain text generation, and evaluate the multi-task solving capability on different domain tasks, i.e., natural language inference (nli), question answering (qa), and summarization. since we use domain datasets, the model specificity on domains can also be evaluated. in addition, we fine-tune the pruned model on domain datasets to further evaluate the generality and specificity. we evaluate d-pruner on the llama2 model family, which is the most used open-source llm. we mainly apply our pruning method and baseline methods to llama2-7b and llama2-13b to show our results. our method can also be easily applied to other llms with different sizes and architectures. for instance, appendix b shows further experiment on bloom model @xcite . ## results and analysis our results and analysis aim to answer the following research questions: ## conclusion we introduce d-pruner, an innovative unstructured dual-pruning method for domain-specific compression on llm. it is able to extract a compressed, domain-specific, and task-agnostic llm by identifying weights that are pivotal for both generality and specificity. more specifically, the general weight importance is first assessed by quantifying the error incurred upon their removal with the help of open-domain calibration data. then, we utilize this general weight importance to refine our training loss, so that it considers generality when fitting into a specific domain. moreover, by efficiently approximating weight importance with the refined training loss on a domain-specific calibration dataset, we obtain a pruned model emphasizing general capabilities and domain-specific knowledge. our comprehensive experiments across various tasks in different domains show the effectiveness of d-pruner in domain-specific pruning.