News

Date: 2022-12-30来源:沐曦AI解决方案 董华楠

一、引言


知识蒸馏是模型压缩的常用方法,需要构建两个模型分别作为教师模型和学生模型,利用教师模型来引导学生模型获取接近于教师模型的精度。一般情况下,教师的网络结构复杂具有丰富的知识,而学生的网络结构简单,教师的性能是明显优于学生的,但近日微软AI(Swin transformer原团队)在论文Contrastive Learning Rivals Masked Image Modeling in Fine-tuning via Feature Distillation 中提出了一种简单的蒸馏方法,在不改变网络结构的情况下使得学生模型的性能优于教师模型1~2个百分点,该方法还夺得了视觉目标检测冠军。本文主要介绍该特征蒸馏方法以及我们在ASR(Automatic Speech Recognition)的前处理模块VAD(Voice Activity Detection)中的相关实践。

二、特征蒸馏方法介绍


基于图像掩码(MIM, Masked Image Modeling)的自监督方法在很多任务上都有不错的微调fine-tuning的性能,这导致基于对比学习方法得到的预训练模型黯然失色。原论文通过简单的特征蒸馏(FD, feature distillation)将基于对比学习方法得到的模型转换成和MIM相似性能的新模型,使得新模型的fine-tuning性能可以与MIM相媲美甚至更好。表1是论文中对不同的对比学习方法的预训练模型进行特征蒸馏后再fine-tuning的效果。特征蒸馏均在ImageNet-1K训练集上进行。从表中可知各个对比学习的模型在特征蒸馏之后在ImageNet-1K上fine-tuning的性能可提升1~2%,在ADE20K上可提升1~3.3mIoU,特征蒸馏后的性能提升相当可观。

表1:特征蒸馏前后效果对比
来源:[2205.14141] Contrastive Learning Rivals Masked Image Modeling in Fine-tuning via Feature Distillation (arxiv.org)


论文中提出的特征蒸馏方法非常简单,如图1所示,其中预训练模型是教师模型,而新的模型是学生模型。与一般的蒸馏有所不同,文中的教师模型和学生模型的网络结构是相同的,只是教师模型的参数是预训练好的,而学生模型的参数是随机初始化的。

图1:特征蒸馏方法示意图
来源:[2205.14141] Contrastive Learning Rivals Masked Image Modeling in Fine-tuning via Feature Distillation (arxiv.org)


蒸馏过程主要有下面4个要点:

1、以特征图为蒸馏目标

采用模型输出的特征图作为目标,使得没有logits输出的预训练模型也可以进行特征蒸馏,如CLIP。为了让教师模型和学生模型的特征图具有可比性,每个原始图像采用相同的增强视图,并且在学生网络上加入了projector,即1x1的卷积层,使得该方法具有普适性,即使教师模型和学生模型输出不一致也可以使用。采用特征图作为蒸馏目标相比logits和GAP特征具有更好的性能,见表2。

2、白化(Whitening)教师模型的特征

不同的预训练模型具有不同的特征量级,使得不同的模型在蒸馏时的超参数都需要调整,故引入白化操作来避免此现象。文中的白化操作是将特征归一化同一量级,具体是用不含参数的layer Normalization层(去掉scaling和bias)实现的。文中使用 smooth l1 loss作为损失函数。s和t分别表示学生模型和教师模型输出的特征,g表示1x1的卷积,β是数值为2的常量


3、共享相对位置编码

在原始ViT中,相对位置编码并没有比绝对位置编码展现出更好的性能,但在特征蒸馏中,作者对比了不同的策略:相对位置编码(RPB),绝对位置编码(APE)以及共享相对位置编码,实验发现采用共享相对位置编码效果更好。

4、采用非对称的 drop path rate

作者发现对学生模型和教师模型使用不同的drop path rate有利于学生模型学到更好的表征。教师不采用drop path,即drop path rate设置为0,而学生模型在{0.1, 0.2, 0.3}中选择最优的drop path rate的效果是最优的。这个策略的效果是显然的,学生模型中引入drop path rate,让学生模型中的某几层网络随机失活,避免过拟合,使得学生更好地学习教师的知识。

表2是上述蒸馏要点的消融实验结果,从表中可知上述4个要点都可在一定程度上提升学生模型的性能。表中的蒸馏和finetune实验都是在ImageNet-1K的数据集上进行的。但此处有个小缺点,作者没有做这4个要点相互叠加的消融实验。

表2:消融实验结果
来源:[2205.14141] Contrastive Learning Rivals Masked Image Modeling in Fine-tuning via Feature Distillation (arxiv.org)


三、蒸馏前后的不同表征


论文通过对特征蒸馏前后模型的特征可视化分析来说明该方法的有效性。

1. 特征蒸馏多样化 attention heads

图2中左边两列从上到下3行分别为DINO,DeiT和CLIP模型每个layer的每个head的平均attention距离,第一列为蒸馏前,第二列为蒸馏后。右侧两列为每个layer的attention head的余弦相似度。平均attention距离表示attention head的感受野大小。从图2中不难看出,特征蒸馏之前,随着层数的加深,不同head的attention距离范围变得更加接近,最后落在很小的范围内,说明不同的head学到了相似的特性,这是对模型能力的浪费。但蒸馏之后,attention距离变得更加多样或更加均匀地分布,尤其在更深层中,这一现象更加明显。在图2右侧的余弦相似度中也呈现出类似的效果。

图2:ViT-B的每个layer的每个head的平均attention距离 / ViT-B的每个layer的attention head的余弦相度
来源:[2205.14141] Contrastive Learning Rivals Masked Image Modeling in Fine-tuning via Feature Distillation (arxiv.org)


2. 特征蒸馏改变attention模式

图3为特征蒸馏前后CLIP ViT-B模型每个layer的平均attention map。显然attention map有两种模式:对角线模式和列模式。其中对角线模式表示某些固定相对位置的图像path之间的关系,而列模式表示某些位置的图像patch对其他位置的影响。可以明显看到特征蒸馏后,attention map的对角线模式更加突出,这说明特征蒸馏之后的模型更加依赖于相对位置关系图像path的视觉建模,从而表明蒸馏后的模型具有更好的平移不变性。

图3:特征蒸馏前后CLIP ViT-B模型每个layer的平均attention map
来源:[2205.14141] Contrastive Learning Rivals Masked Image Modeling in Fine-tuning via Feature Distillation (arxiv.org)


3. 特征蒸馏得到更好的loss/accuracy landscapes

作者使用参考文献[3]中的方法对不同模型的loss/accuracy landscapes进行可视化。在该方法中,模型的训练好的权重会被不同程度的高斯噪声扰动。为了考虑不同模型的权重幅值的影响,每个噪声的大小用滤波器的l2范数来定义。如图4所示,特征蒸馏后的loss/accuracy landscapes较蒸馏前更为平坦,说明网络更容易优化。

图4:不同预训练模型的loss/accuracy landscapes(上面是蒸馏前,下面是蒸馏后)
来源:[2205.14141] Contrastive Learning Rivals Masked Image Modeling in Fine-tuning via Feature Distillation (arxiv.org)


作者也对MAE(MIM模型)的预训练模型进行了特征蒸馏,发现效果提升并不明显(83.6%->83.8%)。进行特征可视化(见图5)发现,MAE在蒸馏前的平均attention 距离就比较发散,loss/accuracy landscapes也比较平坦。这说明特征蒸馏带来的fine-tuning性能的提升与MIM模型在功能上有一定重叠,将对比学习方法得到的模型进行特征蒸馏可以得到与MIM模型相似的性质,从而提升fine-tuning的性能。

图5:MAE蒸馏前后的平均attention距离和loss/accurcy landscapes
来源:[2205.14141] Contrastive Learning Rivals Masked Image Modeling in Fine-tuning via Feature Distillation (arxiv.org)


四、实践应用


文中是对MIM模型和对比学习的模型进行了实现,我们在ASR(Automatic Speech Recognition)的前处理VAD(Voice activity Detection)的分类模型MarbleNet中尝试使用了上述方法,也得到了提升fine-tuning性能的效果。下面先简单介绍一下任务背景,以及相关的模型和数据集。

4.1 ASR简介


ASR是指自动语音识别技术,将人的语音转换为文本,在语音搜索、字幕生成、会议记录等场景中广泛应用。我们所研究的应用场景为字幕生成,即给定一段视频或者音频文件生成对应的字幕文件。图6为我们目前构建的ASR的处理流程:输入一段视频或者音频文件,如果是视频文件的话需要将其中的音频剥离出来,再对音频文件进行语音增强(去噪、人声增强等),然后进行长语音切分(VAD),将较长的音频文件切分成一小段一小段包含人声的音频,接着将小段的音频送入ASR模型中转化成文字,最后用后处理纠错模块对文字的语法语义等进行矫正输出字幕文件。

图6:ASR的pipeline


4.2 VAD方案


如果把一整段长音频送入到ASR模型中处理速度会很慢,而且模型关注到的信息是有限的,输入太长不仅会导致有效信息被忽略,还会被很多无效信息干扰影响精度,因此需要对音频数据进行切分,把其中有效的人声送入ASR模型中,去除中间非人声的部分。VAD就是负责上述工作,是ASR中较为关键的前处理模块,其主要任务是从带有噪声的语音中准确定位出话音的开始点和结束点。图7为VAD的示意图,其中蓝色的曲线为声波,红色的线取值为0时表示背景音,为1表示人声,VAD就需要将音频按照图中红色线段的方式切割出来。

图7:VAD示意图


我们使用的VAD方案是参考NVIDIA Nemo中的方法, 先将长音频按一定的规则分割成短音频(以滑动窗口采样,窗口长度为630毫秒,步长为80毫秒),然后再用MarbleNet模型对短音频进行分类,类别为人声和非人声,最后再将分类结果进行后处理得到类似图7中的背景音或人声片段。因此VAD主要分成两个模块,分类模型和后处理。我们使用的模型为NVIDIA Nemo提出的MarbleNet-3x2x64,模型结构如图8所示,模型的输入是滑动窗口采样得到的小段音频对应的80维的Mel谱特征,用Cross Entropy Loss作为损失函数。由于分类模型的准确率不能达到100%,而且送入分类的音频片段之间存在重叠,因此需要引入后处理模块来做最终的切割。后处理模块的大致逻辑如下:设正样本阈值和负样本阈值, 分类概率大于正样本阈值的为正样本(人声),小于负样本阈值为负样本(非人声),位于正样本阈值和负样本阈值之间的样本以前一时刻的样本类别为准则(即前一时刻样本的类别即为当前样本的类别)。最后对按时间排列的样本做类似聚类的操作,如果连续一段时间内的样本大部分是正样本只有很少一部分为负样本,则将该时间段的音频取出来当作ASR的输入;如果连续一段时间内的样本大部分是负样本只有很少一部分为正样本,则该时间段的音频不需要经过ASR。

图8:MarbleNet
来源:[2010.13886] MarbleNet: Deep 1D Time-Channel Separable Convolutional Neural Network for Voice Activity Detection (arxiv.org)


4.3 实验结果


我们在VAD的分类模型MarbleNet上进行了论文中的蒸馏实验。由于MarbleNet并不是transformer结构的,因此实践过程中并没有使用共享相对位置编码的策略。教师模型参数是Nvidia Nemo提供的vad_multilingual_marblenet,模型蒸馏和fine-tuning使用的数据集是笔者基于公开数据集和网上爬取的数据构造的,训练集为4801259个样本,验证集为600868个样本,测试集为92500个样本。由于是二分类任务,直接用分类准确率无法评判模型的好坏,故用auc作为评价指标。

实验结果如表3所示,Vad_multilingual_marblenet_finetune和FD-Vad_multilingual_marblenet_finetune分别表示蒸馏前和蒸馏后finetune的结果,AUC为在测试集上的结果,最后一列是ASR使用对应的VAD模型后在71集电视剧上的结果(CER越小越好)。从表中可知,对MarbleNet模型引入特征蒸馏后其fine-tuning性能提升约0.5个点,说明论文中提出的特征蒸馏的方法不仅在对比学习的预训练模型上适用,对普通的规模较大的数据训练的模型同样适用,具有较好的普适性。VAD模型蒸馏后在ASR上的效果提升并不是很明显,这是由于VAD分类模型和ASR的结果并不一定是正相关的,只有VAD模型将某些关键时间点上的样本分类正确才可以明显提升ASR的性能。

表3:MarbleNet特征蒸馏结果


五、小结


本文介绍了一种简单的特征蒸馏方法,通过特征蒸馏后处理可以普遍提高许多视觉预训练模型的fine-tuning性能,使得对比学习方法得到的预训练模型与MIM在fine-tuning方面一样具有竞争力。笔者在语音VAD任务上也进行了尝试,特征蒸馏之后模型性能提升约0.5个点,说明该方法并不只局限于对比学习方法得到的预训练模型,在很多场景中亦是适用的。

这篇论文启示我们:简单的策略也能做到显著的效果提升,而复杂的策略可能会带来一些难以预知的额外影响。在平时解决问题的过程中,寻找影响效果的本质因素是一个重要的环节,只有对问题本质有了解,才能更好地把握前进的方向。

参考资料

[1] Wei Y, Hu H, Xie Z, et al. Contrastive Learning Rivals Masked Image Modeling in Fine-tuning via Feature Distillation[J]. arXiv preprint arXiv:2205.14141, 2022.
[2] Jia F, Majumdar S, Ginsburg B. Marblenet: Deep 1d time-channel separable convolutional neural network for voice activity detection[C]//ICASSP 2021-2021 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP). IEEE, 2021: 6818-6822.
[3] Li H, Xu Z, Taylor G, et al. Visualizing the loss landscape of neural nets[J]. Advances in neural information processing systems, 2018, 31.
[4] He K, Chen X, Xie S, et al. Masked autoencoders are scalable vision learners[C]//Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2022: 16000-16009.
[5] Caron M, Touvron H, Misra I, et al. Emerging properties in self-supervised vision transformers[C]//Proceedings of the IEEE/CVF International Conference on Computer Vision. 2021: 9650-9660.
[6]  Li C, Yang J, Zhang P, et al. Efficient self-supervised vision transformers for representation learning[J]. arXiv preprint arXiv:2106.09785, 2021.
[7] Touvron H, Cord M, Douze M, et al. Training data-efficient image transformers & distillation through attention[C]//International Conference on Machine Learning. PMLR, 2021: 10347-10357.
[8] Radford A, Kim J W, Hallacy C, et al. Learning transferable visual models from natural language supervision[C]//International Conference on Machine Learning. PMLR, 2021: 8748-8763.
[9] Liu Z, Hu H, Lin Y, et al. Swin transformer v2: Scaling up capacity and resolution[C]//Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2022: 12009-12019.

  • China Business CooperationBusiness@metax-tech.com
  • International Business CooperationInternational.Business@metax-tech.com
  • Media CooperationPR@metax-tech.com