基于多层级知识蒸馏预训练语言模型自动压缩方法及平台

摘要

本发明公开了一种基于多层级知识蒸馏的预训练语言模型自动压缩方法及平台，所述方法包括如下步骤：步骤一、构建多层级知识蒸馏，在自注意力单元、隐藏层状态、嵌入层三个不同层级上蒸馏大模型的知识结构；步骤二、训练元学习的知识蒸馏网络，生成多种预训练语言模型的通用压缩架构；步骤三、基于进化算法搜索最佳压缩结构。首先，研究基于元学习的知识蒸馏生成多种预训练语言模型的通用压缩架构；其次，在已训练好的元学习网络基础上，通过进化算法搜索最佳压缩结构，由此得到与任务无关的预训练语言模型的最优通用压缩架构。

说明书

技术领域

本发明属于语言模型压缩领域，尤其涉及一种基于多层级知识蒸馏的预训练语言模型自动压缩方法及平台。

背景技术

大规模预训练语言模型在自然语言理解和生成任务上都取得了优异的性能，然而，将具有海量参数的预训练语言模型部署到内存有限的设备中仍然面临巨大挑战。在模型压缩领域，已有的语言模型压缩方法都是针对特定任务的语言模型压缩。面向下游其它任务时，使用特定任务知识蒸馏生成的预训练模型仍需要重新微调大模型以及生成相关的大模型知识。大模型微调费时费力，计算成本也很高。为了提高压缩模型面向多种下游任务使用过程中的灵活性和有效性，研究与任务无关的预训练语言模型的通用压缩架构。而且，已有的知识蒸馏方法主要是人工设计的知识蒸馏策略。由于受计算资源等限制，人工设计所有可能的蒸馏结构并且寻找最优结构几乎不可能。受神经网络架构搜索的启发，尤其是在少样本的情况下，本发明基于多层级知识蒸馏生成面向多任务的预训练语言模型的通用压缩架构。

发明内容

本发明的目的在于针对现有技术的不足，提供一种基于多层级知识蒸馏的预训练语言模型自动压缩方法及平台。本发明首先构建一种多层级的知识蒸馏，在不同层级上蒸馏大模型的知识结构。而且，引入元学习，生成多种预训练语言模型的通用压缩架构。具体地，设计一种结构生成器的元网络，基于多层级知识蒸馏方法构建知识蒸馏编码向量，利用该结构生成器生成与当前输入的编码向量对应的蒸馏结构模型。同时，提出伯努利分布采样的方法训练结构生成器。每轮迭代时，利用伯努利分布采样各个编码器迁移的自注意力单元，组成对应的编码向量。通过改变输入结构生成器的编码向量和小批次的训练数据，联合训练结构生成器和对应的蒸馏结构，可以学得能够为不同蒸馏结构生成权重的结构生成器。同时，在已训练好的元学习网络基础上，通过进化算法搜索最佳压缩结构，由此得到与任务无关的预训练语言模型的最优通用压缩架构。

本发明的目的是通过以下技术方案实现的：

一种基于多层级知识蒸馏的预训练语言模型自动压缩方法，包括如下步骤：

步骤一、构建多层级知识蒸馏，在自注意力单元、隐藏层状态、嵌入层三个不同层级上蒸馏大模型的知识结构；

步骤二、训练元学习的知识蒸馏网络，生成多种预训练语言模型的通用压缩架构；

步骤三、基于进化算法搜索最优压缩结构。

进一步的，步骤二中设计一种结构生成器的元网络，基于步骤一的多层级知识蒸馏构建知识蒸馏编码向量，利用结构生成器生成与当前输入的编码向量对应的蒸馏结构模型；同时，采用伯努利分布采样的方法训练结构生成器，每轮迭代时，利用伯努利分布采样各个编码器迁移的自注意力单元，组成对应的编码向量；通过改变输入结构生成器的编码向量和小批次的训练数据，联合训练结构生成器和对应的蒸馏结构，得到为不同蒸馏结构生成权重的结构生成器。

进一步的，步骤三中在已训练好的元学习网络基础上，通过进化算法搜索最优压缩架构，得到与任务无关的预训练语言模型的最优通用压缩架构。

进一步的，步骤一中将自注意力分布知识、隐藏状态知识和嵌入层知识编码为一个蒸馏网络，采用知识蒸馏实现大模型向小模型的压缩。

进一步的，步骤一中包括自注意力知识蒸馏、隐藏层状态知识蒸馏和嵌入层知识蒸馏。

进一步的，步骤二中所述结构生成器的元网络，由两个全连接层组成，输入一个自注意力知识蒸馏编码向量，输出结构生成器的权重矩阵；

结构生成器的训练过程如下：

步骤1：构造知识蒸馏编码向量，包括层采样向量、多头剪枝向量、隐藏层降维向量和嵌入层降维向量；

步骤2： 基于结构生成器构建蒸馏网络架构，利用该结构生成器构建与当前输入的编码向量对应的蒸馏结构模型，调整结构生成器输出的权重矩阵的形状，与自注意力编码向量对应的蒸馏结构的输入输出的自注意力单元数目一致；

步骤3：联合训练结构生成器和蒸馏结构模型：通过伯努利分布采样的方法训练结构生成器，通过改变输入结构生成器的自注意力编码向量和一个小批次的训练数据，联合训练结构生成器和对应的蒸馏结构，可以学得能够为不同蒸馏结构生成权重的结构生成器。

进一步的，步骤三中，将网络编码向量输入训练好的结构生成器，生成对应蒸馏网络的权重，在验证集上对蒸馏网络进行评估，获得对应蒸馏网络的精度；具体如下：

满足特定约束条件下，首先选取一系列蒸馏网络编码向量作为蒸馏网络的基因，通过在验证集上评估获得对应蒸馏网络的精度；然后，选取精度最高的前k个基因，采用基因重组和变异生成新的基因,通过进一步重复前k个最优基因选择的过程和新基因生成的过程来迭代获得满足约束条件并且精度最高的基因。

进一步的，所述进化算法的具体步骤如下：

步骤一、将知识蒸馏编码向量定义为蒸馏结构模型的基因G，随机选取满足约束条件C的一系列基因作为初始种群；

步骤二、评估现有种群中各个基因

步骤三、利用步骤二选取的精度最高的前k个基因进行基因重组和基因变异生成新的基因，将新基因加入现有种群中；

步骤四、重复迭代N轮步骤二和步骤三，选择现有种群中前k个精度最高的基因并生成新基因，直到获得满足约束条件C并且精度最高的基因。

一种基于多层级知识蒸馏的预训练语言模型自动压缩平台，包括如下组件：

数据加载组件：用于获取登陆用户上传的待压缩的包含具体自然语言处理下游任务的BERT模型和面向多任务的预训练语言模型的训练样本，所述训练样本是满足监督学习任务的带标签的文本样本；

自动压缩组件：用于将面向多任务的预训练语言模型自动压缩，包括知识蒸馏向量编码模块、蒸馏网络生成模块、结构生成器和蒸馏网络联合训练模块、蒸馏网络搜索模块、特定任务微调模块；

所述知识蒸馏向量编码模块包括 Transformer 的层采样向量、自注意力单元的多头剪枝向量、隐藏层降维向量、嵌入层降维向量，前向传播过程中，将蒸馏网络编码向量输入结构生成器，生成对应结构的蒸馏网络和结构生成器的权重矩阵；

所述蒸馏网络生成模块基于结构生成器构建与当前输入的编码向量对应的蒸馏网络，调整结构生成器输出的权重矩阵的形状，使其与自注意力编码向量对应的蒸馏结构的输入输出的自注意力单元数目一致；

结构生成器和蒸馏网络联合训练模块端到端地训练结构生成器，将多层级知识蒸馏编码向量和小批次的训练数据输入蒸馏网络，更新蒸馏结构的权重和结构生成器的权重矩阵；

所述蒸馏网络搜索模块采用进化算法搜索满足特定约束条件的最高精度的蒸馏网络；

所述特定任务微调模块是在所述自动压缩组件生成的预训练模型蒸馏网络上构建下游任务网络，利用蒸馏网络的特征层和输出层对下游任务场景进行微调，输出最终微调好的学生模型，即登陆用户需求的包含下游任务的预训练语言模型压缩模型，将所述压缩模型输出到指定的容器，供所述登陆用户下载，并在所述平台的输出压缩模型的页面呈现压缩前后模型大小的对比信息；

推理组件：登陆用户从所述平台获取预训练压缩模型，利用所述自动压缩组件输出的压缩模型在实际场景的数据集上对登陆用户上传的自然语言处理下游任务的新数据进行推理。并在所述平台的压缩模型推理页面呈现压缩前后推理速度的对比信息。

进一步的，登陆用户可直接下载训练好的预训练语言模型，根据用户对具体某个自然语言处理下游任务的需求，在所述平台生成的已压缩的预训练模型架构基础上构建下游任务网络并进行微调，最后部署在终端设备，或直接在所述平台上对自然语言处理下游任务进行推理。

本发明的有益效果是：本发明的基于多层级知识蒸馏的预训练语言模型自动压缩方法及平台法及平台，首先，研究基于元学习的知识蒸馏生成多种预训练语言模型的通用压缩架构；其次，在已训练好的元学习网络基础上，通过进化算法搜索最佳压缩结构，由此得到与任务无关的预训练语言模型的最优通用压缩架构。

本发明的基于多层级知识蒸馏的预训练语言模型自动压缩平台，压缩生成面向多任务的预训练语言模型的通用架构，充分利用已压缩好的模型架构提高下游任务的压缩效率，并且可将大规模自然语言处理模型部署在内存小、资源受限等端侧设备，推动了通用深度语言模型在工业界的落地进程。

附图说明

图1是基于多层级知识蒸馏的预训练语言模型自动压缩平台的整体架构图。

图2是多层级知识蒸馏示意图。

图3是结构生成器的训练流程图。

图4是编码器模块的隐藏层和输入的嵌入层的降维结构示意图。

图5是基于结构生成器构建蒸馏网络的架构图。

图6是结构生成器和蒸馏网络联合训练过程流程图。

图7是基于进化算法的蒸馏网络搜索架构示意图。

具体实施方式

下面结合附图对本发明作进一步说明。

如图1所示，本发明包括基于元学习的知识蒸馏以及基于进化算法的蒸馏网络自动搜索。将面向多任务的大规模预训练语言模型自动压缩生成满足不同硬约束条件（如，浮点数运算次数），与任务无关的通用架构。

具体方案如下：本发明的基于多层级知识蒸馏的预训练语言模型自动压缩方法整个过程分为三个阶段，第一个阶段是构建多层级的知识蒸馏，在自注意力单元、隐藏层状态、嵌入层三个不同层级上蒸馏大模型的知识结构；第二个阶段是训练元学习的知识蒸馏网络，生成多种预训练语言模型的通用压缩架构。具体地，设计一种结构生成器的元网络，基于第一阶段提出的多层级知识蒸馏方法构建知识蒸馏编码向量，利用结构生成器生成与当前输入的编码向量对应的蒸馏结构模型。同时，提出伯努利分布采样的方法训练结构生成器。每轮迭代时，利用伯努利分布采样各个编码器迁移的自注意力单元，组成对应的编码向量。通过改变输入结构生成器的编码向量和小批次的训练数据，联合训练结构生成器和对应的蒸馏结构，可以学得能够为不同蒸馏结构生成权重的结构生成器；第三个阶段是基于进化算法搜索最佳压缩结构，在已训练好的元学习网络基础上，通过进化算法搜索最优压缩架构，由此得到与任务无关的预训练语言模型的最优通用压缩架构。具体过程如下：

第一阶段：多层级知识蒸馏

本发明将自注意力分布知识、隐藏状态知识和嵌入层知识编码为一个蒸馏网络，如图2

所示。采用知识蒸馏实现大模型向小模型的压缩，并最大程度地将大模型的自注意力知识迁移给小模型。

注意力函数是由 queries、keys 和 values 三个矩阵计算所得，分别表示为矩阵Q、矩阵 K 和矩阵 V。自注意力函数定义如下：

其中，

其中，h是注意力头的数量，

学生模型学习模仿教师网络的多头注意力知识，其中损失函数定义如下：

其中，

识蒸馏，即对 Transformer 层输出的知识进行迁移。隐藏层状态知识蒸馏的损失函数如下：

其中，矩阵

其中，矩阵

第二阶段：元学习的知识蒸馏

设计一个结构生成器，结构生成器是一个元网络，由两个全连接层组成。图3显示了结构生成器的训练过程。输入一个自注意力知识蒸馏编码向量，输出结构生成器的权重矩阵。

结构生成器的训练过程如下：

步骤一：构造知识蒸馏编码向量。前向传播过程中，将蒸馏网络编码向量输入结构生成器，输出结构生成器的权重矩阵。蒸馏网络编码向量由 Transformer 层采样向量、多头注意力剪枝向量、隐藏层降维向量和嵌入层降维向量组成。蒸馏网络编码向量的具体方案如下：

利用随机变量X依次对BERT的12个 Transformer 单元进行伯努利采样，生成一个由12 个0或1元素组成的向量。当随机变量

利用以上伯努利采样方式对BERT包含的所有Transformer层依次进行层采样，组成网络编码向量中的层采样向量：

值得注意的是，

此时，对Transformer进行知识蒸馏时，通过网络编码向量中的层采样向量，将教师模型和学生模型建立一对一的映射关系，之后根据网络编码向量中的层采样向量生成对应的蒸馏网络结构。为了加速整个训练过程，使用教师模型中层采样向量对应元素为1的Transformer 的权重来初始化学生模型迁移的 Transformer 模块。

其中，

因此，基于注意力多头剪枝进行知识蒸馏，每次生成一个蒸馏网络结构时，对于进行迁移的 Transformer 模块，即Transformer层采样编码向量中值为1的元素，生成一个多头剪枝编码向量，表示当前迁移的所有Transformer层中进行自注意力知识迁移的注意力头的数量。

其中，

因此，基于隐藏层进行知识蒸馏生成蒸馏网络时，每次采用简单随机抽样从上述列表采样一个降维因子，生成当前蒸馏网络对应的隐藏层的维度大小。

因此，基于嵌入层进行知识蒸馏生成蒸馏网络时，每次嵌入层的维度与当前隐藏层的维度的大小保持一致即可。

步骤二： 基于结构生成器构建蒸馏网络架构。利用该结构生成器构建与当前输入的编码向量对应的蒸馏结构模型，而且需要调整结构生成器输出的权重矩阵的形状，使其与自注意力编码向量对应的蒸馏结构的输入输出的自注意力单元数目一致。图5所示由结构生成器构建的蒸馏结构的网络架构。

步骤三：联合训练结构生成器和蒸馏结构模型。将自注意力知识蒸馏编码向量和一个小批次的训练数据输入蒸馏结构模型。值得注意的是，反向传播的过程中，蒸馏结构的权重和结构生成器的权重矩阵都会一起更新。构生成器的权重可以使用链式法则计算，因此，可以端到端的训练结构生成器。

同时，提出伯努利分布采样的方法训练结构生成器。每轮迭代时，利用伯努利分布采样各层 Transformer 迁移的自注意力单元，组成对应的编码向量。通过改变输入结构生成器的自注意力编码向量和一个小批次的训练数据，联合训练结构生成器和对应的蒸馏结构，可以学得能够为不同蒸馏结构生成权重的结构生成器。图6所示结构生成器的训练过程。

第三阶段：基于进化算法的蒸馏网络搜索

将网络编码向量输入训练好的结构生成器，生成对应蒸馏网络的权重，在验证集上对蒸馏网络进行评估，获得对应蒸馏网络的精度。然后，为了搜索出满足特定约束条件的最高精度的蒸馏网络，提出采用进化算法搜索满足特定约束条件的最高精度的蒸馏网络，图7所示了基于进化算法的蒸馏网络搜索架构。

在元学习蒸馏网络中采用的进化搜索算法中，每个蒸馏网络是由包含嵌入层蒸馏、隐藏层蒸馏和自注意力知识蒸馏三个蒸馏模块的网络编码向量编码生成，所以将蒸馏网络编码向量定义为蒸馏网络的基因。在满足特定约束条件下，首先选取一系列蒸馏网络编码向量作为蒸馏网络的基因，通过在验证集上评估获得对应蒸馏网络的精度。然后，选取精度最高的前k个基因，采用基因重组和变异生成新的基因。基因变异是指，通过随机改变基因里一部分元素值来进行变异。基因重组是指，随机地将两个父辈的基因进行重组，产生后代。而且，可以很容易地通过消除不合格的基因来加强约束。通过进一步重复前k个最优基因选择的过程和新基因生成的过程来迭代几轮，就可以获得满足约束条件并且精度最高的基因。第三阶段：如图7所示为基于进化算法的蒸馏网络搜索的过程：

在第二阶段训练好的元学习的知识蒸馏网络基础上，将多个满足特定约束条件的知识蒸馏编码向量输入结构生成器生成对应的权重矩阵，得到多个蒸馏结构模型；在验证集上对每个蒸馏结构模型进行评估，获得对应的精度；采用进化算法搜索其中满足特定约束条件（如浮点数运算次数）的精度最高的蒸馏结构模型，由此得到与任务无关的预训练语言模型的通用压缩架构。进化搜索算法的具体步骤如下：

步骤一、每个蒸馏结构模型是由基于Transformer层采样的知识蒸馏编码向量生成的，所以将知识蒸馏编码向量定义为蒸馏结构模型的基因G，随机选取满足约束条件C的一系列基因作为初始种群。

步骤二、评估现有种群中各个基因

步骤三、利用步骤二选取的精度最高的前k个基因进行基因重组和基因变异生成新的基因，将新基因加入现有种群中。基因变异是指通过随机改变基因里一部分元素值来进行变异；基因重组是指随机地将两个父辈的基因进行重组产生后代；而且可以很容易地通过消除不合格的基因来加强约束C。

步骤四、重复迭代N轮步骤二和步骤三，选择现有种群中前k个精度最高的基因并生成新基因，直到获得满足约束条件C并且精度最高的基因。

本发明的基于多层级知识蒸馏的预训练语言模型自动压缩平台，包括以下组件：

数据加载组件：用于获取登陆用户上传的待压缩的包含具体自然语言处理下游任务的BERT模型和面向多任务的预训练语言模型的训练样本，所述训练样本是满足监督学习任务的带标签的文本样本。

自动压缩组件：用于将面向多任务的预训练语言模型自动压缩，包括知识蒸馏向量编码模块、蒸馏网络生成模块、结构生成器和蒸馏网络联合训练模块、蒸馏网络搜索模块、特定任务微调模块。

知识蒸馏向量编码模块包括 Transformer 的层采样向量、自注意力单元的多头剪枝向量、隐藏层降维向量、嵌入层降维向量。前向传播过程中，将蒸馏网络编码向量输入结构生成器，生成对应结构的蒸馏网络和结构生成器的权重矩阵。

蒸馏网络生成模块是基于结构生成器构建与当前输入的编码向量对应的蒸馏网络，调整结构生成器输出的权重矩阵的形状，使其与自注意力编码向量对应的蒸馏结构的输入输出的自注意力单元数目一致。

结构生成器和蒸馏网络联合训练模块是端到端的训练结构生成器，具体地，将多层级知识蒸馏编码向量和一个小批次的训练数据输入蒸馏网络。更新蒸馏结构的权重和结构生成器的权重矩阵。

蒸馏网络搜索模块是为了搜索出满足特定约束条件的最高精度的蒸馏网络，提出进化算法搜索满足特定约束条件的最高精度的蒸馏网络。将网络编码向量输入训练好的结构生成器，生成对应蒸馏网络的权重，在验证集上对蒸馏网络进行评估，获得对应蒸馏网络的精度。在元学习蒸馏网络中采用的进化搜索算法中，每个蒸馏网络是由包含嵌入层知识蒸馏、隐藏层状态的知识蒸馏和自注意力知识蒸馏三个蒸馏模块的网络编码向量编码生成，所以将蒸馏网络编码向量定义为蒸馏网络的基因。在满足特定约束条件下，首先选取一系列蒸馏网络编码向量作为蒸馏网络的基因，通过在验证集上评估获得对应蒸馏网络的精度。然后，选取精度最高的前k个基因，采用基因重组和变异生成新的基因。通过进一步重复前k个最优基因选择的过程和新基因生成的过程进行迭代，获得满足约束条件并且精度最高的基因。

特定任务微调模块是在所述自动压缩组件生成的预训练模型蒸馏网络上构建下游任务网络，利用蒸馏网络的特征层和输出层对下游任务场景进行微调，输出最终微调好的学生模型，即登陆用户需求的包含下游任务的预训练语言模型压缩模型。将所述压缩模型输出到指定的容器，可供所述登陆用户下载，并在所述平台的输出压缩模型的页面呈现压缩前后模型大小的对比信息。

推理组件：登陆用户从所述平台获取预训练压缩模型，用户利用所述自动压缩组件输出的压缩模型在实际场景的数据集上对登陆用户上传的自然语言处理下游任务的新数据进行推理。并在所述平台的压缩模型推理页面呈现压缩前后推理速度的对比信息。

登陆用户可直接下载本发明平台提供的训练好的预训练语言模型，根据用户对具体某个自然语言处理下游任务的需求，在所述平台生成的已压缩的预训练模型架构基础上构建下游任务网络并进行微调，最后部署在终端设备。也可以直接在所述平台上对自然语言处理下游任务进行推理。

下面将以电影评论进行情感分类任务对本发明的技术方案做进一步的详细描述。

通过所述平台的数据加载组件获取登陆用户上传的单个句子的文本分类任务的BERT模型和情感分析数据集SST-2；

通过所述平台的自动压缩组件，生成面向多任务的 BERT 预训练语言模型；

通过所述平台加载自动压缩组件生成的BERT预训练模型，在所述生成的预训练模型上构建文本分类任务的模型；

基于所述自动压缩组件的特定任务微调模块所得的学生模型进行微调，利用自动压缩组件生成的BERT预训练模型的特征层和输出层对下游文本分类任务场景进行微调，最终，平台输出登陆用户需求的包含文本分类任务的BERT模型的压缩模型。

将所述压缩模型输出到指定的容器，可供所述登陆用户下载，并在所述平台的输出压缩模型的页面呈现压缩前后模型大小的对比信息，压缩前模型大小为110M，压缩后为53M，压缩了51.8%。如下表1所示。

表1：文本分类任务BERT模型压缩前后对比信息

通过所述平台的推理组件，利用所述平台输出的压缩模型对登陆用户上传的SST-2测试集数据进行推理，并在所述平台的压缩模型推理页面呈现压缩后比压缩前推理速度加快1.95倍，并且推理精度从压缩前的91.5%提升为92.3%。