Transformer模型的训练后剪枝方法

Law-Yao

已于 2022-10-07 09:39:09 修改

阅读量3.7k

点赞数 1

分类专栏： Transformer 模型压缩优化加速文章标签：剪枝深度学习模型压缩 transformer 自然语言处理

于 2022-10-07 09:32:17 首次发布

本文链接：https://blog.csdn.net/nature553863/article/details/127190452

版权

该文介绍了一种针对Transformer模型的训练后剪枝方法，通过Fisher信息的Mask搜索、重排和微调算法，实现了高效压缩，减少了FLOPS并保持了精度。在BERT-base和DistilBERT上，压缩比例达2x，推理速度提升1.56x，精度损失小于1%，训练后剪枝只需3分钟，压缩过程轻量且资源友好。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

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的FLOPS降解、1.56