Transformer模型的训练后剪枝方法_transformer剪枝

Paper地址：https://arxiv.org/abs/2204.09656

GitHub链接：https://github.com/WoosukKwon/retraining-free-pruning

前言

剪枝是一种经典的模型压缩方法，包括非结构剪枝、结构剪枝，它通过减少模型的冗余参数、以达到降低模型计算复杂度的目的。传统剪枝的实现方式包括训练感知方式（如Network Slimming、SFP、Taylor-prune等）与结构搜索方式（如NetAdapt、AMC等），包含三个阶段，即模型预训练、模型剪枝、模型重训练。因此传统剪枝的压缩成本相对较高，依赖训练资源、且训练的时间开销相对较高。

为了避免较重的训练开销，近年来训练后剪枝逐渐成为研究热点。类似于Post-training量化，仅需少量无标注数据的校准（特征对齐、最小化重建误差等校准方法），通过训练后压缩便能获得较好的压缩效果，且压缩成本可控（对训练资源的依赖较轻、时间代价也较低）。

Transformer模型因其特殊的模型结构（MHA+FFN），本文提出了针对性的训练后结构剪枝方法：1）基于Fisher信息设计Mask搜索算法（Mask search），寻找重要性最低的Attention heads或FFN neurons；2）通过Mask重排算法（Mask rearrangement），决定最终的（0-1取值）的剪枝Mask；3）基于少量无标注数据的知识蒸馏实现Mask微调（Mask tuning），获得Soft-mask以保持剪枝后模型的精度。

最终，针对BERT-base与DistilBERT，通过本文提出的训练后剪枝方法能够实现有效的结构剪枝，在GLUE与SQuAD数据集上，能够实现2x的F