TinyBERT

部分内容转载自机器之心

---

TinyBERT的创新点主要在：

新型蒸馏方式

不同于普通的知识蒸馏（knowledge distillation，KD）方法只针对输出层logits的soft label进行student模型的学习，TinyBERT应用的KD方法还对embedding层、每个transformer layer中的多头注意力层、以及每个transformer layer输出的hidden states和最后一层预测层这四部分进行蒸馏。

TinyBERT蒸馏概览

通过以上几个蒸馏目标，可以整合 teacher 和 student 网络之间对应层的蒸馏损失(具体每个目标的loss见转载原文)：

TinyBERT loss function

两段式学习框架

BERT 的应用通常包含两个学习阶段：预训练和微调。BERT 在预训练阶段学到的大量知识非常重要，并且迁移的时候也应该包含在内。因此，研究者提出了一个两段式学习框架，包含通用蒸馏和特定于任务的蒸馏，（说白了就是预训练阶段也进行蒸馏）如下图所示：

TinyBERT 的两段式学习图示

通用蒸馏可以帮助 student TinyBERT 学习到 teacher BERT 中嵌入的丰富知识，对于提升 TinyBERT 的泛化能力至关重要。特定于任务的蒸馏赋予 student 模型特定于任务的知识。这种两段式蒸馏可以缩小 teacher 和 student 模型之间的差距。

1. 通用蒸馏。在通用蒸馏中，研究者使用原始 BERT 作为 teacher 模型，而且不对其进行微调，利用大规模文本语料库作为学习数据。通过在通用领域文本上执行 Transformer 蒸馏，他们获取了一个通用 TinyBERT，可以针对下游任务进行微调。然而，由于隐藏/嵌入层大小及层数显著降低，通用 TinyBERT 的表现不如 BERT。
2. 针对特定任务的蒸馏。研究者提出通过针对特定任务的蒸馏来获得有竞争力的微调 TinyBERT 模型。而在蒸馏过程中，他们在针对特定任务的增强数据集上（如图 2 所示）重新执行了提出的 Transformer 蒸馏。具体而言，微调的 BERT 用作 teacher 模型，并提出以数据增强方法来扩展针对特定任务的训练集。

3. 数据增强。使用BERT和Glove对句子序列进行词的替换，即通过BERT预测mask住的词，再使用Glove选择和这个词最相近的词（cosine similarity），进行替换，从而达到数据增强的目的。

   
   
   The Proposed Data Augmentation

消融实验

本文提出的两段式 TinyBERT 学习框架包含三个关键步骤：TD（特定于任务的蒸馏）、GD（通用蒸馏）和 DA（数据蒸馏）。每个学习步骤的影响如下表所示：

不同学习步骤的影响

研究者还探索了不同蒸馏目标对 TinyBERT 学习的影响，结果如下表：

不同蒸馏目标的影响

研究者探究了teacher模型和student模型transformer层的不同映射函数 n = g(m) 对于 TinyBERT 学习的影响，比较结果见下表：

不同映射策略的影响