MultiPruner: Balanced Structure Removal in Foundation Models

📄 arXiv: 2501.09949v1 📥 PDF

作者: J. Pablo Muñoz, Jinjie Yuan, Nilesh Jain

分类: cs.LG, cs.AI

发布日期: 2025-01-17

🔗 代码/项目: GITHUB

💡 一句话要点

MultiPruner：一种用于基础模型中平衡结构移除的多维度剪枝方法

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

关键词: 模型剪枝 预训练模型 结构化剪枝 无训练剪枝 多维度剪枝

📋 核心要点

  1. 现有无训练剪枝方法在大型预训练模型上压缩效果有限，尤其是在保持结构平衡方面存在挑战。
  2. MultiPruner采用多维度迭代剪枝策略，依次压缩残差块、MLP通道和注意力头，恢复模型结构平衡。
  3. 实验表明，MultiPruner在零样本精度和模型压缩率方面优于其他无训练剪枝方法，降低了计算和内存需求。

📝 摘要（中文）

最近，用于剪枝大型预训练模型（LPMs）的先进方法表明，在Transformer中无训练地移除非关键残差块对于减小模型尺寸是可行的，并且优于之前的无训练剪枝方法。受此启发，我们扩展了BlockPruner (Zhong et al., 2024)，并提出了MultiPruner，一种通过采用多维度、迭代、细粒度的剪枝策略来超越最近的无训练剪枝方法。在MultiPruner中，多维度剪枝通过依次压缩三个维度来恢复块剪枝模型中的结构平衡：i) 残差块，ii) 多层感知机（MLP）的通道，以及 iii) 注意力头。与其他技术相比，该解决方案增强了下游任务的零样本精度，同时提高了模型压缩率，从而产生了计算和内存需求更少的压缩模型。大量的实验证明了所提出的方法在各种大型预训练模型上的优势。代码和剪枝配置可在https://github.com/IntelLabs/Hardware-Aware-Automated-Machine-Learning上找到。

🔬 方法详解

问题定义：论文旨在解决大型预训练模型（LPMs）的剪枝问题，特别是如何在不进行训练的情况下，有效地减小模型尺寸，同时保持或提升模型在下游任务上的性能。现有无训练剪枝方法，如BlockPruner，虽然可以移除非关键残差块，但可能导致模型结构失衡，限制了压缩率和性能提升。

核心思路：MultiPruner的核心思路是通过多维度剪枝来恢复模型结构平衡。它认为，简单地移除残差块会导致模型在不同维度上的不平衡，因此需要同时考虑残差块、MLP通道和注意力头三个维度进行剪枝，以实现更优的压缩效果和性能。

技术框架：MultiPruner的整体框架是一个迭代的多维度剪枝流程。它包含以下主要阶段：1) 残差块剪枝：首先移除不重要的残差块。2) MLP通道剪枝：对MLP层的通道进行剪枝，减少参数量。3) 注意力头剪枝：移除不重要的注意力头。这三个阶段依次进行，每次剪枝后都评估模型性能，并根据性能指标调整剪枝策略。整个过程迭代进行，直到达到预定的压缩目标。

关键创新：MultiPruner最重要的技术创新点在于其多维度剪枝策略。与以往只关注单一维度（如残差块）的剪枝方法不同，MultiPruner同时考虑了残差块、MLP通道和注意力头三个维度，从而能够更全面地优化模型结构，恢复模型平衡，提升压缩效果和性能。

关键设计：MultiPruner的关键设计包括：1) 剪枝顺序：先剪枝残差块，再剪枝MLP通道，最后剪枝注意力头。这种顺序是基于经验观察和实验结果确定的。2) 剪枝比例：每个维度上的剪枝比例需要根据具体模型和任务进行调整。论文提供了在不同模型上的剪枝配置。3) 性能评估指标：使用下游任务的零样本精度作为评估指标，指导剪枝过程。4) 迭代次数：迭代次数需要根据压缩目标和性能要求进行调整。过多的迭代可能会导致性能下降。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，MultiPruner在多个大型预训练模型上取得了显著的压缩效果和性能提升。例如，在某些模型上，MultiPruner可以在保持或提升零样本精度的前提下，将模型尺寸压缩到原来的50%以下。与BlockPruner等基线方法相比，MultiPruner在压缩率和性能方面均有明显优势。

🎯 应用场景

MultiPruner可应用于各种需要部署大型预训练模型的场景，例如移动设备、边缘计算设备等资源受限的环境。通过减小模型尺寸和降低计算需求，MultiPruner可以使这些模型在资源有限的设备上高效运行，从而推动人工智能技术在更广泛领域的应用。此外，该方法还可以用于加速模型训练和推理，提高开发效率。

📄 摘要（原文）

Recently, state-of-the-art approaches for pruning large pre-trained models (LPMs) have demonstrated that the training-free removal of non-critical residual blocks in Transformers is viable for reducing model size, achieving results that outperform previous training-free pruning approaches. Motivated by these findings, we extend BlockPruner (Zhong et al., 2024) and propose MultiPruner, a pruning approach that surpasses recent training-free pruning methods by adopting a multidimensional, iterative, fine-grained pruning strategy. In MultiPruner, multidimensional pruning reinstates the structural balance in block-pruned models by sequentially compressing along three dimensions: i) residual blocks, ii) channels of multilayer perceptrons (MLP), and iii) attention heads. This solution enhances zero-shot accuracy on downstream tasks compared to other techniques while improving model compression ratios, producing compressed models with fewer computing and memory requirements. Extensive experiments demonstrate the advantages of the proposed method across various large pre-trained models. The code and pruning configurations are available at https://github.com/IntelLabs/Hardware-Aware-Automated-Machine-Learning.