大型语言模型：TinyBERT – 为自然语言处理提取BERT知识
超级语言模型：TinyBERT - 将BERT知识提升到自然语言处理的巅峰
解锁在LLMs中使用Transformer蒸馏的力量
引言
近年来，大型语言模型的发展迅猛。BERT成为最受欢迎和高效的模型之一，能够以高准确性解决各种NLP任务。在BERT之后，一系列其他模型也相继出现，展示出色的结果。
一个明显的趋势是，大型语言模型（LLMs）随着时间的推移变得越来越复杂，通过指数级增加其参数和训练数据。深度学习研究表明，这种技术通常会带来更好的结果。不幸的是，机器学习领域已经遇到了与LLMs相关的几个问题，可扩展性成为有效训练、存储和使用它们的主要障碍。
考虑到这个问题，已经制定了特殊的方法来压缩LLMs。在本文中，我们将重点介绍Transformer蒸馏，这导致了BERT的一个小版本称为TinyBERT的开发。此外，我们将理解TinyBERT的学习过程以及使TinyBERT如此强大的一些微妙之处。本文基于官方的TinyBERT论文。
主要思想
最近，我们已经介绍了DistilBERT中蒸馏的工作原理：简单来说，损失函数目标被修改为使学生和教师的预测相似。在DistilBERT中，损失函数比较了学生和教师的输出分布，并且还考虑了两个模型的输出嵌入（用于相似性损失）。
大型语言模型：DistilBERT – 更小、更快、更便宜和更轻便
解锁BERT压缩的秘密：最大效率的学生-教师框架
towardsdatascience.com
在TinyBERT中，与DistilBERT相比，蒸馏框架在表面上没有太大的变化：损失函数再次被修改为使学生模仿教师。然而，TinyBERT在这一点上更进了一步：损失函数不仅考虑两个模型的输出，还考虑了预测是如何获得的。根据论文，TinyBERT的损失函数由三个组成部分组成，涵盖了两个模型的不同方面：
- 嵌入层的输出
- Transformer层导出的隐藏状态和注意力矩阵
3. 预测层输出的逻辑。
比较两个模型的隐藏状态有什么意义？ 包括隐藏状态和注意力矩阵的输出使得学生能够学习教师的隐藏层，从而构建类似于教师的层。这样，蒸馏模型不仅模仿原始模型的输出，还模仿其内部行为。
复制教师的行为为什么很重要？ 研究人员声称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的设置将在本文稍后进行讨论。
Transformer知识蒸馏
1. 嵌入层蒸馏
在将原始输入传递给模型之前，首先对其进行分词处理，然后将其映射到已学习的嵌入层。这些嵌入层随后被用作模型的第一层。所有可能的嵌入可以表示为一个矩阵。为了比较学生和教师嵌入之间的差异程度，可以使用一个标准的回归指标应用于它们各自的嵌入矩阵 E。例如，Transformer蒸馏使用均方误差(MSE)作为回归度量。
由于学生和教师嵌入矩阵的大小不同，无法使用MSE逐个元素地比较它们。这就是为什么将学生嵌入矩阵乘以一个可学习的权重矩阵W的原因，从而使得得到的矩阵与教师嵌入矩阵具有相同的形状。
由于学生和教师的嵌入空间是不同的，矩阵W在线性变换学生的嵌入空间以使之与教师的嵌入空间相匹配中起着重要作用。
2. Transformer层蒸馏
2A. 注意力层蒸馏
在Transformer中，多头注意力机制产生了包含丰富语言知识的多个注意力矩阵。通过从教师那里传输注意力权重，学生也能理解重要的语言概念。为了实现这个思想，使用损失函数计算学生和教师注意力权重之间的差异。
在TinyBERT中，考虑了所有的注意力层，并且每个层的损失值等于各自学生和教师注意力矩阵的所有头部的MSE值之和。
用于注意力层蒸馏的注意力矩阵A是非归一化的，而不是它们的softmax输出softmax(A)。根据研究人员的说法，这个细微差别可以导致更快的收敛和更好的性能。
2B. 隐藏层蒸馏
在捕获丰富的语言知识的思想下，蒸馏还应用于Transformer层的输出。
权重矩阵W起到了与嵌入层蒸馏中描述的相同的作用。
3. 预测层蒸馏
最后，为了使学生能够复制老师的输出，考虑了预测层损失。这包括计算两个模型预测的逻辑向量之间的交叉熵。
有时，逻辑向量会被温度参数 T 整除，这控制了输出分布的平滑度。在 TinyBERT 中，温度 T 被设定为 1。
损失方程
在 TinyBERT 中，每个层根据其类型都有自己的损失函数。为了赋予某些层更多或更少的重要性，相应的损失值会乘以一个常数 a. 最终的损失函数等于 TinyBERT 所有层上的损失值的加权和。
在众多实验证明，三个损失组成部分中，变压器层蒸馏损失对模型性能的影响最大。
训练
重要的是要注意，大多数 NLP 模型（包括 BERT）是在两个阶段中开发的:
- 模型在大型语料库上进行预训练，以获得对语言结构的一般知识。
- 模型在另一个数据集上进行微调，以解决特定的下游任务。
按照相同的范式，研究人员开发了一种框架，其中 TinyBERT 的学习过程也包括两个阶段。 在两个训练阶段中，都使用 Transformer 蒸馏将 BERT 的知识传递给 TinyBERT。
- 通用蒸馏。TinyBERT 从预训练的 BERT（没有微调）中获得了丰富的通用语言结构知识，此阶段后，TinyBERT 的性能通常比 BERT 差。
- 任务特定蒸馏。这次，微调版本的 BERT 扮演了老师的角色。为了进一步提高性能，研究人员提议在训练数据集上应用数据增强方法。结果显示，在任务特定蒸馏之后，TinyBERT 的性能与 BERT 相当。
数据增强
针对任务特定蒸馏，开发了一种特殊的数据增强技术。它包括从给定数据集中获取序列，并以以下两种方式替换一定比例的单词：
- 如果单词被分词为相同的单词，则使用 BERT 模型预测该词，并将预测的词替代序列中的原始词。
- 如果单词被分词为多个子词，则将这些子词替换为最相似的 GloVe 嵌入。
尽管模型尺寸大大减小，但所描述的数据增强机制通过让 TinyBERT 学习更多样化的示例对 TinyBERT 的性能产生了重大影响。
模型设置
由于仅有 14.5M 参数，TinyBERT 的尺寸约为 BERT base 的 7.5 倍小。它们的详细比较如下图所示:
对于层映射，作者提出了一种统一的策略，即将每个TinyBERT层映射到每三个BERT层：g(m) = 3 * m。也研究了其他策略（如采用所有底部或顶部BERT层），但统一策略显示出最佳结果，这似乎是合理的，因为它允许从不同的抽象层传递知识，使传递的信息更加丰富多样。
谈到训练过程，TinyBERT是在英文维基百科（2500M字）上进行训练的，并且它的大部分超参数与BERT base相同。
结论
Transformer蒸馏是自然语言处理中的一个重要步骤。考虑到基于Transformer的模型目前在机器学习中最强大，我们可以通过应用Transformer蒸馏来更有效地压缩它们。TinyBERT是一个很好的例子，它是从BERT base中压缩了7.5倍。
尽管参数大幅减少，但实验证明TinyBERT与BERT base具有可比性能：在GLUE基准测试中获得了77.0%的分数，TinyBERT离BERT的79.5%分数并不远。显然，这是一个惊人的成就！最后，还可以对TinyBERT应用其他流行的压缩技术，如量化或剪枝，使其变得更小。
资源
除非另有说明，所有图片都由作者提供
