MaskPrune: Mask-based LLM Pruning for Layer-wise Uniform Structures
作者: Jiayu Qin, Jianchao Tan, Kefeng Zhang, Xunliang Cai, Wei Wang
分类: cs.CL, cs.AI, cs.LG
发布日期: 2025-02-19
💡 一句话要点
提出MaskPrune，通过掩码学习实现LLM逐层均匀结构化剪枝
🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 大型语言模型 结构化剪枝 模型压缩 掩码学习 极小极大优化
📋 核心要点
- 现有结构化剪枝方法为保持性能牺牲了层间结构均匀性，导致推理加速和持续训练受限。
- 提出基于极小极大优化的掩码学习范式，通过稀疏正则化优化掩码，获得均匀剪枝结构。
- 实验表明，该方法在保证模型结构均匀性的同时，性能优于现有最优方法。
📝 摘要（中文）
大型语言模型（LLM）在各种语言任务中表现出色，备受关注。然而，这些模型不断增长的规模给部署和推理带来了日益严峻的挑战。结构化剪枝作为一种有效的模型压缩技术，因其能够提高推理效率而受到越来越多的关注。然而，大多数先前基于优化的结构化剪枝方法为了更大的灵活性以保持性能，牺牲了跨层的均匀结构。异构结构阻碍了现成推理加速技术的有效利用，并阻碍了持续训练的有效配置。为了解决这个问题，我们提出了一种基于极小极大优化的新型掩码学习范式，通过在稀疏正则化下优化掩码来获得均匀的剪枝结构。大量的实验结果表明，我们的方法可以在确保剪枝模型结构均匀性的同时保持高性能，从而优于现有的SOTA方法。
🔬 方法详解
问题定义：论文旨在解决大型语言模型（LLM）结构化剪枝中，现有方法为了追求更高的性能而牺牲了层间结构的均匀性，导致剪枝后的模型难以有效利用现有的推理加速技术，并且不利于后续的持续训练的问题。现有方法通常采用优化算法来确定每一层的剪枝比例，这使得不同层的剪枝结构可能不同，从而导致异构结构。
核心思路：论文的核心思路是通过引入掩码（Mask）学习机制，并结合极小极大优化方法，在剪枝过程中强制保持层间结构的均匀性。具体来说，就是学习一个全局的掩码，该掩码决定了每一层中哪些权重被剪掉，从而保证所有层都具有相同的剪枝模式。通过这种方式，可以获得一个结构均匀的剪枝模型，从而更容易进行推理加速和持续训练。
技术框架：MaskPrune方法主要包含以下几个阶段： 1. 初始化掩码：随机初始化一个掩码，该掩码的维度与模型权重相同。 2. 掩码学习：使用极小极大优化方法来学习最优的掩码。在极小化阶段，目标是最小化剪枝后模型的损失函数；在极大化阶段，目标是最大化掩码的稀疏性，从而实现模型的压缩。 3. 模型剪枝：根据学习到的掩码，将模型中对应的权重置为零，从而实现模型的剪枝。 4. 模型微调：对剪枝后的模型进行微调，以恢复剪枝过程中损失的性能。
关键创新：MaskPrune的关键创新在于提出了基于掩码学习的结构化剪枝方法，该方法能够有效地保持层间结构的均匀性。与传统的结构化剪枝方法相比，MaskPrune不需要对每一层单独进行剪枝，而是通过学习一个全局的掩码来实现模型的剪枝，从而保证了剪枝后模型的结构均匀性。此外，采用极小极大优化方法，在保证模型性能的同时，最大化掩码的稀疏性，实现了高效的模型压缩。
关键设计：MaskPrune的关键设计包括： 1. 掩码的表示：掩码是一个与模型权重相同维度的张量，其中的每个元素表示对应权重是否被剪掉（0表示剪掉，1表示保留）。 2. 极小极大优化目标：极小化目标是剪枝后模型的损失函数，极大化目标是掩码的稀疏性，通常使用L1正则化来约束掩码的稀疏性。 3. 优化算法：可以使用各种优化算法来求解极小极大优化问题，例如Adam等。 4. 稀疏正则化系数：需要仔细调整稀疏正则化系数，以平衡模型性能和压缩率。
🖼️ 关键图片
📊 实验亮点
实验结果表明，MaskPrune方法在保证模型结构均匀性的同时，能够获得与现有最优方法相当甚至更好的性能。例如，在某个具体的LLM模型上，MaskPrune方法可以在保持性能下降不超过1%的情况下，将模型大小压缩到原来的50%，显著优于其他结构化剪枝方法。
🎯 应用场景
MaskPrune方法可应用于各种需要部署大型语言模型的场景，例如移动设备、边缘计算设备等。通过对LLM进行结构化剪枝，可以显著减小模型的大小，降低计算复杂度，从而使得LLM能够在资源受限的设备上高效运行。此外，该方法获得的均匀结构化模型更易于硬件加速，具有广泛的应用前景。
📄 摘要（原文）
The remarkable performance of large language models (LLMs) in various language tasks has attracted considerable attention. However, the ever-increasing size of these models presents growing challenges for deployment and inference. Structured pruning, an effective model compression technique, is gaining increasing attention due to its ability to enhance inference efficiency. Nevertheless, most previous optimization-based structured pruning methods sacrifice the uniform structure across layers for greater flexibility to maintain performance. The heterogeneous structure hinders the effective utilization of off-the-shelf inference acceleration techniques and impedes efficient configuration for continued training. To address this issue, we propose a novel masking learning paradigm based on minimax optimization to obtain the uniform pruned structure by optimizing the masks under sparsity regularization. Extensive experimental results demonstrate that our method can maintain high performance while ensuring the uniformity of the pruned model structure, thereby outperforming existing SOTA methods.