模型压缩-方案(四)：模型蒸馏【先用完整复杂模型基于训练集训练出来一个老师模型，然后设计一个小的学生模型，再固定老师模型的权重参数，然后设计一系列Loss，让学生在蒸馏中的预测精度逐渐逼近老师】

u013250861

已于 2022-07-19 21:30:36 修改

阅读量5.7k

点赞数 2

分类专栏： AI/模型量化

于 2022-03-05 02:41:02 首次发布

本文链接：https://blog.csdn.net/u013250861/article/details/123288581

版权

模型蒸馏 BERT 模型压缩 Loss设计 teacher-student模型

关键词由CSDN通过智能技术生成

AI/模型量化专栏收录该内容

18 篇文章 3 订阅

订阅专栏

首先介绍模型蒸馏的概念，模型蒸馏是一种模型压缩的方法，由 Hinton 在论文《Distilling the Knowledge in a Neural Network》中提出，下面是模型蒸馏的要点：

首先需要训练一个大的模型，这个大模型也称为 teacher 模型。
利用 teacher 模型输出的概率分布训练小模型，小模型称为 student 模型。
训练 student 模型时，包含两种 label，soft label 对应了 teacher 模型输出的概率分布，而 hard label 是原来的 one-hot label。
模型蒸馏训练的小模型会学习到大模型的表现以及泛化能力。

名词解释

teacher - 原始模型或模型ensemble
student - 新模型
transfer set - 用来迁移teacher知识、训练student的数据集合
soft target - teacher输出的预测结果（一般是softmax之后的概率）
hard target - 样本原本的标签
temperature - 蒸馏目标函数中的超参数
born-again network - 蒸馏的一种，指student和teacher的结构和尺寸完全一样
teacher annealing - 防止student的表现被teacher限制，在蒸馏时逐渐减少soft targets的权重

模型蒸馏的主要流程是：

先用完整复杂模型使用训练集训练出来一个teacher模型，
然后设计一个小规模的student模型，
再固定teacher模型的权重参数，然后设计一系列Loss，让student模型在蒸馏学习的过程中逐渐向teacher模型的表现特性靠拢，使得student模型的预测精度逐渐逼近teacher模型。

其中，专门针对Bert模型的蒸馏方法有很多，如tinybert，distillBert,pkd-bert等等。虽然有这么多蒸馏方法，但是仔细研究也能发现它们或多或少都有一些共同点，例如：

1、在预训练阶段使用蒸馏方法通常能够取得较好的效果。
2、设计的Loss都有一些共通性。
3、会将模型架构模块化，然后对模型不同的模块设计不同的Loss。

下面我就从上述几个角度分别总结一下。

主要关注的是将蒸馏方法仅作用在Finetune阶段。经过实验，发现仅在对Finetune后的原始模型进行蒸馏，很难保持原始的精度，或多或少都会有一定程度的精度损失。

我们能做的就是在inference性能和inference精度两边做一定的平衡。

比如要考虑裁剪的bert层数，裁剪的中间层神经元数、注意力头数等，通常裁剪得越多，inference的精度损失就越大。

按照我之前两篇文章中的方法，使用layerdrop裁剪一半的层数会有8-10个百分点的下降，而使用bert-theseus方法裁剪一半层数，会有2个百分点的下降。对于bert-theseus来说，完全可以应用到实际的项目服务中。

本次使用蒸馏方式在Finetune阶段裁剪模型，在裁剪一半层数的情况下，精度下降的幅度从1个百分点到5个百分点之间浮动，下面会分别具体介绍不同方法带来的结果。

一、蒸馏的Loss设计

bert蒸馏中的Loss设计可以说是其精髓，这里就结合上面2,3两点一起来介绍一下。总结一下，对于当前所有的bert模型，主要设计的Loss的模块集中在Output层的Logits输出（或者softmax概率化后的输出），中间层的hidden_Output,attention_Output,embedding神经元等。

1、Logits的Loss设计

对于Logits来说，通常使用mean squared error来计算两个Logits之间的差异性。

但是对于模型的不同组件，计算两个Logits的Loss方式也有细微的不同。

以NER这种序列标注任务为例，在对中间隐层和Output层的Logits计算MSE时，只需要考虑正常的batch中的序列mask，不要将所有序列step中的padding部分都计算MSE就可以了。

但是对于attention的Output（在bert代码中为attention_scores,即softmax概率化之前的attention计算结果），其shape为 [batch_size,head,seq_len,seq_len]，需要考虑最后两个维度上的mask，

这里参考了TextBrewer（TextBrewer）中的实现，其官方的pytorch代码如下：

def att_MSE_Loss(attention_S, attention_T, mask=None):
    '''
    * Calculates the MSE Loss between `attention_S` and `attention_T`.
    * If the `inputs_mask` is given, masks the positions where ``input_mask==0``.

    :param torch.Tensor Logits_S: tensor of shape  (*batch_size*, *num_heads*, *length*, *length*)
    :param torch.Tensor Logits_T: tensor of shape  (*batch_size*, *num_heads*, *length*, *length*)
    :param torch.Tensor mask: tensor of shape  (*batch_size*, *length*)
    '''
    if mask is None:
        attention_S_select = torch.where(attention_S <= -1e-3, torch.zeros_like(attention_S), attention_S)
        attention_T_select = torch.where(attention_T <= -1e-3, torch.zeros_like(attention_T), attention_T)
        Loss = F.MSE_Loss(attention_S_select, attention_T_select)
    else:
        mask = mask.to(attention_S).unsqueeze(1).expand(-1, attention_S.size(1), -1) # (bs, num_of_heads, len)
        valid_count = torch.pow(mask.sum(dim=2),2).sum()
        Loss = (F.MSE_Loss(attention_S, attention_T, reduction='none') * mask.unsqueeze(-1) * mask.unsqueeze(2)).sum() / valid_count
    return Loss

大家注意一个细节，当给定mask时，在计算valid_count时，作者使用的是 $\sum(\sum_{k}^{}{mask[:,:,k]})^2$ 。而在tinybert中的实现则是这样的：

for student_att, teacher_att in zip(student_atts, new_teacher_atts):
    student_att = torch.where(student_att <= -1e2, torch.zeros_like(student_att).to(device),
                                                  student_att)
    teacher_att = torch.where(teacher_att <= -1e2, torch.zeros_like(teacher_att).to(device),
                                                  teacher_att)

    tmp_Loss = Loss_MSE(student_att, teacher_att)

这个与TextBrewer中未给出mask时的计算方式相同。两种方式我都尝试过，对于最终模型的精度基本上没有太大区别，两种实现方式效果都差不多。

2、概率分布的Loss设计

也有一些方法专门针对概率化后的信息输出计算其Loss，比如每一层attention_score概率化后的alignment、最终模型输出的概率化结果。

通常来说，会使用交叉熵或者KL-divergence等方法计算两个概率分布之间的差异。

对于Output的概率输出来说，在计算两个模型输出之间的交叉熵之前，需要先对模型的概率分布进行一个flat操作。原因在于我们常规的模型学习完成后，它学习到的概率分布都是比较陡的，即某一个或者极少一部分类别的概率会非常大，其余类别的会非常小，因为模型已经学到了一些成熟的信息。在蒸馏时，我们要让student模型学习到teacher模型的概率输出，如果还保持之前的概率分布，那么会让大部分的概率信息无法被学习。因此，通常在对Logits进行概率化之前，要先对Logits除以一个temperature，让不同类别的概率差异稍微变小一点。同样参考TextBrewer中的代码：

def kd_ce_Loss(Logits_S, Logits_T, temperature=1):
    '''
    Calculate the cross entropy between Logits_S and Logits_T

    :param Logits_S: Tensor of shape (batch_size, length, num_labels) or (batch_size, num_labels)
    :param Logits_T: Tensor of shape (batch_size, length, num_labels) or (batch_size, num_labels)
    :param temperature: A float or a tensor of shape (batch_size, length) or (batch_size,)
    '''
    if isinstance(temperature, torch.Tensor) and temperature.dim() > 0:
        temperature = temperature.unsqueeze(-1)
    beta_Logits_T = Logits_T / temperature
    beta_Logits_S = Logits_S / temperature
    p_T = F.softmax(beta_Logits_T, dim=-1)
    Loss = -(p_T * F.log_softmax(beta_Logits_S, dim=-1)).sum(dim=-1).mean()
    return Loss

对于attention概率输出来说，通常不需要对概率分布进行平滑操作，只需要进行正常的交叉熵或者KL-divergence操作，同时要考虑到mask，对于Logits来说很小的负值代表是一个被mask的维度，而对于概率分布来说就不是这个情况了，这个时候最好是能够提供序列的mask。

3、Finetune任务自身Loss

除了几个蒸馏的Loss之外，将下游任务的Loss也加入到模型蒸馏的整体任务中，也能让student模型学习到下游任务的信息。

二、消融实验

我设计了一个消融对比实验，实验因子包括两大类，一类是对不同模型组件的蒸馏，一类是具体的Loss方法，总共包含如下几种情况：

1、是否使用Finetune任务自身Loss。

2、是否使用Attention Output输出Logits的MSE

3、是否使用hidden Output输出Logits的MSE

4、对比使用Output的输出概率的cross entropy和Logits的MSE

5、对比使用Attention Output输出概率的ce和Logits的MSE

其中，我将Output输出Logits的MSE作为基础的蒸馏方法，该Loss一直存在。

在我实验的任务上，结果如下：

1、使用Finetune任务自身的Loss是有效的，但是效果不大，大概能够提升0.2个百分点。

2、使用Attention Output输出Logits的MSE效果甚微，基本没有太大提升。我推测可能是当前对于序列标注任务来说，attention的学习提升不大。建议使用更多不同的任务来实验。

3、使用hidden Output输出Logits的MSE是非常有效的，能够提升1个百分点。

4、使用概率输出做蒸馏和使用Logits输出做蒸馏差距不大，并不能看到显著的区别，建议用更多不同的任务来实验。

二、其他技巧

1、Bert层的映射设计

在设计Attention Output的Loss时，由于会对bert的层数进行裁剪，所以需要对student的encode层和原始模型中的encode层进行映射。

之前有同学介绍过微软的一篇论文：miniLM ,它只使用bert最后一层的value概率输出和attention概率输出做蒸馏，省去了设计映射的工作。我实验过，并没有得到很好的效果，其有效性还待验证，后续会用其他类型的任务来验证。至于如何设计层的映射，目前还没有一个方法论，通常和任务还是相关的，但是有一些指导意见还是可以参考。如bert的每一层所学习存储的信息重点都是不一样的，越接近embedding的底层会倾向于学习通用基础的语言学知识。而接近下游任务分类的上层，则会倾向于学习下游任务中的具体信息。另外间隔的层之间的连通性比较好，因此通常会以间隔的方式建立层映射，如student中的0-5层可以分别对应原模型中的1,3,5,7,9,11层。

2、尽量沿用teacher模型的权重

在进行模型蒸馏时，通常会初始化student模型的权重从头开始训练。但是，如果能让student模型在一开始就用teacher模型的部分权重进行初始化，不仅能够提升学习效率，最后得到的精度也是不错的。通过实验发现，使用teacher模型权重初始化student模型，至少能够带来5个百分点的性能提升。

然而，这种方法也为蒸馏带来的局限性，即我们只能对模型进行模块化的裁剪，如只裁剪整个层或者整个注意力头。如果要裁剪隐层神经元个数，就不能使用这个方法了。如果实际项目服务对于精度要求还是比较高的，那么建议使用这种方式。

3、一步到位不一定有效

这个技巧是看了论文miniLM发现的，论文中它的最终目标是将模型裁剪到4层，hidden_size裁剪一半。

实际操作时，它并非直接使用蒸馏训练一个最小模型，而是先用原始模型蒸馏一个中介模型，其层数为4层，但是hidden_size不变，然后使用这个中介模型作为teacher模型来蒸馏得到最终的模型。

我尝试了这种方式，发现有一定的效果，为了蒸馏得到4层的模型，我先将原始模型蒸馏到6层，然后再蒸馏到4层。

这种方式比直接蒸馏小模型能够有3-4个百分点的提升。

当然，我这里要说明一点，我比较的是训练相同epoch数下的两个模型的精度，也有可能是一步到位蒸馏小模型需要更多的训练步数才能达到收敛，并不能直接断定一步到位为训练法一定就比较差，但至少在相同的训练成本下，采用中介过渡是更有效的。

参考资料：
模型压缩实践收尾篇——模型蒸馏以及其他一些技巧实践小结
 深度神经网络模型蒸馏Distillation
深度学习模型压缩方法（4）-----模型蒸馏（Distilling）与精细模型网络
 模型蒸馏（Model Distillation）
知识蒸馏的过程是怎样的？与迁移学习的区别在哪里？
模型蒸馏（Distil）及mnist实践
 BERT 模型蒸馏 Distillation BERT

u013250861

关注

2
点赞
踩
44

收藏

觉得还不错? 一键收藏
0
评论
模型压缩-方案(四)：模型蒸馏【先用完整复杂模型基于训练集训练出来一个老师模型，然后设计一个小的学生模型，再固定老师模型的权重参数，然后设计一系列Loss，让学生在蒸馏中的预测精度逐渐逼近老师】

名词解释其中，专门针对Bert模型的蒸馏方法有很多，如tinybert，distillBert,pkd-bert等等。下面我就从上述几个角度分别总结一下。主要关注的是将蒸馏方法仅作用在Finetune阶段。经过实验，发现仅在对Finetune后的原始模型进行蒸馏，很难保持原始的精度，或多或少都会有一定程度的精度损失。我们能做的就是在inference性能和inference精度两边做一定的平衡。...
复制链接

扫一扫

专栏目录