[论文解读] MLPruning: A Multilevel Structured Pruning Framework for Transformer-based Models

Zhewei Yao, Linjian Ma+3 more|arXiv (Cornell University)|May 30, 2021

Topic Modeling参考文献 39被引用 7

一句话总结

MLPruning 提出了一种针对基于 Transformer 的模型的多级结构化剪枝框架，通过使用可学习的 Top-k 阈值实现自适应正则化，结合了头剪枝、行剪枝和块级稀疏剪枝。该方法在 GLUE 基准测试上实现了高达 3.69 倍的推理加速，仅保留原始模型约 20% 的参数量，且准确率损失极小。

ABSTRACT

Pruning is an effective method to reduce the memory footprint and computational cost associated with large natural language processing models. However, current approaches either only explore head pruning, which has a limited pruning ratio, or only focus on unstructured pruning, which has negligible effects on the real inference time and/or power consumption. To address these challenges, we develop a novel MultiLevel structured Pruning (MLPruning) framework, which uses three different levels of structured pruning: head pruning, row pruning, and block-wise sparse pruning. We propose using a learnable Top-k threshold, which employs an adaptive regularization to adjust the regularization magnitude adaptively, to select appropriate pruning ratios for different weight matrices. We also propose a two-step pipeline to combine block-wise pruning with head/row pruning to achieve high structured pruning ratios with minimum accuracy degradation. Our empirical results show that for \bertbase, with extapprox20\% of remaining weights, \OURS can achieve an accuracy that is comparable to the full model on QQP/MNLI/\squad, with up to extapprox3.69x speedup. Our framework has been open sourced~\cite{codebase}.

研究动机与目标

• 为解决现有剪枝方法的局限性，这些方法要么仅关注头剪枝（剪枝率较低），要么采用非结构化剪枝（在实际应用中效率提升有限）。
• 开发一种结构化剪枝框架，实现在保持模型准确率的同时实现高剪枝率。
• 将多个剪枝层级——头级、行级和块级——整合到统一且高效的 Transformer 模型剪枝流水线中。
• 设计一种基于可学习 Top-k 阈值的自适应正则化机制，以动态调整不同权重矩阵的剪枝强度。

提出的方法

• 该框架采用三级结构化剪枝：头剪枝用于移除注意力头，行剪枝用于消除前馈层中的整行，块级稀疏剪枝则以结构化方式在权重矩阵中强制实现稀疏性。
• 引入可学习的 Top-k 阈值，以在训练过程中自适应地控制正则化强度，实现对每个权重矩阵的最优剪枝率的动态选择。
• 设计了两步剪枝流水线：首先应用头剪枝和行剪枝，然后将二者与块级剪枝结合，以在最小化准确率下降的同时最大化稀疏度。
• 该框架采用基于权重矩阵重要性的自适应正则化机制，确保训练过程稳定且剪枝效果高效。

实验结果

研究问题

• RQ1 多级结构化剪枝框架是否能在保持模型准确率的同时，实现高于仅头剪枝或非结构化剪枝方法的剪枝率？
• RQ2 可学习的 Top-k 阈值在根据不同权重矩阵自适应调整正则化强度方面，对实现最优剪枝性能有多大的有效性？
• RQ3 在两步流水线中结合头级、行级和块级剪枝，对推理速度和准确率的影响如何？
• RQ4 结构化剪枝在多大程度上能减少模型大小和推理时间，同时在 NLP 基准测试上保持极低的准确率损失？

主要发现

• MLPruning 在 BERT-base 模型上实现了高达 3.69 倍的推理加速，且仅保留原始模型 20% 的参数量。
• 在 QQP、MNLI 和 SQuAD 基准测试上，剪枝后的模型在显著减少参数量的情况下，仍保持与完整模型相当的准确率。
• 使用可学习的 Top-k 阈值实现了自适应正则化，提升了剪枝效率和训练过程中的模型稳定性。
• 两步流水线有效结合了头/行剪枝与块级稀疏性，实现了高结构化剪枝率，且准确率损失极小。