Press "Enter" to skip to content

达摩院李雅亮:大规模预训练模型的压缩和蒸馏

 

作者 | 李雅亮博士

 

阿里巴巴 来源 | DataFunTalk

 

 

 

导读: 本次分享的主题为 大规模预训练模型的压缩和蒸馏 ,主要是从 自动机器学习的角度,介绍大规模预训练模型的压缩和蒸馏。将介绍阿里巴巴达摩院关于模型压缩的三个连续承接性的工作:

 

工作1:AdaBERT:Task-AdaptiveBERT Compression with Differentiable Neural Architecture Search. IJCAI’2020.

 

工作2:L2A:Learning toAugment for Data-Scarce Domain BERT Knowledge Distillation. AAAI’2021.

 

工作3:Meta-KD:MetaKnowledge Distillation Framework for Language Model Compression across Domains.ACL’2021.

 

01

 

背景介绍

 

自2018年BERT预训 练横空出世,NLP领域出现了一个全新的范式,大规模预训练模型给我们带来了很多新的能力。 通俗来讲,我们会先从很多其他语料或者说很多大规模数据里面做一个预训练的模型,然后在具体任务的训练语料微调预训练的模型。 基于“预训练-微调“范式,许多任务的性能都得到了极大的提升。 如图1,文本分类(Text Classification), 命名实体识别(Named Entity Recognition),基于机器阅读理解的问答( Question Answering, MachineReading Comprehension )等任务的性能得到显着提升。 在BERT论文中他们指出BERT在11个经典的自然语言任务中超越当前的state-of-the-art的性能,“一石激起千层浪”,自此之后,大量的预训练模型层出不穷——GPT系列预训练模型,相比BERT炼制更加成熟的RoBERTa, 缓解序列长度的限制XLNET, 融入知识的 K-BERT, 领域预训练模型SCIBERT等等。

 

 

图1  BERT模型在不同任务上微调

 

然而,预训练的大模型,顾名思义,他们的参数量会非常大,如12层transformer堆叠的wwmpytorch版本的预训练模型大小为400M左右,24层transformer的版本大概1.2G左右。但是在现实生活中,我们的计算资源有限的情况下,至少会引起两个问题,其一,我们设备终端无法加载这幺大的模型,其二,参数太多,推断速度太慢。比如说在线搜索这些需要实时响应的应用,或者说我们自己的手机、车载系统,它们的计算资源是非常受限的,所以我们需要有一些小的模型来替代它的功能。

 

所以,为了解决模型上述的痛点,模型的压缩技术在工业界和学术界应运而生。如图2,现在的BERT预训练模型压缩的基本研究思路如下:

 

 

图2 BERT压缩常规研究方法

 

模型蒸馏(distillation):强迫一个小的模型去通过去模仿一个大模型的输入输出,如DistilBERT, TinyBERT, MiniLM等工作。

 

量化(quantization):对整个大模型里面的参数做量化,也就把很多不重要的参数,我们可以用比较少的精度来存储, 如Q-BERT等工作。

 

参数共享(parameter sharing):模型不同的组件(component)之间去共享参数,如ALBERT的工作。

 

02

 

AdaBERT介绍

 

1. 研究动机

 

上述模型压缩的几个基本的研究思路(distillation,quantization, parameter sharing等)都是和下游任务(downstream)不相关的,即模型压缩的结果都是大模型统一压缩成小模型,然后再利用小模型去微调具体任务。

 

但是,在2019年,Ian Tenney等人(如图3)揭示了——BERT学的knowledge非常的多,换一句说,BERT学到的非常多的general knowledge,但是不同下游任务需要的knowledge可能不太一样,比如做NER是你可能更关注的是细粒度的信息,做reading comprehension可能关注的是需要一些high level的信息。受此启发,我们提出AdaBERT:针对一个具体的下游任务,通过自动压缩一个小模型,来适配该任务,即小模型是专门为该任务服务的。

 

 

图3 自适应任务模型压缩的研究启发

 

2. 模型介绍

 

我们利用神经网络结构搜索(NAS)来自动化搜索小模型的网络结构。下面我们从三个方面介绍模型的一些设计。

 

① 如何定义Search Space

 

 

图4 Search space 设计

 

不同于以往BERT的压缩——将Transformer结构压缩成更少参数transformer的结构。在这里我们采用基于CNN结构的architecture。原因如下:

 

首先,很多场景下的优化,基于CNN的硬件支持可能会更好,有利于实际场景中的部署;

 

另外,主要考虑到在AutoML中NAS研究中关于绝大部分研究都会是和CNN相关的,方便借鉴已有的成果,如differentiable, Gumbel trick。

 

②  Search目标

 

Search的目标包含两个方面,一方面,需要小模型尽可能保留BERT中有用的knowledge,另一方面,需要权衡推断效率,即小模型的size。 权衡保留大模型的有用的knowledge和 efficiency,故总体的Loss设计参见图5。 其中L  为具体下游任务的Loss。 , 是超参数。

 

 

图5 权衡knowledge-efficiency损失函数设计

 

 

图6  knowledge distillation Loss

 

 

图7 efficiency-aware Loss

 

③ 如何去搜

 

因为定义的CNN的网络架构是离散化的,从传统的方法来看的话,如采用随机采样的方式去尝试不同的architecture,费时费力,所以我们利用differential方法去优化。其思想是:首先,我们预定义好一些operation, 对于每个网络层,并行组织所有的operation,形成super net,即每个super net中不是存在一个operation,而是存在所有可能operation的一个集合,通过加权形成所有可能的operation,这样我们就可以用differential的方法去优化整个大网络的结构了。训练architecture的parameters使用Gumbel trick。Differential的更多的具体原理细节,可以参考Liu, Simonyan等人提出DARTS。

 

3. 实验效果

 

 

图8 实验结果

 

可以看出AdaBERT的performance基本都处于次优的结果,而参数量最少,inference speedup有明显提升。

 

 

图9 验证Task-AdaBERT的实验

 

如图9,为了验证网络结果符合任务自适应,将训练的任务自适应的小模型应用到其他任务上,完全符合我们的设想,自适应任务的模型在自己的任务上表现最优。

 

03

 

L2A介绍

 

1. 研究动机

 

通过上述介绍,AdaBERT可以work的一个前提假设是我们存在大量的领域训练数据,然而,真实世界往往是缺乏训练数据的;如果仍然按照AdaBERT一套做模型蒸馏,在数据稀少的情况下,会导致学生(小)模型性能退化。

 

对于数据缺乏,NLP领域有大量的data augmentation的方法,如对于真实的文本序列做token 的增,删,改,插或者来回逆向翻译。然而从自动化的角度来看,这些data augmentation的方法需要花费大量时间成本去尝试,所以我们提出一种BERT蒸馏过程中自动学习增强数据缺少领域的方法。

 

2. 模型介绍

 

 

图10 L2A 模型结构

 

如图10,模型的主要的思想基于transfer learning和对抗网络。我们通过对source Domain 和 target Domain的data generator自动产生generated data,形成的training data去给Teacher Model 和 Student Model在压缩的时候去用,通过distillation Loss 和 Evaluation Metric通过Reinforcement learning的方式去迭代修改Generator行为。通过这样的方式去做筛选,把一些有用的信息留下来,相反把一些没用的信息就会就给剔除掉。

 

3. 实验结果

 

我们选取了四个任务,如图11,对于每组target dataset,我们都挑选合理有效的source datase。这里的缺陷即为,对于每个target dataset,需要我们人工挑选合适的source。这个问题我们会在下面Meta-KD缓解它。

 

 

图11  L2A四组任务

 

 

图12 L2A实验评估结果

 

如图12, 我们可以在所有的任务上都得到一个比较显着的一个提升。特别地,一个非常有趣的事情,我们找到的小模型很可能比大模型还要好,如NLI, Text classification这两个任务。

 

 

图13 L2A的消融实验

 

分析消融实验结果,我们可以发现L2A在数据量越小,改善越明显,他与teacher model之间的性能gap越大。这表明L2A可以有效地帮助数据稀缺领域的知识提炼。

 

04

 

Meta-KD介绍

 

1. 研究动机

 

上文提到L2A利用迁移学习的思想,可以在选择的source domain和target domain来训练一个适应任务的小模型,然而,人工选择source domain数据仍然存在一定的时间成本。为此我们提出Mate-KD:通过一个meta-teacher model 学习各种cross-domain的knowledge, student model在meta-teacher的“指导”下学习一个具体domain的knowledge。

 

2. 模型介绍

 

阶段一 :meta-teacher learning。学习一个整合了cross-domainmeta-teacher model;

 

阶段二 :meta-distillation。在 meta-teacher model的指导下,student model学习一个具体的数据集的knowledge, 具体的数据集可能从来没有在source domain出现过。

 

 

图14 Meta-KD 模型结构

 

3. 实验结果

 

我们在两个大的数据集上(MNLI, Amazon)做了实验,每个数据集都有四~五个domain,都有performance都有提升。然后比较有趣的,如图在Mate-KD在无特定领域的前置knowledge情况下,performance表现非常好。

 

 

图15 Meta-KD在MNLI数据集上的实验结果

 

 

图16 Meta-KD在Amazon数据集上的实验结果

 

 

图17 Meta-KD在无domain data情况下的实验结果

 

值得注意的一点是,如图17所示,当Meta-teacher在训练模型的过程中没有引入任何fiction domain的知识,然而实验效果却非常好。

 

05

 

精彩问答

 

Q: 关于小模型的性能比大模型性能更好的意外收获,下一步的计划?

 

A: 我们猜想大模型可能有些knowledge对于具体的task来说是noise,小模性可能去掉这个noise, 故表现更好。 但是如何针对性地让所有小模型比大模型表现更好,我们认为有potential,但是目前还没有展开研究。

 

Q : 模型压缩都是基于分类任务做的,考虑生成式任务上模型压缩吗?

 

A:目前我们大部分涉及到NLP 任务主要是要判断意图,分类,抽取信息,搜索之类,目前还没有涉及到generation 任务。

 

Be First to Comment

发表回复

您的电子邮箱地址不会被公开。 必填项已用*标注