All-in-One Tuning and Structural Pruning for Domain-Specific LLMs

📄 arXiv: 2412.14426v2 📥 PDF

作者: Lei Lu, Zhepeng Wang, Runxue Bao, Mengbing Wang, Fangyi Li, Yawen Wu, Weiwen Jiang, Jie Xu, Yanzhi Wang, Shangqian Gao

分类: cs.CL, cs.AI

发布日期: 2024-12-19 (更新: 2024-12-20)

备注: Updated a typo in the author list;

💡 一句话要点

提出ATP：面向领域LLM的端到端调优与结构化剪枝方法

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 大型语言模型 结构化剪枝 领域特定模型 微调 低秩适应 模型压缩 剪枝决策

📋 核心要点

  1. 现有两阶段剪枝方法在领域微调时无法动态调整剪枝决策，导致性能下降。
  2. ATP通过可训练的剪枝决策生成器，在微调过程中动态识别最优子结构。
  3. 实验表明，ATP在法律和医疗领域任务中，显著优于现有两阶段剪枝方法。

📝 摘要（中文）

现有的领域特定大型语言模型（LLM）剪枝技术通常采用两阶段流程：首先剪枝预训练的通用LLM，然后在特定领域上微调剪枝后的LLM。然而，剪枝决策来源于预训练权重，在微调期间保持不变，即使权重已经更新。因此，这种剪枝决策与微调权重的结合可能不是最优的，导致显著的性能下降。为了解决这些限制，我们提出了ATP：All-in-One Tuning and Structural Pruning，一种统一的单阶段结构化剪枝和微调方法，通过可训练的剪枝决策生成器在微调阶段动态识别当前最优的子结构。此外，考虑到领域特定应用数据的有限性，低秩适应（LoRA）成为微调LLM的常用技术。在ATP中，我们引入了LoRA感知的正向传播和稀疏性正则化，以确保与学习到的剪枝决策相对应的子结构可以在ATP过程后直接移除。在法律和医疗领域的任务中，ATP优于最先进的两阶段剪枝方法。更具体地说，当剪枝LLaMA2-7B和LLaMA3-8B模型40%的参数时，ATP分别恢复了密集模型高达88%和91%的性能。

🔬 方法详解

问题定义：现有针对领域特定LLM的剪枝方法通常采用两阶段流程，即先对预训练的通用LLM进行剪枝，然后再在特定领域的数据集上微调剪枝后的模型。这种方法的痛点在于，剪枝决策是在预训练阶段确定的，并且在后续的微调阶段保持不变，无法适应领域数据的特性，导致剪枝后的模型性能下降。

核心思路：ATP的核心思路是将剪枝和微调过程统一到一个阶段中进行。通过引入一个可训练的剪枝决策生成器，在微调过程中动态地调整剪枝决策，从而找到在当前微调权重下的最优子结构。这种方法能够更好地适应领域数据的特性，提高剪枝后模型的性能。

技术框架：ATP的技术框架主要包括以下几个模块：1) 预训练的LLM；2) 可训练的剪枝决策生成器；3) LoRA模块；4) LoRA感知的正向传播和稀疏性正则化。整个流程如下：首先，使用预训练的LLM作为基础模型。然后，在微调过程中，剪枝决策生成器根据当前的权重动态地生成剪枝决策。同时，使用LoRA模块来降低微调的计算成本。最后，通过LoRA感知的正向传播和稀疏性正则化来保证剪枝后的模型可以直接移除相应的子结构。

关键创新：ATP最重要的技术创新点在于将剪枝和微调过程统一到一个阶段中进行，并通过可训练的剪枝决策生成器动态地调整剪枝决策。与现有方法的本质区别在于，ATP能够根据领域数据的特性动态地调整剪枝决策，从而找到在当前微调权重下的最优子结构，而现有方法的剪枝决策是在预训练阶段确定的，无法适应领域数据的特性。

关键设计：ATP的关键设计包括：1) 可训练的剪枝决策生成器的设计，需要考虑如何有效地生成剪枝决策；2) LoRA感知的正向传播和稀疏性正则化的设计，需要保证剪枝后的模型可以直接移除相应的子结构；3) 损失函数的设计，需要平衡剪枝的稀疏性和模型的性能。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，ATP在法律和医疗领域的任务中，显著优于现有的两阶段剪枝方法。具体来说，当剪枝LLaMA2-7B模型40%的参数时，ATP能够恢复密集模型88%的性能；当剪枝LLaMA3-8B模型40%的参数时，ATP能够恢复密集模型91%的性能。这些结果表明，ATP能够有效地减少模型参数量，同时保持较高的模型性能。

🎯 应用场景

该研究成果可广泛应用于各种领域特定的大型语言模型，例如法律、医疗、金融等。通过在保证模型性能的前提下减少模型参数量，可以降低模型的部署成本和推理延迟，使其更易于在资源受限的环境中使用。此外，该方法还可以促进领域知识的迁移和融合，提高模型的泛化能力。

📄 摘要（原文）

Existing pruning techniques for large language models (LLMs) targeting domain-specific applications typically follow a two-stage process: pruning the pretrained general-purpose LLMs and then fine-tuning the pruned LLMs on specific domains. However, the pruning decisions, derived from the pretrained weights, remain unchanged during fine-tuning, even if the weights have been updated. Therefore, such a combination of the pruning decisions and the finetuned weights may be suboptimal, leading to non-negligible performance degradation. To address these limitations, we propose ATP: All-in-One Tuning and Structural Pruning, a unified one-stage structural pruning and fine-tuning approach that dynamically identifies the current optimal substructure throughout the fine-tuning phase via a trainable pruning decision generator. Moreover, given the limited available data for domain-specific applications, Low-Rank Adaptation (LoRA) becomes a common technique to fine-tune the LLMs. In ATP, we introduce LoRA-aware forward and sparsity regularization to ensure that the substructures corresponding to the learned pruning decisions can be directly removed after the ATP process. ATP outperforms the state-of-the-art two-stage pruning methods on tasks in the legal and healthcare domains. More specifically, ATP recovers up to 88% and 91% performance of the dense model when pruning 40% parameters of LLaMA2-7B and LLaMA3-8B models, respectively.