导读
本文详细介绍了4种主流的压缩与加速技术:结构优化、剪枝、量化、知识蒸馏,作者分别从每个技术结构与性能表现进行陈述。
近年来,深度学习模型在CV、NLP等领域实现了广泛应用。然而,庞大的参数规模带来的计算开销、内存需求,使得其在计算能力受限平台的部署中遇到了巨大的困难与挑战。因此,如何在不影响深度学习模型性能的情况下进行模型压缩与加速,成为了学术界和工业界的研究热点。
1、简介
深度学习模型压缩与加速是指利用神经网络参数和结构的冗余性精简模型,在不影响任务完成度的情况下,得到参数量更少、结构更精简的模型。被压缩后的模型对计算资源和内存的需求更小,相比原始模型能满足更广泛的应用需求。(事实上,压缩和加速是有区别的,压缩侧重于减少网络参数量,加速侧重于降低计算复杂度、提升并行能力等,压缩未必一定能加速,本文中我们把二者等同看待)
压缩与加速,大体可以从三个层面来做:算法层、框架层、硬件层,本文仅讨论算法层的压缩与加速技术。
2、主流技术
主流的压缩与加速技术有4种:结构优化、剪枝(Pruning)、量化(Quantization)、知识蒸馏(Knowledge Distillation)。
通过优化网络结构的设计去减少模型的冗余和计算量。常见方式如下:
剪枝是指在预训练好的大型模型的基础上,设计对网络参数的评价准则,以此为根据删除“冗余”参数。根据剪枝粒度粗细,参数剪枝可分为非结构化剪枝和结构化剪枝。非结构化剪枝的粒度比较细,可以无限制去掉网络中期望比例的任何“冗余”参数,但会带来裁剪后网络结构不规整难以有效加速的问题。结构化剪枝的粒度比较粗,剪枝的最小单位是filter内参数的组合,通过对filter或者feature map设置评价因子,甚至可以删除整个filter或者某几个channel,使网络“变窄”,可以直接在现有软硬件上获得有效加速,但可能带来预测精度的下降,需要通过对模型微调恢复性能。
量化是指用较低位宽表示典型的32bit浮点型网络参数。网络参数包括权重、激活值、梯度和误差等等, 可以使用统一的位宽(如16bit,8bit,2bit和1bit等),也可以根据经验或一定策略自由组合不同的位宽。量化的优点在于:1).能够显著减少参数存储空间与内存占用空间,如,将参数从32bit浮点型量化到8bit整型能够减少75%的存储空间,这对于计算资源有限的边缘设备和嵌入式设备进行深度学习模型的部署和使用都有很大帮助;2).能够加快运算速度,减少设备能耗,读取32bit浮点数所需的带宽可以同时读入4个8bit整数,并且整型运算相比浮点型运算更快,自然能降低设备功耗。但其仍存在一定的局限性,网络参数的位宽减少损失一部分信息量,造成推理精度下降,虽然能通过微调恢复部分精确度,但带来时间成本;量化到特殊位宽时,很多现有的训练方法和硬件平台不再适用,需要设计专用的系统架构,灵活性不高。
知识蒸馏在Hinton于2015年发表的《Distilling the Knowledge in a Neural Network》(https://arxiv.org/pdf/1503.02531.pdf)中被提出。
同其它压缩与加速技术不同,知识蒸馏需要两种类型的网络:Teacher Model和Student Model。前者参数量大、结构复杂,后者参数量较小、结构相对简单。二者可以是不同的网络结构,但是采用相似的网络结构,蒸馏效果会更好。训练流程如下图所示:
首先训练Teacher Model,然后用其指导Student Model的训练(训练集可以保持一致)。具体指导方案,将Teacher Model在Softmax层的输出作为数据的soft label(熵更高,信息量更大),Student Model的loss function将是对soft label预测和hard label预测的loss的线性加权和(具体可以参考https://zhuanlan.zhihu.com/p/102038521)。Student Model训练好后,按照常规模型使用即可。
对soft label预测的loss,可以计算两个Model在Softmax层的输出的Cross Entropy,也可以计算两个Model在Softmax层的输入的MSE。
知识蒸馏通过将Teacher Model的知识迁移到Student Model中,使Student Model达到与Teacher Model相当的性能,同时又能起到模型压缩的目的。其局限性在于,由于使用Softmax层的输出作为知识,所以一般多用于具有Softmax层面的分类任务,在其它任务上的表现不好。
3、应用实例
以下列举几个知识蒸馏的应用实例。
DistillBERT在Hugging Face于2019年发表的《DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter》(https://arxiv.org/pdf/1910.01108.pdf)中被提出。
DistillBERT的细节如下:
,其中,
是soft label之间的KL散度(非交叉熵),
同BERT,
是隐层向量之间的cosine值
DistillBERT在GLUE数据集上的性能表现:
DistillBERT在参数量和推理速度上的表现:
TinyBERT在华科+华为于2019年发表的《TinyBERT: Distilling BERT for Natural Language Understanding》(https://arxiv.org/pdf/1909.10351.pdf)中被提出。
TinyBERT的训练流程如下:
整体是个两阶段的蒸馏过程:
TinyBERT的结构如下:
TinyBERT的loss如下:
其中,
表示embedding layer,
表示prediction layer,
有三种形式,代表三种不同的loss:
都是MSE loss,
是Cross Entropy loss,具体计算公式可参考论文。
TinyBERT的性能表现如下:
FastBERT在北大+腾讯于2020年发表的《FastBERT: a Self-distilling BERT with Adaptive Inference Time》(https://arxiv.org/pdf/2004.02178.pdf)中被提出。
FastBERT的结构如下:
它的创新点在于self-distillation mechanism和sample-wise adaptive mechanism:
FastBERT的训练流程如下:
作为soft label,计算Student Classifier的输出
和
之间的KL散度,将所有层的KL散度之和作为total loss:
为了实现Adaptive inference,FastBERT采用了归一化的熵作为Student Classifier决策结果的置信度指标:
其中,
是Student Classifier输出的概率分布,
是类目数量。论文中,给Uncertainty设定的阈值是Speed,二者取值均介于0、1之间。
FastBERT的性能表现如下(当前仅应用于分类任务):