[论文解读] Rethinking Network Pruning -- under the Pre-train and Fine-tune Paradigm

Dongkuan Xu, Ian En-Hsu Yen+2 more|arXiv (Cornell University)|Apr 18, 2021

Topic Modeling参考文献 27被引用 6

一句话总结

本文提出 SparseBERT，一种在 BERT 微调阶段应用的、具备知识感知的稀疏剪枝方法，通过使用微调后的教师模型进行知识蒸馏，以保留通用语言知识和任务特定知识。该方法在 GLUE 上实现 20 倍参数量和 FLOPs 压缩，平均准确率仅下降 1.4%，优于先前的剪枝方法，包括结构化剪枝和基于蒸馏的模型。

ABSTRACT

Transformer-based pre-trained language models have significantly improved the performance of various natural language processing (NLP) tasks in the recent years. While effective and prevalent, these models are usually prohibitively large for resource-limited deployment scenarios. A thread of research has thus been working on applying network pruning techniques under the pretrain-then-finetune paradigm widely adopted in NLP. However, the existing pruning results on benchmark transformers, such as BERT, are not as remarkable as the pruning results in the literature of convolutional neural networks (CNNs). In particular, common wisdom in pruning CNN states that sparse pruning technique compresses a model more than that obtained by reducing number of channels and layers (Elsen et al., 2020; Zhu and Gupta, 2017), while existing works on sparse pruning of BERT yields inferior results than its small-dense counterparts such as TinyBERT (Jiao et al., 2020). In this work, we aim to fill this gap by studying how knowledge are transferred and lost during the pre-train, fine-tune, and pruning process, and proposing a knowledge-aware sparse pruning process that achieves significantly superior results than existing literature. We show for the first time that sparse pruning compresses a BERT model significantly more than reducing its number of channels and layers. Experiments on multiple data sets of GLUE benchmark show that our method outperforms the leading competitors with a 20-times weight/FLOPs compression and neglectable loss in prediction accuracy.

研究动机与目标

• 为解决在 BERT 等预训练语言模型中，稀疏剪枝效果不如结构化剪枝的问题。
• 在剪枝过程中同时保留预训练阶段的通用语言知识和微调阶段的任务特定知识。
• 证明稀疏剪枝在 BERT 中可实现远高于通道/层剪枝的压缩率。
• 开发一种在极端稀疏度下仍能保持高性能的剪枝策略，尤其适用于资源受限环境。
• 在标准基准（GLUE、SQuAD）和支持稀疏推理的硬件平台上验证该方法。

提出的方法

• 剪枝在微调阶段应用，以预训练 BERT 作为学生模型，以微调后的 BERT 作为教师模型。
• 通过 logits 蒸馏和损失最小化，使用知识蒸馏将教师模型中的任务特定知识迁移至剪枝后的学生模型。
• 对自注意力和前馈网络中的线性变换施加稀疏性，这些组件被识别为参数量过大且计算密集。
• 避免剪枝注意力头或嵌入层，以保持结构和表征完整性。
• 将剪枝过程整合进微调循环，实现稀疏度与性能的联合优化。
• 在 Moffett AI 的 ANTOM 平台（支持原生稀疏张量）上进行硬件评估，以测量实际推理加速效果。

实验结果

研究问题

• RQ1 在预训练与微调范式下，稀疏剪枝能否在 BERT 中实现高于结构化剪枝（如通道剪枝）的压缩率？
• RQ2 如何在剪枝过程中保留来自预训练和微调阶段的知识，以避免性能下降？
• RQ3 在微调阶段结合知识蒸馏，能否实现高效稀疏剪枝并同时在极端稀疏度下保持高准确率？
• RQ4 与在预训练阶段或无蒸馏的微调阶段剪枝相比，稀疏剪枝的性能表现如何？
• RQ5 在具备原生稀疏张量支持的硬件平台上，稀疏 BERT 模型的实际推理加速效果如何？

主要发现

• SparseBERT 在 BERT-base 上实现 20 倍压缩（保留 5% 参数），在 GLUE 开发集上平均准确率仅下降 1.4%。
• 在 MRPC 数据集上，SparseBERT 将参数量从 8553 万降至 484 万，FLOPs 从 10.87 GFLOPs 降至 0.54 GFLOPs，实现 20 倍压缩。
• 在 GLUE 基准测试中，SparseBERT 在相似或更低稀疏度下优于所有主流竞争模型，包括 TinyBERT、DistilBERT 和 BERT-of-Theseus。
• 在 SQuAD v1.1 和 v2.0 上，SparseBERT 持续优于在下游微调阶段应用的剪枝方法，尤其在高稀疏度下表现更优。
• 在 Moffett AI 的 ANTOM 硬件上，SparseBERT 在 20 倍压缩下实现超过 10 倍推理加速，且在 4 倍压缩内接近线性加速。
• 与在预训练阶段剪枝相比，SparseBERT 展现出更强的拟合能力，表现为在相同稀疏度下训练损失更低、RTE 数据集上开发集准确率更高。