Smaller, faster, cheaper, lighter: Introducing DistilBERT, a distilled version of BERT 翻译

paper: https://arxiv.org/pdf/1910.01108v2.pdf

code: https://github.com/huggingface/transformers

Time: 2019.10

在过去的18个月中，在过去的18个月中，几乎所有的自然语言处理任务都从大型语言模型进行迁移学习这一方式达到了SOTA效果。通常基于Vaswani等人的Transformer体系结构，这些经过预先训练的语言模型会变得越来越大，并在更大的数据集上进行训练。 Nvidia的最新模型具有83亿个参数：比BERT-large大24倍，比GPT-2大5倍，而来自Facebook AI的最新作品RoBERTa受过160GB文本的训练

 

在Hugging Face，我们亲身体验了这些模型的日益普及，因为封装了大多数模型的NLP库在短短几个月内被安装了40万次以上。但是，随着这些模型进入更大的NLP社区，一个重要而具有挑战性的问题开始出现。 我们应该如何将这些怪物投入生产？ 我们如何在低延迟约束下使用如此大的模型？ 我们需要（昂贵的）GPU服务器进行大规模服务吗？

For many researchers and developers, these can be deal-breaking issues

为了构建更多个人可用的系统，我们注意到越来越需要使机器学习系统在边缘运行，而不是调用云API并将可能的私有数据发送到服务器。在智能手机devices等设备上运行模型还需要轻巧，响应迅速且节能的模型！

最后同样重要的一点是，我们越来越关注按比例缩放这些模型的计算需求的环境成本。

So, how can we reduce the size of these monster models⁉️

有许多技术可以解决前面的问题。 最常用的工具包括量化(quantization)（以较小的精度近似估计网络的权重）和权重修剪(weights pruning)（删除网络中的某些连接）。 对于这些技术，您可以看一下Rasa出色的关于BERT量化的博客文章。我们决定专注于蒸馏(distillation)：一种可以用来将大型模型（称为教师）压缩为较小模型（称为学生）的技术。

知识蒸馏-转移归纳能力

知识蒸馏（有时也称为师生学习）是一种压缩技术，是训练了一个小模型来重现一个较大模型（或模型集合）的行为。它是由Bucila等人提出的，并由Hinton等人概括。 我们将遵循后一种方法。

在监督学习中，一般训练分类模型以通过使用对数似然概率最大化（softmax of logits）来预测目标分类。 在许多情况下，良好的性能模型将有正确的类别具有高概率，而其他类别的概率接近零的输出分布。

But, some of these “almost-zero” probabilities are larger than the others, and this reflects, in part, the generalization capabilities of the model.

例如，桌椅可能会误认为扶手椅，但通常不应误认为是蘑菇。 这种不确定性有时被称为“暗知识”。

理解蒸馏的另一种方式是，它防止模型对其预测过于确定（类似于标签平滑）。

这是在实践中看到此想法的示例。 在语言建模中，我们可以通过查看词汇表的分布轻松观察到这种不确定性。 以下是BERT完成卡萨布兰卡电影中这句名言的前20个猜测：

 

我们如何复制这些暗知识？

在师生训练(teacher-student training )中，我们训练学生网络以模仿教师网络的全部输出分布（其知识）。

We are training the student to generalize the same way as the teacher by matching the output distribution.

与其对硬目标（one-hot encoding of the gold class ）进行交叉熵训练，我们对软标签(soft-label)（教师的概率输出）进行交叉熵训练，将知识从教师转移到学生。 因此，我们的训练损失为：

 

这种损失是一个更丰富的训练信号，因为单个示例比单个硬目标要执行的约束要多得多。

为了进一步揭示类别上分布的质量，Hinton等人引入了softmax-temperature：

 

当T→0时，该分布成为Kronecker（相当于一个one-hot目标矢量），当T→+∞时，它成为均匀分布。 在训练时，将相同的温度参数应用于学生和老师，从而进一步揭示了每个训练示例的更多信号。 在推断时，T设置为1并恢复标准Softmax。

Hands-on coding in PyTorch — 压缩 BERT

我们想使用蒸馏来压缩大型语言模型。 为了进行提炼，我们将使用Kullback-Leibler loss，因为这些优化是等效的：

 

在计算关于q（学生分布）的梯度时，我们获得相同的梯度。 它使我们能够利用PyTorch实现更快地进行计算：

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.optim import Optimizer
​
KD_loss = nn.KLDivLoss(reduction='batchmean')
​
def kd_step(teacher:nn.Module,
            student:nn.Module,
            temperature:float,
            inputs:torch.tensor,
            optimizer:Optimizer):
    teacher.eval()
    student.train()
    
    with torch.no_grad():
        logits_t = teacher(inputs=inputs)
    logits_s = student(inputs=inputs)
    
    loss = KD_loss(input=F.log_softmax(logits_s/temperature, dim=-1),
                   target=F.softmax(logits_t/temperature, dim=-1))
    
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()

使用老师的信号，我们可以在BERT（使用BERT的English bert-uncased版本）的监督下，训练一个较小的语言模型，称为DistilBERT。

按照Hinton等人的方法，训练损失是蒸馏损失(distillation loss)和掩盖语言建模损失(language modeling loss)的线性组合。 我们的学生是BERT的一个小版本，在其中我们删除了令牌类型的嵌入和pooler（用于下一个句子分类任务），并保持了体系结构的其余部分不变，同时将层数减少了两倍。

总体而言，我们的提炼模型DistilBERT具有大约BERT基础参数总数的一半，并且在语言理解基准GLUE上保留了BERT 95％的性能。

Note 1 — Why not reducing the hidden size as well?

将其从768减少到512确实可以使参数总数减少。 但是，张量的最后一个维度（隐藏维度）的变化对Transformer体系结构中使用的大多数操作（线性层和层归一化）影响很小。在我们的实验中，层数是推理时间的决定因素，而不是隐藏层的大小。所以减少hidden size没啥效果。

Note 2 — Some works on distillation like Tang et al. use the L2 distance as a distillation loss directly on downstream tasks.

我们的早期实验表明，在我们的案例中，交叉熵损失导致性能显着提高。 我们假设在语言建模设置中，输出空间（词汇）显着大于下游任务输出空间的维度。 因此，对数可以在L2损失中彼此补偿。

训练子网不仅涉及架构，这也涉及为子网收敛找到正确的初始化方式。因此，我们利用老师和学生之间的hidden size相同这一特点，从老师BERT两层中选出一层来初始化学生DistilBERT。

我们还使用了最近RoBERTa论文中的一些训练技巧，这些技巧表明BERT的训练方式对其最终性能至关重要。 继RoBERTa之后，我们利用梯度累积、动态屏蔽对DistillBERT进行了大批量训练（每批最多4000个示例），并删除了下一个句子预测任务。

Our training setup is voluntarily limited in terms of resources. We train DistilBERT on eight 16GB V100 GPUs for approximately three and a half days using the concatenation of Toronto Book Corpus and English Wikipedia (same data as original BERT). （8块16GB的V100GPU，羡慕...）

模型性能-测试DistilBERT

性能比较：

时间比较：

 

In terms of inference time, DistilBERT is more than 60% faster and smaller than BERT and 120% faster and smaller than ELMo+BiLSTM , 确实。

Downstream task: Distillation & transfer-learning

略。