[发明专利]基于自注意力的知识蒸馏方法、装置和计算机设备

说明书

本申请涉及人工智能领域，揭示了基于自注意力的知识蒸馏方法，包括：将输入数据输入第一模型得到第一模型的中间层输出的第一特征矩阵，将输入数据输入第二模型得到第二模型的中间层输出的第二特征矩阵，其中，第一模型为训练好的老师模型，第二模型为待训练的学生模型；根据第一特征矩阵计算老师模型对应的第一自注意力权重分布，根据第二特征矩阵计算学生模型对应的第二自注意力权重分布；计算第一自注意力权重分布和第二自注意力权重分布之间的分布差异；将分布差异，作为老师模型和学生模型之间的知识蒸馏损失函数；根据知识蒸馏损失函数，将老师模型的中间层的数据映射关系迁移至学生模型的中间层上，能满足不同任务类型模型的知识蒸馏训练。

技术领域

本申请涉及人工智能领域，特别是涉及到基于自注意力的知识蒸馏方法、装置和计算机设备。

背景技术

知识蒸馏（Knowledge Distillation）是一种特殊的迁移学习方法，目的主要是在保证训练效果的同时对待训练模型的体积进行压缩。通过训练好的老师模型来指导小体积的待训练学生模型的学习，通过训练小模型学习到大模型的知识，相比于直接训练小模型效果更好，速度更快。

目前用于知识蒸馏的损失函数更多的是针对分类模型进行的，其要求大模型和小模型的类别数或者网络输出特征维度要一致，限制了知识蒸馏的应用范围，不能满足不同任务类型模型的知识蒸馏训练。

发明内容

本申请的主要目的为提供基于自注意力的知识蒸馏方法，旨在解决现有知识蒸馏的损失函数设计，不能满足不同任务类型模型的知识蒸馏训练的技术问题。

本申请提出一种基于自注意力的知识蒸馏方法，包括：

将输入数据输入第一模型得到所述第一模型的中间层输出的第一特征矩阵，将所述输入数据输入第二模型得到所述第二模型的中间层输出的第二特征矩阵，其中，所述第一模型为训练好的老师模型，所述第二模型为待训练的学生模型，所述第一特征矩阵和所述第二特征矩阵具有相同的序列长度；

根据所述第一特征矩阵计算所述老师模型对应的第一自注意力权重分布，根据所述第二特征矩阵计算所述学生模型对应的第二自注意力权重分布；

计算所述第一自注意力权重分布和所述第二自注意力权重分布之间的分布差异；

将所述分布差异，作为所述老师模型和所述学生模型之间的知识蒸馏损失函数；

根据所述知识蒸馏损失函数，将所述老师模型的中间层的数据映射关系迁移至所述学生模型的中间层上。

优选地，所述根据所述第一特征矩阵计算所述老师模型对应的第一自注意力权重分布的步骤，包括：

将所述第一特征矩阵进行转置计算，得到所述第一特征矩阵对应的第一转置矩阵；

根据所述第一特征矩阵和所述第一转置矩阵，计算所述第一特征矩阵的内部特征关系；

对所述第一特征矩阵的内部特征关系进行softmax函数计算，得到所述老师模型对应的第一自注意力权重分布。

优选地，所述老师模型使用多头注意力机制，所述根据所述第一特征矩阵计算所述老师模型对应的第一自注意力权重分布的步骤，包括：

将所述第一特征矩阵按照所述多头注意力机制对应的头数进行均分，得到多个子矩阵；

对第一子矩阵进行转置计算，得到所述第一子矩阵对应的第一转置子矩阵，其中，所述第一子矩阵为所述第一特征矩阵对应的多个子矩阵中的任一个；

根据所述第一子矩阵和所述第一转置子矩阵，计算所述第一子矩阵的内部特征关系；

根据所述第一子矩阵的内部特征关系的计算方式，计算所述第一特征矩阵的各子矩阵分别对应的内部特征关系；