Press "Enter" to skip to content

AAAI 2022 | 北大 & 阿里达摩院:基于对比学习的预训练语言模型剪枝压缩

近年来,预训练语言模型迅速发展,模型参数量也不断增加。为了提高模型效率,各种各样的模型压缩方法被提出,其中就包括模型剪枝。

 

然而,现有的模型剪枝方法大多只聚焦于保留任务相关知识,而忽略了任务无关的通用知识的保留。

 

针对这个问题,北京大学与阿里达摩院合作,在今年 AAAI 2022 中的一项工作中工作提出了一种 基于对比学习的新的剪枝框架 ,适用于各种结构化剪枝与非结构化剪枝方法,在各种压缩比下都提升了压缩后模型的表现。

 

论文标题:

 

From Dense to Sparse: Contrastive Pruning for Better Pre-trained Language Model Compression

 

论文链接:

 

https:// arxiv.org/abs/2112.0719 8

 

代码链接:

 

https:// github.com/RunxinXu/Con trastivePruning

 

本期 AI Drive,由北京大学计算语言所二年级硕士生许润昕详解这一项工作。主要内容分别是:

背景:预训练模型与模型压缩(简单介绍现在的预训练模型的发展,以及预训练模型压缩的常用方法)
预训练模型剪枝存在的问题与解决方法(介绍预训练模型,特别是剪枝方法的现存问题,以及给出针对问题的解决方法)
实验探究和结论

1. 预训练模型和模型压缩

 

首先,预训练模型的发展在各个领域都发挥了很大的作用。

 

随着时间的推移,越来越多不同的预训练模型被提出:从 ELMo 到 GPT,再到 2018 年的 BERT,之后,针对不同的目的,又有了衍生出各种各样的预训练模型,比如多语言、多模态等,从 GPT 到 GPT-2、GPT-3 实现基本上 zero-shot 的做法,越来越多、各种各样的预训练模型都被提出来,并发挥了很大的作用。

 

使用预训练的基本范式,包括预训练与微调两个步骤。

 

预训练模型目前流行的使用范式如下:先假设给出初始化模型,再给出自监督的训练目标,比如 mask language model,然后在无监督无标注的大规模语料里做预训练,之后把预训练模型迁移到下游的、有标注的具体任务的数据,在其数据基础上进行微调,如此一来便可以达到比较好的效果。

 

其中的典型代表是 BERT 模型。BERT 堆叠了一系列 Transformer encoder 的 blocks。在预训练阶段中, 通过Masked Language Model(MLM任务),以及 Next Sentence Prediction(NSP 任务),去自监督地预训练模型。

 

在微调阶段中,可以根据下游不同任务制定不同的 Fine-tuned 策略。比如单句分类、句子分类或情感分类。直接输入句子,然后用模型进行编码,用 [cls] 进行预测。

 

而对预训练模型进行模型压缩的原因则在于,预训练模型在各个领域都发挥了巨大作用,因此模型越做越大,其参数量越来越大、计算成本越来越高。比如从 BERT base 到 RoBERTa base,从 GPT-2 到 Turing-NLG 等,还有为众人熟悉的 GPT-3,参数量同样越来越大。指数级的增长见下图。

 

虽然大模型效果好,但是计算效率比较低,因此越来越多的人希望能够给预训练模型瘦身。

 

模型瘦身的必要性有以下几点:

因为模型越来越大,所以希望缩小模型给预训练模型来进行瘦身,如此一来可以加快推理计算速度,满足最低延迟限制;
模型瘦身可减少所消耗的内存,对于端侧更易于部署,比如手机;
模型瘦身便于训练/微调。

那幺,为什幺选择压缩大的模型而不直接做小模型?

 

2020 年 ICML 的一篇文章针对此问题做了研究探讨,其结论表明,适应一个大模型后将其压缩为小模型,这样模型的表现要优于直接训练出小模型,并且,大模型具有敏感度更低、更容易压缩的优点。

 

下图为 paper 里的实验图,其横轴表示参数量,纵轴表示在 mnli 这项任务中的 accuracy,不同的线表示的是不同模型结构各自的表现。实验图上面还有一些数字,例如 0% 、15%、30%,表示的是这个模型压缩的比例,0% 为原始模型,15% 为压缩 15% 的参数之后模型的表现。最后由下图数据分析得出结论:压缩大模型后所得到的小模型,效果比直接训练一个小模型好。

 

模型瘦身方案也有许多种,比如蒸馏、量化、剪枝、矩阵分解等,下文主要介绍知识蒸馏和剪枝。

 

方法一:知识蒸馏

 

将大模型的知识蒸馏到小模型中,使得较小的模型能够达到类似于大模型的效果,从而达到瘦身的目的。

 

如下左图,teacher 对应 big model,student 对应 smallmodel,teacher 先去学习书中的知识转移为 teacher 所得,经 teacher 消化后将其精华教给 student,知识迁移为学生所得,使得学生所学可以达到类似于老师所获得的效果,从而达到知识蒸馏的目的。

 

如下右图,对于模型实现,即输入数据时,数据同时经过 teacher 的 model 和 student 的 model,student 的 model 既可以从 hard label 运行,也可以跟 teacher model 学,但是其学的方法是在温度系数的调整下做 softmax,最后去对齐 model 预测的标签分布。比如 teacher 的 model 预测是正类 0.7,负类是 0.3,那幺需要 student model 预测出的 label 的分布也与其数据接近,从而达到约束 kl 闪度的效果。如此一来则是 student 通过 soft label 的方式将 teacher 所学知识迁移到其自身中。

 

扩展:TinyBERT

 

在上图模型中,只能通过最上面的 softlabel 进行知识迁移,而在下图的整个模型中,从第一层、第二层、第三层到最上面一层都可以做知识迁移,所以 Tinybert提出:从 embedding的时候,对其teacher和student的 embedding,而中间每个 transformer 的 block,都可以去对齐 hidden states,也可以对齐 teacher 的 hidden stage 和 student 的 hidden state;以及attention的 magic,同样可以去要求 student 计算出来 attention 的举证,尽量去跟 teacher 的 attention 举证是尽量相似的,如此就相当于 student 全方位模仿 teacher,并且更好的吸收 teacher 的知识。

 

拓展:MiniLM

 

Minilm 与 tiny bert 类似,它去蒸馏了最上面的 transformer 的 block,再通过一些 attention,去要求 student attention 尽量的去跟 teacher 的 attention相似。一方面,普通的 attention 是由 query 和 keys 计算出来的分数,之后可以约束它;然后另一方面, minilm 也加入了一个是 value 乘以 value 的操作,即用每一个位置的 value 去当 query,用每个位置的 value 去当 keys 同样可以计算出来一个类似于 attention 的矩阵;同样,去对齐 student 跟 teacher 的 attention 的矩阵,这样的话可以更好的去做蒸馏。

 

方法二:剪枝

 

剪枝的思路与知识蒸馏的思路本质上差不多,相同之处是目的都是压缩模型,而不同之处则在于剪枝的实现方式更加灵活。

 

剪枝的意思是将模型里面里面不重要的、多余的参数移除。比如下图,图中是一个全连接的网络,有一些参数连接该网络对于整个模型来说是可有可无的,那幺此时就可以对这些多余部分进行剪枝。

 

剪枝的经典方法分为两类,即结构化剪枝和非结构化剪枝。

 

1. 结构化剪枝: 一个参数相当于一个矩阵,如果一个矩阵有一定的空间规律,比如每一次移除参数都是移除某一行、某几行或者某几列、某一个 block,像这样有规律的剪枝就是结构化剪枝。

 

2. 非结构化剪枝: 如果参数是以神经元形式呈零散分布,从整体的矩阵来说没有表现出有规律的分布,剪枝时将重要的参数保留,非重要参数则直接除去,如下图的灰色圆点表示的就是被裁剪掉的参数,裁剪时并不需要考虑整个空间的结构性,像这样没有规律性的剪枝就是非结构化剪枝。

 

在剪枝工作中,有两个很重要的问题:

 

第一,什幺时候剪枝;第二,用什幺标准剪枝。

 

对于第一个问题,判断模型从什幺时候开始剪枝,方法可以分为很多种,具体最主要有三种:

 

1. One–shot pruning,即首先去训练一个模型,等模型收敛之后去识别不重要的参数并将其去除,最后再将其模型训练一定时间后让它收敛,这是最简单的方法。

 

2. Iteration Pruning,此方法是一种最常见的方法即迭代式的 pruning situation 的 pruning,具体做法就是训练一个模型,然后去掉一些参数,后面再不断地重复前面的步骤,直到达到目的。正如下图所示,迭代式的发展就是迭代式的剪枝,而对模型的每一次剪枝,都会带来效果的降低,但是在后续的训练过程中,performance 又会随之上升,所以呈现出图中的阶梯式分布。

 

3. Spars training,此方法的具体呈现是一开始便吸收一个网络进行训练,中途会进行 regrowth,然后进行剪枝,把不重要的参数去除,形成一个由少变多再变少的过程。但是这对于预训练模型的剪枝来说并不是很适用,因为预训练模型是从一开始就拿到了一个很大的模型,要进行的是对大模型密集的网络进行剪枝,所以对于预训练模型更适用的方法是迭代式的剪枝方法,所以,最常见的就是迭代式的剪枝方法。

 

对于第二个问题,此处为大家介绍两种比较经典的剪枝方法。

 

第一种方法是基于参数对模型的贡献。

 

它的直观思路是如果将一个模型的参数置 0,此时 loss 变化很大,说明该参数对于模型来说很重要,而模型对于该参数则是很敏感的;但如果将一个模型的参数置 0 后,loss 的变化很微弱或者基本上不变,那幺就说明这个参数是可有可无的;而对某个或某部分参数hi进行具体计算,将 hi 置顶和没有置顶之间的 lost 的绝对值进行展开,可以推出一个式子,即该参数乘以参数本身以及该参数的梯度的绝对值,此式子可作为衡量该参数的重要性的指标。此处便不再详细展开。

 

第二种方法是,基于参数在训练过程中的变化方向以确定参数的重要性。

 

它的直观思路是如果参数在训练过程中倾向于往 0 移动,那幺该参数可能是不太重要的;反之,如果参数在训练过程中倾向于远离 0,那幺该参数可能是比较重要的。用下图表征,对于参数 Wij,它远离 0 有两种情况,第一种情况是参数本身大于零,梯度小于零,而梯度小于零就意味着参数要往大于零的方向移动,如此便是远离零;如果参数小于零,而该参数的梯度大于零,那幺就意味着该参数要往负方向移动,如此也是远离零。这就是参数远离零的两种情况。

 

以上提到的可以用下图的式子表征参数的重要性,因为式子是参数本身乘以参数本身的梯度,这两个是不同的符号,相乘得负数,而且越小出名物的重要程度越高,所以需要在式子前加一个负号,这样就有 S 可以来表征参数重要性的分数。而 s 是一个可以调整的超参数,由此可以看到下图中出现 y 等于 x 的线,而黄色部分是保留参数的部分,即远离零的参数都被保留而靠近零的参数基本都被移除了。

 

基于 movementpruning,还有一个加强版本叫做 soft movement pruning。

 

上文提到的 movement 的方法就是计算出参数的重要性分数,再人为的去指定要保留前百分之七十或者百分之六十;而 soft movement pruning,并不采用保留 top k 的方式,而是直接采用计算出的分数去取一个阈值,然后把小于该阈值的参数尽数除去。除此之外,我们还需要在 loss 里面加入一个正则项对分数进行约束,将分数约束得越小越好,如此一来便可以裁剪去比较多得参数。上文 Lambda MVP 的意思其实是一个正则项的系数,超参越大说明该分数的约束的正则化越重,那幺会有越多的分数比较低,最后得到的模型稀疏度就会越高。

 

以上便是模型压缩的经典方法,接下来,介绍预训练模型剪枝存在的问题与解决方法。

 

2. 预训练模型剪枝存在的问题与解决方法

 

上文提到了压缩预训练模型的剪枝,它的流程可如下所述:首先拿到一个预训练模型,该模型的 sparsity ratio 为 0%,相当于没有剪枝,然后该模型通过迭代式的剪枝一步一步进行一定的裁剪,比如裁剪到 10%、20%、30% 或者 70% 等。

 

每一次裁剪的过程中,在计算模型优化的 loss 以及计算模型参数的重要性分数时,都是依靠下游的数据,并且涉及到用下游的 task、specific data 来对模型求导,然后计算出参数的重要性以及模型的优化损失。

 

而运用这样的方式在整个压缩过程中,该模型只考虑到了任务相关的知识,而忽略了与任务无关的通用知识的保留。比如下图中一开始的 Pretrained Model,它也包含了一些较为广泛的知识,这些也是很重要的,如果在整个剪枝压缩过程中也可以保留这部分知识,那幺这部分知识是有助于保留提高剪枝后的模型的一些淡化性能,以及防止一些灾难性遗忘问题的。

 

此外,在迭代式的裁剪的过程中,会产生一系列的模型,比如下图中的分别剪了 10%、20%、30% 的模型。通常来说,之前的一些方法只是简单地把这一系列模型抛弃,但是实际上产生地这一系列各个不同系数比例地模型也可以指导现在的剪枝模型的一些知识的学习,因为这一系列模型蕴含着一些中间知识,这一系列模型叫做 snapshots,在剪枝过程中可以从 snapshots 学到相应知识的话有助于保留之前产生的一些知识。相应的,我们还可以从 Fine-tuned Model 学习,这个是与任务更加相关的一个知识。

 

总的来说,目前预训练模型剪枝存在的问题,是在压缩过程中只考虑了任务相关的知识,而忽略了与任务无关的通用知识的保留,以及迭代过程中产生的中间知识。

 

从上述问题入手,对于当前的模型剪枝,这项研究最直接的 motivation 就是希望在剪枝过程中可以学习到一些与任务无关的通用知识,比如说 PretrainedModel、Snapshots Model 以及 Fine-tunedModel 之间的知识,即在剪枝压缩过程中要保留各个方面的知识,这就是核心的 motivation。

 

如下图所示,右边部分图是模型的架构图。

 

该架构图的直观思路:从 InputTensor 开始,首先是最上部的 Pre-trained Model,该模型储存与任务无关的通用知识;接着是底部 Fine-trained Model,Fine-tuned Model 主要是储存与任务相关的知识;而中间则是产生的一系列 Snapshots Model,其下标表示的目前的系数比列,比如裁到 20%、40% 或者 60%。对于裁剪了 60% 参数的模型,绿色的箭头表示它可以学习 Pre-trained Model 的知识,蓝色的线则对应学习 Fine-tuned Model 的知识,以及黄色的线则对应学习中间模型的知识。这就是该架构图的直观思路,它们本质上分别对应了图中所提出的三个不同模块——PrC、SnC 以及 FiC,而这三个模块则分别负责了对不同的目标模型的学习。

 

那幺,如何保留各个部分的知识?

 

从模型对于数据的编码表征能力入手。比如希望当前模型的编码表征能力可以尽可能地向目标模型看齐,但并不需要一模一样,只是达到类似的效果、学到相应的保留目标模型的编码能力即可。

 

这个场景就适用于运用对比学习的方式去学习目标模型的编码表征能力。

 

因为对比学习是一种非常有效的学习方法。以下简单介绍对比学习。

 

比如对两张钞票进行真假辨别,我们并不需要记住真钞票的所有细节,只需要在大脑中有一个模糊的美元的概念或者抽象的印象,然后去与下图中的两张钞票进行对比,会发现右边的钞票与大脑中的印象更加贴合,所以可以得知右边的钞票是真的钞票,这就是对比学习的一个直观理解。

 

对于对比学习,简单可以理解为一个作为 anchor 的 instance,比如上文提到的头脑中的模糊印象;一堆作为候选集的 instances,比如上文所提到的两张钞票,两张钞票即为候选集,将它们编码以后,与向量空间中的 anchor 进行查询比较,在向量空间中的 anchor 距离近的 instances,就是与 anchor 更像的 instances。

 

例如下图中的动物们,左上图的狗狗作为 Anchor,将 Anchor 编码到向量空间中;而左下的图片是一只与 Anchor 相似的狗狗,那幺这只狗狗就相当于 Anchor 的正类即 Positive,对于正类是要尽可能地与 Anchor 靠近;而对于右边的动物比如猫猫、大象以及其他与 Anchor 并不相似的动物们,则是作为负类,那幺负类则是要尽量远离 Anchor。

 

由上述可知,通过对比学习可以学到编码器的更好的编码表征能力。

 

对比学习的核心思想就是拉近正例、扯远负例。

 

结合上文所提,如下图所示,对比学习的核心思想的具体表征:输入数据 Xi,再用编码器将其数据编码为 Zj,除此之外还有 n 个不同的 candidates,紧接着用 normalize 过的内基除以温度系数去计算它们的相似度并用来表征它的分数,再用 soft max 的方式计算出它和不同样本之间的距离,最后再通过一个对比学习的 loss,从而达到 Anchor 与正例尽量接近而与反例尽量远离的目的。通过这种方式,可以更好地学习编码器的能力。

 

进一步的,如何将对比学习应用于我们的场景?

 

目前需要的是,正在剪枝的模型的表征能力尽量与目标模型的表征能力接近,那幺,首先选用当前的稀疏模型去编码数据作为 Anchor,再选取 n 个不同的数据,用目标模型编码得到 candidates,最后通过要求当前稀疏模型编码的 anchor 与目标模型 candidates中正例距离更近,负例更远,从而拉近当前稀疏模型与目标模型表征能力的 gap。

 

具体以 PrC 模块为例。

 

比如在这个例子中,需要当前的系数模型对预训练模型的表征能力进行学习,如下图中正方形表示由当前的稀疏模型编码的数据,圆形表示预训练模型编码的数据,将其分为两个场景,一个是 Unsupervised,另一个则是 Supervised。

 

在 Unsupervised 中,r 表示当前的稀疏程度,pre 表示 PretrainedModel,对于相同的、数据都是 x1 的模型,我们希望当前稀疏模型编码的向量与预训练编码的向量是尽量接近的,而对于其他的不同数据来说,我们则希望稀疏模型编码出来的向量与其他的由预训练模型编码的向量尽量的疏远,因为它没有用到相应的 Label 的一些信息。在 supervise 中,它与 Unsupervised 一个不同之处是它通过它的 label 来判断是不是正例,如果它们的标签是一样的话,那幺就可以认为它们是相同的正例。

 

如上右图,有由当前稀疏模型编码的和 pretrained model 编码出来的四个不同的数据,x1 与 x4 的 label 都是 yi 等于 1,可以认为它们是正例,那幺就希望它们是尽量靠近的,对于其他的 yi 等于 0 的数据,相当于是不同的标签,那幺就希望它们是尽量疏远的。

 

相应的,比如下图中的 snc 模块,sn 表示的是 snapshots,即意味着去学习中间模型产生的知识,本质上与上文所提是一样的,只是现在编码的这些数据的模型不是 pretrained model 而是换成了中间模型。同样这里仍有 unsupervised 与 supervised 的场景,此处不一一展开。

 

还有 FiC 模块,该模块的主要思路就是向 fine-tuned model 学习,将 candicates 的编码模型换成 fine-tune 的编码模型将数据编码出来后,再去做相应的对比学习。

 

应用对比学习总的来说有三个模块,每个模块负责各自相应的不同的目标模型,它们都有 supervise 和 unsupervised 的用法。

 

简单来说,unsupervised 就是不需要去看它的 label 或者是直接把数据的 id 作为 label,而 supervise 则是要去看具体的 label,如果 label 是相同的那幺就可以当作是正例。而对比学习的核心目的,就是拉近当前的稀疏模型与各个目标模型的表征能力的 gap。

 

如果当前的系数模型涉及到了许多个目标模型,那幺是不是需要存很多的模型到内存里?

 

实际上它的可行性并不强。naïve 想法就是同时 load pre-trained/fine-tuned/ 所有中途产生的 snapshots 模型到内存中,每一次输入一个 batch 再去选取 n 个 batch size 样本,最后把经由这所有的模型进行编码,但其实这样的话在空间和时间上都是不可行的,因此这个想法并没有可行性。

 

实际上在实际应用时,并不需要保存模型,只需要关心编码后的向量,因为可以在迭代式的训练裁剪过程中将其 pre-encode 并保存,需要的时候 fetch 相应向量到 GPU 即可。

 

比如,对于 BERT base,其参数量为 110 M,n 表示对比学习中 candidates 的样本数量,如下图所示,每一次 fetch 进来的向量只占总参数量的 3% 不到,所以这是接受的消耗。

 

CAP 更多的是一个框架而不是一个具体的剪枝策略,因此它可以与不同的剪枝方法结合起来使用。下文介绍将 CAP 框架应用于不同的剪枝方法中。

 

将 CAP 应用于前面说过的 First-order Taylor Pruning 方法,则得到了 CAP-f 方法;同样地,将 CAP 应用于 movement pruning,以及 soft movement pruning,则得到了 CAP-m 以及 CAP-soft 方法。

 

3. 实践探究与结论

 

实验基于 BERT base 模型以及 MNLI、QQP、SST-2、SQuADv1.0 数据集展开。

 

实验结论一是,通过 CAP 框架可以移除模型大部分参数,但依然取得可比的效果。

 

比如,下图中第一行中蓝框中的数据为未裁剪的 BERT base 的基本模型,而下一个蓝框部分就是做了 50% 的裁剪后的数据,两组数据对比可以看出,做了 50% 的裁剪的 BERT,与原始的 BERT 的表现类似甚至略有提高;再比如另一组数据,裁了 97% 的参数之后,在 qqp 这一任务上依然保持原始 bert 表现的 97.2%。

 

结论二则是,由于 CAP 可以用于不同的剪枝方法,所以,当 CAP 应用在 first order burning 上时,在不同的裁剪比例上都能带来有效的提升,比如应用于 movement running 的方式上,也可以获得比较明显的提升。

 

结论三是,对比其他的剪枝优化的方法,比如 SNIP 或其他不同的剪枝优化方法,CAP 有着明显的优势。

 

结论四是,CAP 对比蒸馏方法也有优势,但是两者本质上并不是为了打败蒸馏方法,因为二者是可以结合的。

 

在上文提到,我们要学习一些通用知识,但这样能否提升模型的泛化性能呢?接下来将介绍一个验证上述猜想的实验。

 

1. 实验方法: 将得到训练剪裁完毕的模型后将其解冻,迁移到新任务上,在模型之上训练一个简单线性分类器,衡量模型产生表示的泛化能力。

 

2. 实验结果: 如下表格,首先左边表示的是先在 MNLI 上训练裁剪,再迁移到不同的任务,右边则表示的是先在 QQP 上训练裁剪,再迁移到不同任务,图中左侧的 f 和 m 表示的就是基于不同的剪枝方法,f 是 first order 的结构化剪枝方法,m 就是 movement Pruning 非结构化剪枝方法,50%、90% 表示的就是模型的裁剪的比例。从图中可以看到,表格中的蓝色部分表示的是有提升,相比 baseline 使用了 CAP 比不使用 CAP 是有提升的,而且迁移的比例会更大一些。

 

3. 实验结论: CAP 提高了模型的泛化性能。

 

消融实验方面,消融实验就是去除不同的模块都发生了 performance 的下降,所以对比学习的三个模块都发挥不可或缺的作用,以及 supervised 跟 unsupervised 的对比学习目标对于模型都很重要。

 

而在不同的裁剪比例下,CAP 都能给不同的剪枝方法带来相应的提升。

 

简而言之,预训练模型越来越大,给模型瘦身变得越来越重要。瘦身的方法包含多种,但是往往在模型剪枝的过程中许多人仅考虑了任务相关的知识,而忽略了与任务无关的中间知识的保留。

 

所以,我们针对上述问题提出了使用基于对比学习的方法,以及在模型裁剪的过程中向 Pretrained 等一系列的中间模型学习的方法,这些方法可以给不同的剪枝策略带来提升,并且提高了模型的泛化性能(扫描下方二维码即可访问此次研究的论文和代码)。

 

Be First to Comment

发表回复

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