模型小快好！微软预训练语言模型通用压缩方法MiniLM助你“事半功倍”

编者按：大规模预训练模型在自然语言理解和生成任务中表现优异，但巨大的参数量和计算成本让其很难直接部署到线上产品中。为此，微软亚洲研究院提出了一种将基于 Transformer 的预训练大模型压缩成预训练小模型的通用方法：深度自注意力知识蒸馏（Deep Self-Attention Distillation），让预训练模型又快又好。预训练模型和微调代码现已开源。

近年来，大规模预训练语言模型在自然语言理解和生成任务上都取得了优异的性能。预训练模型通过自监督学习（Self-supervised Learning）在大规模无标记文本学习通用的语言知识，在下游任务微调后就可以显著提升各个任务的结果。

然而，大规模预训练模型通常包含数以亿计的参数，例如 BERT Base 模型具有1.1亿参数，BERT Large 模型具有3.4亿参数，最近发布的一些大规模预训练模型甚至超过了百亿参数。海量的参数使得这些大模型在下游任务微调和推理解码时速度慢，对计算资源要求高，且计算成本大，因而将大模型直接部署到实际的线上产品时是很大的挑战。

相较于大模型，小模型则具有更少的参数和更快的速度，在实际产品中也更容易上线使用。但是小模型的效果相对大模型较差，如何有效地提升小模型的效果是一个重要的研究课题。Hinton 等人在2015年提出了知识蒸馏（Knowledge Distillation），把大模型当作“老师”，将其知识迁移给作为“学生”的小模型以提升其性能。知识蒸馏是一种有效的模型压缩方法，在学术界和工业界的实际产品中都得到了广泛的应用。目前大规模预训练模型的知识蒸馏主要针对具体的下游任务（Task-specific Distillation），由以下几步组成：

1） 将预训练模型在下游任务微调，得到针对特定任务的大模型。

2） 将微调后的大模型对相关的任务数据进行推理，以得到其类别预测的概率分布（Soft Labels）或是其隐层（Hidden States）以及自注意力矩阵分布/得分（Self-Attention Distributions/Scores）。这些向量以及概率分布蕴含大模型的知识。

3） 将这些知识作为指导信息应用于小模型的训练，训练目标是降低小模型和大模型在上述知识上的差异，使小模型尽可能模仿大模型。

尽管针对特定任务的知识蒸馏（Task-specific Distillation）很有效，但是对大模型微调和推理仍费时费力，计算成本也很高。而且上述过程仅针对某个特定任务，当我们换到另外一个任务时，我们需要重新微调大模型以及生成相关的大模型知识。此外，小模型的效果也受限于具体任务数据的数量。

还有一种做法就是将预训练大模型在大规模无标记数据上直接压缩成一个和下游任务无关的预训练小模型（Task-agnostic Distillation），该小模型可以在下游任务上直接进行微调。微调预训练小模型省时省力，对于计算资源的要求更低，计算成本更小，在不同任务间切换也更加方便。而且，预训练小模型可以用来初始化特定任务知识蒸馏中的“学生”模型，以进一步提升结果。相对于针对特定任务的知识蒸馏，直接将预训练大模型压缩成同下游任务无关的预训练小模型是一种更加理想的做法。

为此，微软亚洲研究院提出了一种将基于 Transformer 的预训练大模型压缩成预训练小模型（更少的层数和更小的隐层维度）的通用方法：深度自注意力知识蒸馏（Deep Self-Attention Distillation）。其核心思想是最大程度上将预训练 Transformer 模型中非常重要的自注意力（Self-Attention）知识迁移给小模型。需要特别提出的是，我们只迁移预训练大模型最后一层的自注意力知识。该方法主要有以下几个优点：

1） 压缩方法简单有效，可以应用在不同的预训练 Transformer 模型上。不同模型上的实验结果也验证了这一点。

2） 只需迁移大模型最后一层的自注意力知识，将迁移的知识进一步简化，提升训练速度。

3） 对小模型层数和隐层大小没有限制，让小模型的结构更加灵活。

我们的模型（以下称为 MiniLM ）在英文自然语言理解和生成任务上都取得了出色的结果。多语言 MiniLM 在跨语言理解任务上也取了很好的效果。现在，我们已将预训练模型和微调代码开源，欢迎大家使用、交流。

GitHub 地址：https://github.com/microsoft/unilm/tree/master/minilm

MiniLM 论文链接：https://arxiv.org/abs/2002.10957

Transformer 网络

目前，Transformer 网络被广泛地应用于自然语言处理领域的各个任务上，也是预训练语言模型的核心网络。每层 Transformer 网络主要由两个子层构成：自注意力层和前馈（Feed-forward）层。Transformer 网络的设计目标是舍弃传统的 RNN 结构，完全依靠自注意力机制来捕捉词与词间的依赖关系。所以自注意力层在 Transformer 网络中有着至关重要的作用。我们论文中的方法的核心思想是深度迁移自注意力层的知识。下面我们先简单介绍一下自注意力机制：

图1：自注意力运行机制

如图1所示，上一层的输出向量被送入三个不同的线性网络，转换得到三组向量：Queries，Keys 和 Values。然后 Queries 和 Keys 通过点积操作，得到自注意力得分矩阵（Attention Scores），每一个得分表示两两词之间的依赖或是关联程度。最后我们把这些分数指定到其对应的 Values 向量，加权求和后得到最终输出。每层 Transformer 会将上述过程并行重复多次（Multi-head）。在整个过程中， Queries，Keys 和 Values 是最基础的也是最重要的向量，我们论文提出的压缩方法也是围绕这三个向量进行知识迁移的。

深度自注意力知识蒸馏

论文中提出的模型压缩方法核心思想是迁移大模型自注意力知识，以让小模型自注意力模块的行为和大模型尽可能相似。但我们只迁移了大模型最后一层的自注意力知识。因为如果每一层都进行知识迁移，那么就需要手动设计大小模型间的层对应关系来完成迁移。相较而言，只迁移最后一层知识简单有效，训练速度更快，而且不需要手动设计大小模型的层对应关系。

图2整体介绍了深度自注意力知识蒸馏的方法，主要由两种知识迁移构成，第一种就是自注意力得分/分布迁移（Self-Attention Distribution Transfer），主要迁移自注意力得分/分布知识（Attention Scores/Distributions）。自注意力得分矩阵由Queries 和 Keys 通过点积操作得到，矩阵中每个值表示两个词的依赖关系。自注意力得分矩阵是自注意力模块中至关重要的知识，我们通过相对熵（KL-Divergence）来计算大模型和小模型自注意力得分矩阵的差异。

图2：深度自注意力知识蒸馏（Deep Self-Attention Distillation）

Queries 和 Keys 的知识已经通过自注意力得分矩阵进行了迁移。为了迁移更深层次的自注意力知识，使小模型可以更深层次地模仿大模型，我们引入了 Values 知识并将其转换为关系矩阵进行迁移。Values 关系矩阵（Value Relation）由 Values 向量间点积得到，可以表示 Values 词与词间的依赖。使用点积操作可以将大小模型不同维度的 Values 向量转换为相同维度大小的关系矩阵，避免引入额外的随机初始化参数对小模型的 Values 向量进行线性变换，以使其和大模型向量具有相同的维度来进行知识迁移。我们也通过相对熵来衡量大小模型间关系矩阵的差异。

在进行知识迁移时，如果大模型和小模型间的大小差距较大，那么可能导致迁移的效率较低。Mirzadeh 等人提出引入一个“助教”/中间模型（Teacher Assistant）来提升迁移效率，首先将大模型压缩成一个中间大小的模型，再将这个中间模型当作“助教”，将其压缩成最终的小模型。我们在将大模型压缩成更小模型的实验中，引入了中间模型来提升其结果。

综合上述，我们的论文提出的模型压缩方法主要有三个关键点：

1） 深度迁移最后一层自注意力知识。

2） 除了迁移自注意力得分矩阵，我们引入了 Values 关系矩阵来迁移更深层的词与词之间的依赖，迁移更深层的自注意力知识。

3） 通过使用中间模型来提升更小模型的压缩。

自然语言理解任务微调

我们在英文和多语言模型上进行了压缩实验，来验证我们方法的有效性。

英文单语言模型实验

我们首先对英文的 UniLMv2 Base 模型进行压缩，我们将12层768维度的 UniLMv2 Base 模型压缩成12层384维度（MiniLMa）和6层384维度（MiniLMb）两个模型，并且在常用的问答/阅读理解基准（SQuAD 2.0）和自然语言理解基准（GLUE）上进行实验。

实验结果如表1所示。12层384维的MiniLM模型相较于 BERT Base 模型（12层768维度），模型大小不到其1/3，推理速度是其2.7倍，在整体实验上取得了相对出色的结果。6层384维的 MiniLM 模型大小是 BERT Base 的1/5，推理速度是其5.3倍，也取得了不错的结果。

表1：英文 MiniLM 在自然语言理解任务上的结果

我们也将 BERT Base 模型当作“老师”，进行不同参数量小模型的实验，实验结果显示我们的方法相对于之前的知识蒸馏方法取得了更好的效果。更多详细信息请参考论文。

多语言模型实验

为了验证该方法的通用性，我们对多语言预训练模型 XLM-RoBERTa Base 进行了压缩。我们将其压缩成12层384维度（MiniLMa）和6层384维度（MiniLMb）两个模型，并且在跨语言推理基准（XNLI）以及跨语言问答（MLQA）基准上进行实验。

如表2和表3的结果所示，12层 MiniLM 模型的隐层维度不到多语言 BERT 和 XLM 模型的1/2，速度是其2.7/4.0+倍，但在两个跨语言任务上都取得了相对更好的结果。6层 MiniLM 模型也取得了不错的结果。

表2：多语言 MiniLM 在跨语言推理（XNLI）上的结果

表3：多语言 MiniLM 在跨语言问答（MLQA）上的结果

自然语言生成任务微调

通过使用 UniLM 提出的针对端到端（Sequence to Sequence）生成模型的特定自注意力掩码（Self-Attention Mask），我们对上文中12层和6层英文 MiniLM 模型在生成任务上进行了微调实验。我们使用 UniLM 在 GitHub 开源的 s2s-ft 工具包在生成任务上进行微调，在生成式摘要（Abstractive Summarization）以及问题生成（Question Generation）两个任务上进行了实验。

表4和表5列举了 MiniLM 模型在生成式摘要和问题生成任务上的结果，12层 MiniLM 模型大小不到一些 Base 模型的1/3，但是取得了相近甚至更好的结果。6层 MiniLM 在生成式摘要上也取得了不错的结果。在问题生成任务上，MiniLM 也有着不错的表现。

表4：英文 MiniLM 在生成式摘要（Abstractive Summarization）上的结果

表5：英文 MiniLM 在问题生成（Question Generation）上的结果

总结

本文介绍了基于 Transformer 预训练语言模型的通用压缩方法：深度自注意力知识蒸馏（Deep Self-Attention Distillation），深度迁移“老师”模型最后一层自注意力知识。该压缩方法简单有效，由不同预训练大模型压缩得到的英文和多语言 MiniLM 预训练模型更小更快，在自然语言理解和生成任务上均取得了出色的结果。

论文及代码

MiniLM 相关的论文、预训练模型以及代码样例，请访问以下链接。

MiniLM 论文链接：https://arxiv.org/abs/2002.10957

UniLMv2 论文链接：https://arxiv.org/abs/2002.12804

MiniLM 代码链接：https://github.com/microsoft/unilm/tree/master/minilm

s2s-ft 代码链接：https://github.com/microsoft/unilm/tree/master/s2s-ft