仰止

V1

2022/12/08阅读:26主题:默认主题

TinyBERT-模型蒸馏

本文介绍下TinyBERT,华为在2020发布的一篇论文,主要内容是对模型进行蒸馏,蒸馏的方法值得学习

论文信息

论文地址:

https://arxiv.org/abs/1909.10351

代码地址:

https://github.com/huawei-noah/Pretrained-Language-Model/tree/master/TinyBERT

主要内容

目前已经有很多模型压缩的技术,如矩阵分解,量化、权重共享、剪枝、以及知识蒸馏,本文的重点在于知识蒸馏。 如上图所示,tinybert的蒸馏步骤可以概括为General DistillationTask-specific Distillation,也就是在大规模语料上对通识知识的蒸馏,这是在预训练阶段的蒸馏,和在指定任务数据上对特定任务的知识蒸馏,并且用通识知识的蒸馏模型对指定任务的蒸馏模型进行初始化,这是在微调阶段的蒸馏。同时,在特征任务上进行知识蒸馏时,会先对数据进行增强。

作者的实验结果是,4层的tinybert可以达到bertbase的96.8%的效果,但是参数量为bertbase的13.3%,推理时间为10.6%,并且比其他蒸馏的效果要好,同时,6层的tinybert和bertbase的表现近似。

1、transformer distillation

对transformer网络层数的蒸馏。假设学生模型有M层,老师模型有N层,自定义一个map函数 ,实现学生层到老师层的map,表示学生模型的第m层从老师模型的第g(m)层学得信息。损失函数如下:

表示的是某一个transformer layer或者是embedding layer的损失函数, 表示第m层的目标函数值, 表示第m层的重要性,为超参数。

transformer distillation包括attention distill、hidden distill、embedding distill、以及prediction distill,如下图所示:

  • attention

其中,attention distill的目标函数为:

h表示注意力头的个数, 表示学生或老师第i个注意力头的attention matrix

同时,作者表明,之所以使用 ,而不是 作为拟合目标,是因为前者的收敛更快,效果更好。

  • hidden

其中,transformer输出蒸馏的损失函数为:

其中 , 表示学生模型的向量维度。 是一个可学习矩阵,用来对学生模型进行线性变化,将其转化为与老师模型相同的维度。

  • embedding

embedding层输出蒸馏的损失函数为:

可以看到基本与transformer输出的蒸馏形式是一样的。

  • prediction 损失函数为:

z表示logits,t表示温度系数,作者实验发现,t=1时效果最好。这部分的损失函数就和distillbert设计的蒸馏损失比较像.

整体模型的损失函数如下:

其中,m表示学生的层数

2、task-specific distillation

该部分先对数据集进行增强,然后进行蒸馏。作者对数据增强的解释为,学生模型在经过增强的数据集上进行训练,可以提高其效果,也就是说,相比于老师模型,学生模型在特定任务上的训练数据是经过增强的,以此来提升学生模型的效果,因此学生就有超过老师的可能。

数据增强 其伪代码如下:

作者结合bert和glove的词嵌入,在word-level上进行替换,以实现数据增强。作者的参数设置如下,

论文并没有对task-specific distillation的蒸馏部分进行阐述,说明其与general distill的蒸馏方式应该是一样的,只是一个处于预训练阶段,一个处于微调阶段。

3、实验结果

实验时,作者使用 进行映射,也就是说4层的tinybert的每层都是从3层的bertbase中学得。

下面是作者对tinybert使用得学习策略和蒸馏方式做的消融实验:

下面是作者针对学生层到老师层的映射做的消融实验:

可以看到,使用均匀映射的效果是最好的,同时,作者也表明,对于一个下游任务,自适应的选择层数是一个具有挑战性的问题,也是未来的工作方向。

相关思考

该论文通过蒸馏方式实现对模型的压缩,整体上的实现分为以下几步:

  • 预训练阶段的蒸馏
  • 数据增强
  • 微调阶段的蒸馏

每个蒸馏,又会进行以下操作:

  • embedding层的蒸馏
  • hidden层的蒸馏
  • attention的蒸馏
  • 预测层的蒸馏

之所以做数据增强,是为了在对具体任务蒸馏时,扩充学生模型的训练集,提高学生模型的表现。

分类:

人工智能

标签:

自然语言处理

作者介绍

仰止
V1