魔王编译

小版BERT也能出奇迹:最火的预训练语言库探索小巧之路

过去一段时间,大模型层出不穷。在大家纷纷感叹「大力出奇迹」的时候,作为调用预训练语言模型最流行的库,HuggingFace 尝试用更少的参数量、更少的训练资源实现同等的性能,于是知识蒸馏版 BERT——DistilBERT 应运而生!

近日,HuggingFace 发布了 NLP transformer 模型——DistilBERT,该模型与 BERT 的架构类似,不过它仅使用了 6600 万参数(区别于 BERT_base 的 1.1 亿参数),却在 GLUE 基准上实现了后者 95% 的性能。

在过去 18 个月中,基于大规模语言模型迁移学习显著提升了自然语言处理任务的 SOTA 成绩。这些预训练语言模型通常基于 Vaswani 等人提出的 Transformer 架构,这些模型的规模越来越大,训练数据集的规模也越来越大。最近英伟达提出的预训练语言模型拥有 83 亿参数:是 BERT-large 参数量的 24 倍、GPT-2 参数量的 5 倍。而 Facebook AI 最近提出的 RoBERTa 模型在 160GB 文本上训练得到。

社区中的一些人质疑训练越来越大 Transformer 的必要性,尤其是考虑到训练的资金成本和环境成本时。该图展示了部分近期大模型及其参数量。

Hugging Face 直接体会到这些模型的流行度,因为其预训练语言库(包含这些模型中的大部分)在近几个月的安装量超过 40 万次。

  • NLP 库地址:https://github.com/huggingface/pytorch-transformers

然而,尽管这些模型被更大的 NLP 社区接受,一个重要且有挑战性的问题出现了。如何将这些庞然大物投入到生产中?如何在低延迟约束下使用这些大模型?我们需要用(昂贵的)GPU 服务器执行大规模服务吗?

为了构建更尊重隐私的系统,Hugging Face 注意到在边缘设备上运行机器学习系统的需求在不断增长,而不是调用云 API,将隐私数据发往服务器。在智能手机等设备上运行的模型需要是轻量级、响应快和能源利用率高的!

最后但同样重要的一点,Hugging Face 越来越担忧这些大模型所需的指数级计算成本增长。

有很多技术可以解决前述问题。最常见的工具是量化(使用更小精度逼近全精度模型)和权重剪枝(移除网络中的部分连接)。想了解更多,可以查看这篇关于 BERT 量化的精彩博客:https://blog.rasa.com/compressing-bert-for-faster-prediction-2/。

Hugging Face 的研究者决定把重点放在知识蒸馏(distillation)上。蒸馏即,将较大模型(教师模型)压缩成较小模型(学生模型)的方法。

知识蒸馏迁移泛化能力

知识蒸馏是一种模型压缩方法,又叫师生学习(teacher-student learning)。它训练一个小模型,使之复制大模型(或模型集成)的行为。知识蒸馏由 Bucila 等人提出,几年后被 Hinton 等人推广(参见论文《Distilling the Knowledge in a Neural Network》)。Hugging Face 研究者使用的是 Hinton 等人的方法。

在监督学习中,分类模型通常用于预测类别,它利用对数似然信号最大化类别概率。在很多案例中,高性能模型预测的输出分布中,正确的类别具备高概率,而其他类别的概率则接近于零。

例如,desk chair(办公椅)可能会被误分类为 armchair(扶手椅),但通常不会被误认为是 mushroom(蘑菇)。这种不确定性被称为「暗知识」。

理解蒸馏的另一种方式是,它阻止模型对预测结果过于自信(类似于标签平滑)。

以下是一个实践示例。在语言建模过程中,我们通过观察词汇分布,可以轻松发现这种不确定性。下图展示了 Bert 对电影《卡萨布兰卡》中某句著名台词的 top 20 补全结果:

BERT_base 对被遮蔽 token 的 top 20 补全结果。语言模型确定了两个概率较高的 token(day 和 life)。

如何复制暗知识?

在师生训练中,我们训练学生网络来模拟教师网络的完整输出分布(它的知识)。

我们使学生网络和教师网络具备同样的输出分布,从而使学生网络实现同样的泛化。

我们不对硬目标类别(正确类别的 one-hot 编码)使用交叉熵来进行训练,而是对软目标类别(教师网络的概率)执行交叉熵,从而将教师网络的知识迁移到学生网络。这样训练损失函数就变成了:

其中 t 表示教师网络的 logit 值,s 表示学生网络的 logit 值。该损失函数具备更丰富的训练信号,因为软目标类别比单个硬目标类别提供更多约束。

为了进一步揭示类别分布的多样性,Hinton 等人提出了 softmax-temperature:

其中T 表示温度参数,当 T → 0 时,分布接近于 one-hot 目标向量,当 T →+∞ 时,则得到均匀分布。在训练过程中对教师网络和学生网络使用同样的温度参数,进而为每一个训练样本提供更多信号。在推断时,T 被设置为 1,恢复标准的 Softmax 函数。

PyTorch 动手实践:压缩 BERT

Hugging Face 研究者想利用知识蒸馏压缩大型语言模型。对于蒸馏,研究者使用 KL 散度作为损失函数,因为最优化过程与交叉熵是等价的:

在计算 q(学生网络的分布)的梯度时,得到了同样的梯度。这允许研究者利用 PyTorch 实现执行更快速的计算:

使用教师网络 BERT 的监督信号,研究者训练得到较小的语言模型——DistilBERT。(研究者使用的是 Bert 的英语 bert-base-uncased 版本)。

按照 Hinton 等人的方法,训练损失是蒸馏损失和遮蔽语言建模损失的线性组合。学生模型是 BERT 的较小版本,研究者移除了 token 类型的嵌入和 pooler(用于下一句分类任务),保留了 BERT 的其余架构,不过网络层数只有原版的 1/2。

备注 1:为什么不减少隐藏层大小呢?将它从 768 减少到 512 即可将参数总量减少约 1/2。但是,在现代框架中,大部分运算是经过高度优化的,张量最后一维(隐藏维度)的变化对 Transformer 架构中使用的大部分运算影响较小。在研究者的实验中,相比隐藏层大小,层数才是推断阶段的决定性因素。

研究者的早期实验表明,在该案例中,交叉熵损失会带来更好的性能。因此,他们假设在语言建模设置中,输出空间(词汇)要比下游任务输出空间的维度大得多。而在 L2 损失中,logit 可能会相互抵消。

训练子网络不仅仅关乎架构,它还需要找出子网络收敛的合适初始化(例如彩票假设论文《The Lottery Ticket Hypothesis: Finding Sparse, Trainable Neural Networks》)。因此,研究者基于教师网络 Bert 对学生网络 DistilBERT 进行初始化,将层数减半,对学生网络使用与教师网络一样的隐藏层大小。

研究者还使用 RoBERTa 论文中的一些训练 trick,这篇论文证明 Bert 的训练方式对最终性能有着重要影响。遵循 RoBERTa 的训练方式,研究者利用梯度累积以非常大的批次(每个批次多达 4000 个样本)训练 DistilBERT,训练使用了动态遮挡(dynamic masking),并移除了下一句预测目标。

该训练设置主动接受资源方面的限制。研究者在 8 块 16GB V100 GPU 上训练 DistilBERT,训练时长接近三天半,训练数据为 Toronto Book Corpus 和英文维基百科(与原版 BERT 的训练数据相同)。

DistilBert 的部分代码来自于 Facebook XLM 的代码,部分代码来自 Google AI BERT 的 Hugging Face PyTorch 版本。这些代码可在 Hugging Face 的 NLP 库中获取,该库还包含多个 DistilBert 训练版本和微调版本,及其复现代码。

模型性能:DistilBERT 测试

研究者在 GLUE 基准的开发集上对比了 DistilBERT 和两个基线模型的性能,基线模型分别是 BERT base(DistilBERT 的教师模型)和来自纽约大学的强大非 transformer 基线模型:ELMo + BiLSTMs。研究者使用纽约大学发布的 ELMo 基线 jiant 库和 BERT 基线模型的 PyTorch-Transformers 版本。

如下表所示,DistilBERT 在参数量分别是基线模型的 1/2 和 1/3 的情况下,性能可与后者媲美。在 9 项任务中,DistilBERT 的性能通常等同于或优于 ELMo 基线(在 QNLI 任务上的准确率超出后者 14 个百分点)。令人惊讶的是,DistilBERT 的性能堪比 BERT:在参数量比 BERT 少 40% 的情况下,准确率达到了后者的 95%。

在 GLUE 基准开发集上的模型对比结果。ELMo 的性能结果来自原论文,BERT 和 DistilBERT 的性能结果是使用不同种子进行 5 次运行后的中位数

至于推断性能,DistilBERT 的推断速度比 BERT 快 60%,规模也比后者小;DistilBERT 的推断速度比 ELMo+BiLSTM 快 120%,规模也比后者小很多。

为了进一步调查 DistilBERT 的加速/规模权衡(speed-up/size trade-off),研究者对比了每个模型的参数量和在 STS-B 开发集上使用一块 CPU、批大小为 1 的情况下完成一个完整 pass 的推断时间,如上表所示。

下游任务:蒸馏和迁移学习

研究者进一步研究了在高效推断约束下,DistilBERT 在下游任务上的应用。研究者对紧凑的预训练语言模型 DistilBERT 进行微调,用于分类任务。这是结合蒸馏预训练和迁移学习的绝妙方式!

研究者使用 IMDB 评论情感分类数据集,该数据集包含 5 万条英文评论(被标注为积极或消极):其中 2.5 万条作为训练数据,另外一半作为测试数据(均类别均衡)。研究者使用一块 12GB K80 GPU 进行训练。

首先,在数据集上训练 bert-base-uncased。该模型达到了 99.98% 的准确率(3 次运行的平均值),几乎完美!

然后使用同样的超参数训练 DistilBERT。该模型达到了 99.53% 的准确率(3 次运行的平均值),在延迟降低 60%、规模减少 40% 的情况下,DistilBERT 的性能仅比原版 BERT 低 0.5%!

少即是多:小模型也能出奇迹

Hugging Face 对 DistilBERT 的潜力非常看好。DistilBERT 只是个开始,它也提出了许多问题:使用知识蒸馏技术,我们可以把大模型压缩到什么程度?这些技术可用于进一步探索和洞察大模型中存储的知识吗?在压缩过程中损失了语言学/语义学的哪些方面?……

目前,HuggingFace 的这项研究已经开源,并进行知识共享。他们认为这是每个人参与 NLP 进展,并收获最新进展果实的最快和最公平的路径。

  • GitHub 地址:https://github.com/huggingface

  • Medium 页面:http://www.medium.com/huggingface

入门模型知识蒸馏
相关数据
机器学习技术

机器学习是人工智能的一个分支,是一门多领域交叉学科,涉及概率论、统计学、逼近论、凸分析、计算复杂性理论等多门学科。机器学习理论主要是设计和分析一些让计算机可以自动“学习”的算法。因为学习算法中涉及了大量的统计学理论,机器学习与推断统计学联系尤为密切,也被称为统计学习理论。算法设计方面,机器学习理论关注可以实现的,行之有效的学习算法。

迁移学习技术

迁移学习是一种机器学习方法,就是把为任务 A 开发的模型作为初始点,重新使用在为任务 B 开发模型的过程中。迁移学习是通过从已学习的相关任务中转移知识来改进学习的新任务,虽然大多数机器学习算法都是为了解决单个任务而设计的,但是促进迁移学习的算法的开发是机器学习社区持续关注的话题。 迁移学习对人类来说很常见,例如,我们可能会发现学习识别苹果可能有助于识别梨,或者学习弹奏电子琴可能有助于学习钢琴。

知识蒸馏技术

Hinton 的工作引入了知识蒸馏压缩框架,即通过遵循“学生-教师”的范式减少深度网络的训练量,这种“学生-教师”的范式,即通过软化“教师”的输出而惩罚“学生”。为了完成这一点,学生学要训练以预测教师的输出,即真实的分类标签。这种方法十分简单,但它同样在各种图像分类任务中表现出较好的结果。

语言模型技术

语言模型经常使用在许多自然语言处理方面的应用,如语音识别,机器翻译,词性标注,句法分析和资讯检索。由于字词与句子都是任意组合的长度,因此在训练过的语言模型中会出现未曾出现的字串(资料稀疏的问题),也使得在语料库中估算字串的机率变得很困难,这也是要使用近似的平滑n元语法(N-gram)模型之原因。

推荐文章
暂无评论
暂无评论~