TinyBERT模型解析——压缩BERT模型的精髓

译者 | 朱先忠

为了帮助尚未通关的玩家，下面是一些解谜技巧，希望对大家有所帮助。

简介

近年来，大型语言模型的发展突飞猛进。BERT成为最受欢迎和最有效的模型之一，可以高精度地解决各种自然语言处理（NLP）任务。继BERT模型之后，一组其他的模型也先后出现并各自展示出优秀的性能。

不难看到一个明显趋势是，随着时间的推移，大型语言模型（LLM）往往会因其训练的参数和数据数量呈指数级增加而变得更加复杂。深度学习研究表明，这种技术通常会带来更好的运行结果。然而，遗憾的是，尽管机器学习世界已经克服了不少关于大型语言模型相关的问题；但是，可扩展性的问题已经成为有效训练、存储和使用大型语言模型的主要障碍。

考虑到上述问题，人们已经开发出不少的压缩大型语言模型的特殊方法。在这篇文章中，我们将重点讨论转换器蒸馏方法，这种方法诞生了名为TinyBERT的一个迷你版本的BERT模型。此外，我们还将介绍TinyBERT模型的学习过程，以及使TinyBERT模型变得如此强大的几个微妙原因。本文基于TinyBERT的官方论文整理而成。

主要思想

在最近的文章中，我们已经讨论了DistilBERT模型中蒸馏技术的工作原理。简而言之，蒸馏技术的主要思想是：修改损失函数目标，以便使学生模型和教师模型的预测结果相似。在DistilBERT模型中，损失函数比较学生模型和教师模型的输出分布，并兼顾两个模型的输出嵌入（针对相似性损失）。

有关DistilBERT模型的更多的细节，请参考文章《Large Language Models: DistilBERT — Smaller, Faster, Cheaper and Lighter》，地址是：
“https://towardsdatascience.com/distilbert-11c8810d29fc?source=post_page-----1a928ba3082b--------------------------------”。此文的主要内容介绍了BERT模型压缩的秘密，目标是实现师生模型框架效率的最大化。

从表面上看，TinyBERT模型中的蒸馏框架与DistilBERT模型没有太大变化：再次修改了损失函数，目标是使学生模型模仿教师模型。然而，在TinyBERT模型的情况下，它更进了一步：损失函数不仅考虑了师生两个模型产生的结果，还考虑了如何获得预测结果的问题。根据TinyBERT模型论文作者介绍，TinyBERT损失函数由三个部分组成，它们涵盖了师生两个模型的不同方面：

1. 嵌入层的输出
2. 从转换器层导出的隐藏状态和注意力矩阵
3. 预测层输出的logits值

转换器蒸馏损失函数示意图

那么，比较师生两种模型的隐藏状态有什么意义呢？通过包括隐藏状态和注意力的输出结果，注意力矩阵使得学生模型有可能学习教师模型的隐藏层内容，从而构建与教师模型相似的层。这样，提取的模型不仅可以模仿原始模型的输出，而且模仿其内部行为。

那么，为什么复制教师模型的行为很重要呢？研究人员声称，通过BERT模型学习到的注意力权重有利于捕捉语言结构。因此，它们对另一种模式的蒸馏也给了学生模型更多获得语言知识的机会。

层映射

TinyBERT模型仅代表一种较小的BERT版本，具有较少的编码器层。现在，不妨让我们将BERT模型层数定义为N，将TinyBERT模型层数定义为M。鉴于层数不同，如何计算蒸馏损失值的问题尚不明确。

为此，引入了一个特殊函数n=g（m）来定义哪个BERT模型层n用于将其知识提取到TinyBERT模型中的相应层m。然后，所选择的BERT层用于训练期间的损失值计算。

引入的函数n=g（m）具有两个推理约束：

• g（0）=0。这意味着，BERT模型中的嵌入层被直接映射到TinyBERT模型中的嵌入图层，这是有意义的。
• g（M+1）=N+1。该等式指示，BERT模型中的预测层被映射到TinyBERT模型中的预测层。对于所有其他TinyBERT模型中满足条件1≤m≤m的那些层，需要映射n=g（m）的相应函数值。现在，假设已经定义好了这样的函数。
  有关TinyBERT模型设置的问题，将在本文稍后进行研究。

转换器蒸馏

1.嵌入层蒸馏

原始输入在被传递到模型之前，首先被标记化，然后被映射到学习的嵌入层。然后，这些嵌入层被用作模型的第一层。所有可能的嵌入层都可以用矩阵的形式表示。为了比较学生模型和教师模型的嵌入层有多大的不同，可以在他们各自的嵌入矩阵E上使用标准回归度量。例如，转换器蒸馏使用均方误差（MSE）作为回归度量。

由于学生模型和教师模型的嵌入矩阵具有不同的大小，因此不可能通过使用均方误差来明智地比较它们的元素。这就解释了为什么学生模型嵌入矩阵乘以可学习的权重矩阵W，从而导致结果矩阵与教师模型嵌入矩阵具有相同的形状。

嵌入层蒸馏损失函数。

由于学生模型和教师模型的嵌入空间是不同的，矩阵W在将学生模型的嵌入空间线性转换为教师模型嵌入空间方面也起着重要作用。

2.转换器层蒸馏

转换器层蒸馏损失函数可视化展示

2A. 注意力层蒸馏

转换器中的多头注意力机制的核心是生成包含丰富语言知识的多个注意力矩阵。通过转移教师模型的注意力权重，学生模型也可以理解重要的语言概念。为了实现这一思想，使用损失函数来计算学生模型和教师模型注意力权重之间的差异。

在TinyBERT模型中，考虑了所有的注意力层，并且每一层的最终损失值等于所有头部的相应学生模型和教师模型注意力矩阵之间的均方误差值之和。

注意层蒸馏损失函数计算公式

值得注意的是，用于注意力层提取的注意力矩阵A是未归一化的，而不是它们的softmax输出softmax（A）。根据研究人员的说法，这种微妙之处有助于更快地收敛并提高性能。

2B. 隐藏层蒸馏

为了实现获取丰富语言知识的想法，蒸馏操作也被应用到转换器层的输出上。

隐藏层蒸馏损失函数计算公式。

这里，权重矩阵W起到与上述用于嵌入层蒸馏的权重矩阵相同的作用。

3.预测层蒸馏

最后，为了使学生模型再现教师模型的输出结果，使用了预测层损失函数。它包括计算两个模型预测的logit向量之间的交叉熵。

预测层蒸馏损失函数计算公式

值得注意的是，有些情况下，logits要除以控制输出分布的平滑度的温度参数T。在TinyBERT模型中，温度参数T设置为1。

损失方程

在TinyBERT模型中，根据其类型特征，每一层都有自己的损失函数。考虑到某些层或多或少的重要性作用，将相应的损失值乘以常数a。最终的损失函数等于所有TinyBERT模型层的损失值的加权和。

TinyBERT模型中的损失函数计算公式

大量实验表明，在三种损失分量中，转换器层蒸馏损失对模型性能的影响最大。

模型训练

需要注意的是，大多数自然语言处理模型（包括BERT）开发过程可大致划分为两个阶段：

1. 在一个大型数据语料库上对模型进行预训练，以获得语言结构的一般知识。
2. 在另一个数据集上对模型进行微调，以解决特定的下游任务。

遵循与此同样的思想，研究人员研发了一个新的框架TinyBERT，它的学习过程也是由类似上面的两个阶段组成的。在这两个训练阶段中，使用转换器蒸馏算法将BERT模型知识转换成TinyBERT模型。

阶段一：普通蒸馏。TinyBERT作为一个教师模型，通过预先训练（无需微调）的BERT模型获得了丰富的语言结构常识。通过使用更少的层和参数，在这个阶段之后，TinyBERT模型的性能通常比BERT模型差一些。

阶段二：特定任务的蒸馏。这一次，微调版的BERT模型扮演了教师模型的角色。为了进一步提高性能，正如研究人员所提出的，在训练数据集上应用了数据增强方法。实验结果表明，经过特定任务的蒸馏操作后，TinyBERT模型在BERT模型方面取得了相当的性能。

TinyBERT模型训练流程示意图

数据增强

针对特定任务的蒸馏，引入了一种特殊的数据增强技术。在这种数据增强技术中，首先从给定的数据集中提取序列，然后以下列两种方式之一替换一定百分比的单词：

• 如果某单词被标记为同一个单词，则该单词由BERT模型预测，并且用预测后的结果单词替换序列中的原始单词。
• 如果单词被标记为几个子单词，那么这些子单词将被最相似的GloVe嵌入（全局向量的词嵌入：Global Vectors for Word Representation）所取代。

尽管模型大小显著减小，但是所描述的数据增强机制通过允许TinyBERT学习更多不同的示例，对其性能产生了很大影响。

数据增强技术示意图

模型设置

由于只有14.5M个参数，TinyBERT模型比基础型BERT模型小约7.5倍。它们的详细比较如下图所示：

基础BERT模型与TinyBERT模型比较

对于层映射，论文作者提出了一种统一的策略。根据该策略，层映射函数将每个TinyBERT层映射到按序排序的每三个为一组的BERT层中的第一个：g（m）=3*m。论文作者还研究了其他的策略（如采用所有底部或顶部BERT层），但仅有统一策略显示出最佳实验结果。这个结论似乎是比较合乎逻辑的，因为它允许从不同的抽象层转移知识，使转移的信息更加多样化。

基于不同的层映射策略情况：图中展示了基于GLUE数据集的性能比较结果

就训练实现过程方面，TinyBERT模型是在英语维基百科（2500M个单词）上训练的，其大多数超参数与BERT模型库中使用的相同。

结论

转换器蒸馏是自然语言处理中的一项重要措施。考虑到基于转换器的模型是目前机器学习中最强大的模型之一，我们可以通过应用转换器蒸馏来有效地压缩它们来进一步开发利用它们。这方面最伟大的例子之一是TinyBERT模型，它在BERT模型基础上压缩了7.5倍。

尽管参数大幅减少，但实验表明，TinyBERT模型的性能与BERT基础模型基本相当：在GLUE基准数据集上的测试结果表明，TinyBERT模型获得77.0%的得分，与得分为79.5%的BERT模型相距并不远。显然，这是一个惊人的成就！最后，其他的一些流行的压缩技术，如量化或修剪等，都可以应用于TinyBERT模型压缩算法，从而使此模型体积变得更小。

除非另有说明，否则本文中所有图片均由作者本人提供。

参考资料

• TinyBERT: Distilling BERT for Natural Language Understanding：https://arxiv.org/pdf/1909.10351.pdf。

译者介绍

朱先忠，51CTO社区编辑，51CTO专家博客、讲师，潍坊一所高校计算机教师，自由编程界老兵一枚。

原文标题：Large Language Models: TinyBERT — Distilling BERT for NLP，作者：Vyacheslav Efimov