[发明专利]一种模型蒸馏方法、文本检索方法及装置

说明书

本申请提供一种模型蒸馏方法、文本检索方法及装置。方法包括：获取老师模型和学生模型，将训练样本分别输入老师模型和学生模型中；利用EMD计算老师模型的第一transformer层的输出与学生模型的第二transformer层的输出的第一蒸馏损失，并分别计算老师模型的第一embedding层的输出与学生模型的第二embedding层的输出的第二蒸馏损失，以及老师模型的第一prediction层的输出与学生模型的第二prediction层的输出的第三蒸馏损失；根据第一蒸馏损失、第二蒸馏损失和第三蒸馏损失对学生模型中的参数进行优化获得训练好的学生模型。本申请能够保证学生模型从老师模型中学习到更多的知识。

技术领域

本申请涉及自然语言处理技术领域，具体而言，涉及一种模型蒸馏方法、文本检索方法及装置。

背景技术

随着深度学习的发展，自然语言处理中深度神经网络的使用越来越多，然而很多模型，如Bert，xlnet，都存在模型复杂，参数量大，训练时间长，内存消耗大，推理时间长的问题，很难直接应用于图形处理器(Graphics Processing Unit，GPU)及智能手机等应用资源受限的设备上。

为了解决这些问题，研究者们提出了许多模型压缩的方法，以减少受训神经网络的冗余，对于nlp未来的发展而言非常有价值。知识蒸馏就是一种对模型压缩的方法，该方法是通过老师模型中的某一层与学生模型的某一层进行对应，通过学习将学生模型的某一层对老师模型的对应层的距离尽量缩小，以实现学生模型对老师模型的学习。这种模型压缩方法学生模型能够学习到老师模型中的知识有限，因此，训练获得的学生模型在文本检索时准确性不高。

发明内容

本申请实施例的目的在于提供一种模型蒸馏方法、文本检索方法及装置，用以解决现有技术获得的学生模型在文本检索时准确性不高的问题。

第一方面，本申请实施例提供一种模型蒸馏方法，包括：获取老师模型和学生模型，其中，所述老师模型包括第一向量embedding层、第一转换器transformer层和第一预测prediction层，所述学生模型包括第二embedding层、第二transformer层和第二prediction层；所述第一transformer层的层数大于所述第二transformer层的层数；获取训练样本，并将所述训练样本分别输入所述老师模型和所述学生模型中；其中，所述老师模型为经过预先训练获得的；利用搬土距离EMD计算所述第一transformer层的输出与第二transformer层的输出之间的第一蒸馏损失，并分别计算所述第一embedding层的输出与第二embedding层的输出之间的第二蒸馏损失，以及所述第一prediction层的输出与第二prediction层的输出之间的第三蒸馏损失；根据所述第一蒸馏损失、所述第二蒸馏损失和所述第三蒸馏损失对所述学生模型中的参数进行优化，获得训练好的学生模型。

本申请实施例中，发明人在研发过程中发现Bert模型中的transformer层对模型的贡献最大，包含的信息最丰富，学生模型在该层的学习能力也最为重要，因此，基于EMD对老师模型进行蒸馏，能够保证学生模型学习到更多的老师模型的知识。

进一步地，所述利用EMD计算所述第一transformer层的输出与第二transformer层的输出之间的第一蒸馏损失，包括：获取所述第一transformer层中各层分别输出的第一注意力attention矩阵以及所述第二transformer层中各层分别输出的第二attention矩阵；根据所述第一attention矩阵和所述第二attention矩阵计算第一EMD距离；获取所述第一transformer层中各层分别输出的第一全连接前馈神经网络FFN隐层矩阵和所述第二transformer层中各层分别输出的第二FFN隐层矩阵；根据所述第一FFN隐层矩阵和所述第二FFN隐层矩阵计算第二EMD距离；根据所述第一EMD距离和所述第二EMD距离获得所述第一蒸馏损失。