基于不确定性估计知识蒸馏的语言模型压缩方法

技术领域

本发明属于神经网络压缩领域，特别涉及一种预训练语言模型的压缩方法，可用于模型部署、模型压缩、减轻模型硬件负担。

背景技术

近些年来，自然语言处理研究社区见证了一场预训练和自我监督模式的革命，随着大规模预训练语言模型的研究和应用，首先在大规模文本数据上进行预训练，然后在下游任务进行迁移学习，预训练加微调逐渐成为自然语言处理解决方案的基本范式。如BERT的出现，显著提高了许多自然语言处理任务的性能。然而，预训练语言模型通常具有昂贵的计算代价和密集的内存消耗，这些模型通常有数亿个参数，在边缘和硬件受限的计算训练或推理预算下操作这些大型模型仍然具有挑战性。虽然在设备上实时操作这些模型有可能实现新颖而有趣的语言处理应用程序，但这些模型日益增长的计算和内存需求可能会阻碍其广泛采用。此外，相关研究也证明了预训练模型存在冗余。因此，在保持预训练语言模型性能的同时降低计算开销和模型存储是至关重要，模型轻量化具有巨大的应用和研究价值。

近年来，网络轻量化技术进行了大量的研究，提出了许多不同角度的网络轻量化技术。目前，常见的网络轻量化方向有模型剪枝、知识蒸馏、数值量化等等。知识蒸馏的思想由Hinton等人在论文Distilling the knowledge in a neural network中提出，并得到广泛的研究和发展。其中由于BERT网络优异的性能，使得有关BERT的压缩技术引起了广泛关注。

Jiao等人在论文TinyBERT:Distilling BERT for Natural LanguageUnderstanding中提出了一种为基于Transformer的模型专门设计的知识蒸馏的方案。通过增加三部分损失函数的约束，来试图获得更好的蒸馏效果。分别对网络中间嵌入层的输出、Transformer层的Attention Map权重进行均方差损失函数计算，对预测结果进行交叉熵约束。最终压缩后的模型大小为原始模型的1/7，而且网络推理速度提高了9倍，但是网络压缩训练过程计算代价非常大，压缩过程花费了750GPU小时，导致大部分研究人员难以承受该计算代价。

Xu等人在论文BERT-of-Theseus:Compressing BERT by Progressive ModuleReplacing中提出了利用模块替换思想来做模型压缩，将原始模型中中具备相似特性的网络子结构进行压缩替换，同时在训练过程中将原始网络模块进行冻结，只训练替换的网络结构。该方案在中文命名实体识别NER任务上进行验证，虽然该压缩方法的计算代价不大，但是压缩后的网络模型精确度与原始网络相比下降了1-2个百分点，并且网络压缩率较低，使得模型运行效率下降，导致大量计算资源被浪费。

上述已存在的网络轻量化方法的不足是：1)缺少对网络中间推理过程的监督，2)对原始网络参数的利用率不够，3)知识蒸馏过程缺少噪声估计。进而导致了网络压缩过程计算负担过大，压缩过程效率过低，压缩后网络的性能精度较低。

发明内容

本发明的目的在于针对上述现有技术的不足，提出了一种基于不确定性估计知识蒸馏的语言模型压缩方法，以减少计算资源消耗，缩短压缩运行时间，提高压缩效率和压缩后网络的性能精度。

为实现上述目的，本发明的技术方案包括如下：

1.基于不确定性估计知识蒸馏的语言模型压缩方法，其特征在于，包括：

(1)对原始语言模型W进行对半压缩，得到压缩后的神经网络模型

(2)对压缩后神经网络

(3)对压缩后的神经网络模型

(3a)定义参数蒸馏损失函数L

L＝L

其中，L

(3b)利用(3a)所定义的损失函数对压缩后的神经网络模型参数进行更新：

(3b1)构建由多层卷积神经网络与激活函数层级联组成的不确定性估计模块；

(3b2)将原始语言模型的偶数层前馈网络结构与压缩后神经网络模型进行对齐蒸馏，并利用参数蒸馏损失函数L

(3b3)利用原始语言模型对输入数据样本x

(3b4)利用自然语言处理的具体任务输入数据样本x

本发明与现有技术相比，具有以下优点：

1.本发明通过合理设计压缩后神经网络的结构，合理初始化压缩后网络的参数，保留了原始语言模型的大部分推理性能。与现有技术相比，提高了压缩后神经网络的初始性能，减少了训练过程中的资源浪费。

2.本发明由于将原始语言模型和压缩后神经网络的前馈网络结构进行对齐，并通过参数蒸馏损失函数L

3.本发明通过不确定性估计模块对训练过程中的噪声进行估计，学习噪声分布的均值μ和方差σ，从而建模了原始语言模型的特征分布，并通过不确定性损失函数L

附图说明

图1为本发明的实现流程图；

图2为本发明原始语言模型结构示意图；

图3为本发明设计的知识蒸馏方案的示意图；

图4为本发明提出的不确定性估计模块的示意图。

具体实施方式

以下结合附图和具体实施例，对本发明作进一步详细描述。

参照图1，基于不确定性估计知识蒸馏的预训练语言模型压缩方法，其实现步骤如下：

步骤1，获取训练和测试的数据集。

获取公开自然语言理解基本任务GLUE中的数据集，该数据集包含常见的自然语言处理的各种任务，能够较好地测试出语言模型的综合性能。

本实例是从该数据集中的以下四类数据集中获取，后续进行实验测试任务：

一是语言可接受性语料库CoLA是单句子分类任务，其语料来自语言理论书籍和期刊，其中的每个单词序列都被标注为是否合乎语法；

二是斯坦福情感树语料库SST-2是单句子分类任务，具体包含电影评论中的句子和句子情感极性的人类注释；

三是微软研究院释义语料库MRPC，其包括语义相似性和释义任务，这些任务是从在线新闻源中自动抽取句子对组成，并且人工注释了句子对是否在语义上等效；

四是成对问题语义判断数据集QQP，其包括句子相似性任务和释义任务，具体来自社区问答网站Quora中的问题集合，并且人工注释了两个问题在语义上是否等效。

步骤2，构建原始语言模型，如图2。

本步骤的具体实现如下：

(2.1)使用Google的HuggingFace团队开源的BERT

所述语言模型，其由12层堆叠的Transformer网络构成，该网络包括12个注意力层，12个前馈网络层和24个归一化层。其结构关系为：输入层→第一注意力层→第一归一化层→第一前馈网络层→第二归一化层→第二注意力层→第三归一化层→第二前馈网络层→第四归一化层→第三注意力层→第五归一化层→第三前馈网络层→第六归一化层→第四注意力层→第七归一化层→第四前馈网络层→第八归一化层→第五注意力层→第九归一化层→第五前馈网络层→第十归一化层→第六注意力层→第十一归一化层→第六前馈网络层→第十二归一化层→第七注意力层→第十三归一化层→第七前馈网络层→第十四归一化层→第八注意力层→第十五归一化层→第八前馈网络层→第十六归一化层→第九注意力层→第十七归一化层→第九前馈网络层→第十八归一化层→第十注意力层→第十九归一化层→第十前馈网络层→第二十归一化层→第十一注意力层→第二十一层归一化层→第十一前馈网络层→第二十二层归一化层→第十二注意力层→第二十三层归一化层→第十二前馈网络层→第二十四层归一化层→输出层组成。

(2.2)语言模型的输入数据由三种向量嵌入加和组成，分别是：词元嵌入、输入分隔嵌入，位置嵌入，其中分隔嵌入是标注两个输入语句是否为连续语句，若语句连续为1，若语句非连续则为0值；

步骤3，微调训练原始语言模型。

利用步骤1中所述自然语言数据集的输入数据样本xi，标签yi，通过交叉熵损失函数对语言模型的参数进行迭代更新，直至损失函数收敛，得到训练好的语言模型。

步骤4，对语言模型W进行对半压缩，得到压缩后神经网络

对语言模型W进行对半压缩，是舍弃原始语言模型的偶数层Transformer网络结构，保留原始语言模型的奇数层Transformer网络结构。压缩后神经网络

第一注意力层的注意力头为12，输出长度为512，隐藏向量维度为768；

第二注意力层的注意力头为12，输出长度为512，隐藏向量维度为768；

第三注意力层的注意力头为12，输出长度为512，隐藏向量维度为768；

第四注意力层的注意力头为12，输出长度为512，隐藏向量维度为768；

第五注意力层的注意力头为12，输出长度为512，隐藏向量维度为768；

第六注意力层的注意力头为12，输出长度为512，隐藏向量维度为768；

所有归一化层采用LayerNorm函数；

压缩后网络中包含的所有激活函数均采用ReLU函数。

复制步骤3中训练完成的语言模型中奇数层Transformer网络参数，作为压缩后神经网络的参数初始化。

步骤5，通过不确定性估计知识蒸馏方式对压缩后神经网络进行训练。

参照图3，所述不确定性估计知识蒸馏方法对多头注意力层进行参数复制和冻结，对原始语言模型和压缩后神经网络的前馈网络结构进行对齐蒸馏，同时对语言模型和压缩后神经网络预测层的软目标结果进行蒸馏约束，具体实现如下：

(5.1)将语言模型和压缩后神经网络的前馈网络结构进行蒸馏对齐。具体参数蒸馏函数L

其中，

(5.2)将原始语言模型和压缩后神经网络两个网络的输出结果通过交叉熵损失函数进行拟合，具体公式如下：

L

其中，z

(5.3)设计不确定性估计模块：

参照图4，将本步骤设计的不确定性估计模块置于原始语言模型和压缩后神经网络之间，利用该模块在预测层前估计特征的不确定性。

不确定性估计模块具体包括4个卷积神经网络层和4个激活函数层的不确定性估计模块，其结构如下：

输入层→第1卷积层→第1激活函数层→第2卷积层→第2激活函数层→第3卷积层

→第3激活函数层→第4卷积层→第4激活函数层→输出层。

其中不确定性估计模块中的各层参数如下：

第1卷积层卷积核大小为1*1，输入通道为1维，输出通道为512维；

第2卷积层卷积核大小为1*1，输入通道为512维，输出通道为1024维；

第3卷积层卷积核大小为1*1，输入通道为1024维，输出通道为512维；

第4卷积层卷积核大小为1*1，输入通道为512维，输出通道为1维；

不确定性估计模块包含的所有激活函数均采用ReLU函数；

(5.4)利用不确定性估计模块学习特征分布的方差σ，具体不确定性估计损失函数L

其中，N是一个批次中的样品数量，d[·,·]是衡量压缩后神经网络和原始语言模型之间差距的距离度量标准，w

在知识蒸馏过程中加入了对蒸馏噪声的建模，我们认为知识蒸馏过程中存在数据噪声，并且困难数据样本在知识蒸馏过程中提供的高价值信息较少，因此我们对每个样本的输出建模为一个高斯分布N(μ,σ

其中，∈表示均值为零方差为一的标准正态分布。同时估计类似原始语言模型特征

其中，

然后使用带Jeffrey’s prior的对数似然函数可以将上式表述如下形式：

在压缩后骨干网络的末端添加了一个新分支(Uncertainty Estimate Module,UEM)来估计特征不确定性。基于上述推导，通过最大似然估计(Maximun Likelihood,ML)制定一个不确定性估计蒸馏损失，以帮助压缩后神经网络感知样本难度：

其中，n表示小批量训练数据集的样本数。通过原始语言模型提取特征z

(5.5)在语言可接受性语料库、斯坦福情感树库、微软研究院释义语料库、问题对数集四个自然语言处理任务数据集上进行压缩后神经网络训练：利用自然语言处理的具体任务输入数据样本x

下面结合仿真实验对本发明的效果做进一步说明。

1.仿真条件：

仿真实验的硬件环境是：NVIDIA GTX 2080Ti型号的GPU，128GB的运行内存；

仿真实验的软件环境是：深度学习框架Pytorch1.8.0。

仿真实验中，对四个任务拥有不同的客观量化评价指标，语义可接受性测试CoLA采用马修斯相关系数MCC指标，可以表示为：

其中，n表示观测总数，χ表示原始数据。

语义情感分类任务SST-2采用准确率指标acc，为了克服样本数据不均衡问题，句子对语义相似性任务QQP和微软语义相似性任务MRPC均采用F1 Score指标,其表示为：

其中，precision表示准确率，recall表示召回率。

2.仿真内容与结果分析。

仿真1，使用近些年的相关语言模型压缩方法和本发明方案对原始语言模型进行压缩仿真，并将所有得到的压缩后神经网络在不同自然语言任务上进行模型语言能力测试，计算各自任务的评价指标，即CoLA任务的MCC指标，QQP和MRPC任务的F1Score指标，SST-2和QNLI任务的准确率指标，结果如表1所示：

表1不确定性估计知识蒸馏的压缩方法实验结果

表1中参与对比的压缩方法分别是：Bert

表1中进行测试的自然语言任务分别是：微软研究院释义语料库MRPC、语言可接受性语料库CoLA、斯坦福情感树语料库SST-2、问答自然语言推断QNLI、成对问题语义判断数据集QQP。各项任务的实验结果平均数作为衡量神经网络综合性能的指标Average。

从表1中可以看出，本发明与同类网络压缩方案相比，压缩后的神经网络在MRPC、CoLA、QNLI三项任务上取得了最优的性能指标，同时，压缩后神经网络的综合性能达到最优。

仿真2，对本发明和所有对比方法进行网络模型的参数量统计和网络推理速度测量，得到所有方法压缩后神经网络的参数和推理时间情况，如表2所示：

表2本发明与对比压缩方案的网络参数对比

从表2中可以看出，本发明方法与同类网络压缩方案相比，参与压缩训练的参数规模较小，网络推理时间较短。

综合表1和表2的结果可以看出，本发明与同类网络压缩方案相比，在压缩后神经网络综合性能最优的情况下，参与训练的网络参数量最小，压缩过程计算负担最小，网络压缩率较高，压缩后神经网络的运行速度较快。