NAACL 2022 | TAMT:通过下游任务无关掩码训练搜索可迁移的BERT子网络

2022 年 6 月 26 日 PaperWeekly


©PaperWeekly 原创 · 作者 | 刘源鑫

单位 | 中科院信工所

研究方向 | 模型压缩




论文标题:

Learning to Win Lottery Tickets in BERT Transfer via Task-agnostic Mask Training

收录会议:

NAACL 2022

论文链接:

https://arxiv.org/abs/2204.11218

代码链接:

https://github.com/llyx97/TAMT




动机

以 BERT 的问世为标志,“预训练+ 下游任务微调”成为了 NLP 任务的一个主流范式。研究趋势显示,持续增长的参数量是预训练语言模型(PLM) 获得成功的一个关键因素。然而,这也增大了 PLM 在下游任务微调和部署的开销。

与此同时,近期关于彩票假设(Lottery Ticket Hypothesis, LTH [1])的研究 [2]发现,基于权重的剪枝(magnitude-based pruning)可以在预训练的 BERT 模型中发现一些 winning ticket 子网络,它们可以迁移到多个下游任务微调,并取得和完整 BERT 相当的效果。

通过初步的实验分析,我们发现这些 winning ticket 子网络良好的下游任务迁移效果,和它们在预训练任务上的效果存在相关性。如图 1 所示,在 50% sparsity 下,oneshot magnitude pruning (OMP) [1]  得到的子网络在预训练 mask language modeling(MLM)任务和下游任务上都明显超过随机子网络。然而,随着 sparsity [2] 持续上升,OMP 子网络在 MLM 任务和下游任务上的效果同时下降。这意味着如果我们能够找到在预训练任务上性能更好的 BERT 子网络,那么它们很可能也有更好的下游任务的迁移能力。


▲ 图1 OMP子网络和随机子网络(Rand)在下游任务(左)和预训练任务(右)的表现




下游任务无关的掩码训练TAMT


▲ 图2 TAMT在预训练任务上(MLM或知识蒸馏)学习子网络结构,然后将其迁移到不同的下游任务进行微调

基于以上动机,我们提出下游任务无关的掩码训练(Task-Agnostic Mask Training,TAMT)方法。如图 2 所示,TAMT 在预训练任务上优化 BERT 子网络的结构(不改变预训练参数值),从而使子网络在预训练任务上有较好的性能。随后搜索到的子网络将被迁移到多种下游任务进行微调训练。
2.1 BERT子网络
对于一个模型 ,我们通过在其参数 上加上二元掩码 得到子网络 。对于 BERT 模型 我们考虑的压缩对象是 Transformer 中的多头自注意力层(Self-Att)和 feed-forward network(FFN)中的参数矩阵,以及词嵌入(word embedding)矩阵。即
2.2 通过二元化和梯度估计进行掩码训练
为了用梯度下降更新二元掩码 ,我们借鉴了训练二元神经网络(binarized neural networks) [3] 的方法。这主要包括前向传播时的二元化,和反向传播时的梯度估计
具体地,在掩码训练过程中,模型中的每个参数矩阵 都会伴随一个二元掩码矩阵   和一个实数值的掩码矩阵 其中 进行二元化得到:



其中 为阈值。
在反向传播时,由于二元化操作不可导,我们用 straight-through estimator [4] 进行梯度估计,从而更新实数值掩码矩阵的参数:



其中 分别为学习率和损失函数。


2.3 子网络结构初始化
由于 OMP 子网络本身具有较好的性能(见图 1)并且 OMP 的计算开销很小,我们用 OMP 子网络来初始化后续的掩码训练。
2.4 TAMT训练目标
子网络结构的学习可以采用任意的下游任务无关(task-agnostic)的目标函数。在本文中我们探究两种训练目标:预训练 BERT 时使用的 MLM,以及基于隐层表示的知识蒸馏 [5] 。具体地,在知识蒸馏中,我们优化学生模型(BERT 子网络)和教师模型(完整 BERT)隐层表示的 cosine 距离:


根据训练目标的不同,TAMT有三种不同的形式。TAMT-MLM 采用 MLM 为训练目标,TAMT-KD 采用隐层表示蒸馏为训练目标,TAMT-MLM+KD 同时采用两种训练目标。




实验

3.1 实验设置

模型:在主实验中,我们选用 BERT-base 和 RoBERTa-base 两种预训练语言模型。在其余实验分析中,我们主要以 BERT-base 为研究对象。

基线方法:我们主要和基于权重剪枝的方法对比,包括 OMP 和在 MLM 任务上进行的迭代权重剪枝(Iterative Magnitude Pruning, IMP)。

数据集: 在子网络搜索阶段(IMP 和 TAMT),我们采用 WikiText-103 数据集。对于下游任务,我们选用自然语言理解 GLUE benchmark 中的 6 个数据集 CoLA, SST-2, RTE, MNLI, MRPC,STS-B,以及 SQuAD v1.1 问答数据集。
3.2 主实验结果


▲ 图3 BERT及其子网络的下游任务表现


▲ 图4 RoBERTa及其子网络的下游任务表现

图 3 和图 4 展示了不同 sparsity 下,不同方法搜索到的 BERT 和 RoBERTa 子网络迁移到下游任务微调的性能。我们发现:
  • 权重剪枝和 TAMT 都明显优于随机子网络,在较低 sparsity 下可以保留完整模型的大部分性能。
  • 采用不同训练目标的三种 TAMT 方法都优于 IMP 和 OMP,在较高的 sparsity 下可以找到迁移能力更强的 BERT/RoBERTa 子网络。
  • 不同的 TAMT 训练目标对于 BERT 子网络影响不大,而对于 RoBERTa,TAMT-KD+MLM > TAMT-MLM > TAMT-KD。这说明 MLM 和 KD 两种目标具有一定的一致性,或者至少不是冲突的。


3.3 预训练效果和下游任务效果的关系


▲  图5 预训练任务(MLM及KD)dev loss和下游任务平均性能的关系

为了验证 TAMT 子网络下游任务性能的提升是否真的来源于预训练任务性能的提升(我们的动机),我们计算了 TAMT 过程中子网络在相应任务上的 dev loss,并且将之和下游任务性能联系起来。如图 5 所示,我们发现:

  • TAMT-MLM 和 TAMT-KD 分别具有最低的 MLM 和 KD loss,说明在 TAMT 过程中,子网络结构的确按照我们设定的目标被优化。
  • 同时,子网络在下游任务上的性能也随着 TAMT 对其结构的优化而逐渐提升。这说明了 BERT 子网络的预训练任务性能和下游任务(微调后)性能之间的确存在相关性,也证明了我们动机的合理性。

3.4 子网络搜索开销和下游任务效果的关系    


▲ 图6 不同预训练步数(左)和时间(右)下搜索到的子网络(70% sparsity)的下游任务平均性能


IMP 和 TAMT 都需要通过一定的预训练来搜索子网络结构。前者交替进行参数训练和剪枝,而后者直接优化子网络结构。图 6 对比了二者的子网络搜索效率。可以看出,随着预训练开销的增加,TAMT 和 IMP 子网络的下游任务微调性能都逐渐提升。相比之下,TAMT 的提升速度更快,在下游任务性能相同(IMP 的最高性能)的情况下,TAMT 的训练开销要少 8 倍多。

3.5 子网络结构初始化的影响

   

▲ 图7 OMP初始化和随机初始化的对比

在以上实验中,我们默认采用 OMP 子网络作为 TAMT 的初始化。为了验证这么做的必要性,我们将 OMP 初始化和随机初始化进行对比。如图 7 所示,当 TAMT 以随机子网络为初始化时,随着训练进行下游任务迁移效果仍然可以稳定提升。但是相比于 OMP 初始化,随机初始化训练收敛后的最终结果要差很多。

更多实验结果和分析可以参见我们的论文。




总结及未来方向

在本文中,我们关注搜索可迁移的 BERT 子网络的问题。我们发现 BERT 子网络的下游任务迁移效果和其在预训练任务上的效果呈正相关。受此启发,我们提出利用下游任务无关掩码训练(TAMT)的方法搜索可迁移的 BERT 子网络。相比于启发式的权重剪枝方法,利用 MLM 或知识蒸馏作为训练目标的 TAMT 可以搜索到下游任务效果更好的 BERT 子网络,这在现有工作的基础上揭示了 BERT 子网络具有更强的下游任务迁移能力。同时 TAMT 在子网络搜索过程中也比迭代权重剪枝更加高效。

在 TAMT 的框架下,仍然有几个方向值得继续改进和探究:
  • 在本文中我们关注非结构化子网络,这需要特殊的硬件支持才能实现实际部署中的运算加速。因此在未来工作中可以探究用 TAMT 搜索结构化 BERT 子网络的效果。

  • TAMT 以及其他方法搜索到的子网络在 90% sparsity 下的效果都还不理想,因此另一个有意义的未来研究方向是探究如何改进 TAMT,从而在更高 sparsity 下搜索到可迁移的 BERT 子网络。

  • 为了实现这个目标,一个可行的方向是探究其他下游任务无关训练目标。

参考文献

[1] Jonathan Frankle and Michael Carbin. The lottery ticket hypothesis: Finding sparse, trainable neural networks. In ICLR 2019.
[2] Tianlong Chen, Jonathan Frankle, Shiyu Chang, Sijia Liu, Yang Zhang, Zhangyang Wang, and Michael Carbin. The lottery ticket hypothesis for pretrained BERT networks. In NeurIPS 2020.
[3] Itay Hubara, Matthieu Courbariaux, Daniel Soudry, Ran El-Yaniv, and Yoshua Bengio. Binarized neural networks. In NIPS 2016.
[4] Yoshua Bengio, Nicholas Léonard, and Aaron C. Courville. 2013. Estimating or propagating gradients through stochastic neurons for conditional computation. CoRR, abs/1308.3432.
[5] Siqi Sun, Yu Cheng, Zhe Gan, and Jingjing Liu. Patient knowledge distillation for BERT model compression. In EMNLP/IJCNLP 2019.


更多阅读




#投 稿 通 道#

 让你的文字被更多人看到 



如何才能让更多的优质内容以更短路径到达读者群体,缩短读者寻找优质内容的成本呢?答案就是:你不认识的人。


总有一些你不认识的人,知道你想知道的东西。PaperWeekly 或许可以成为一座桥梁,促使不同背景、不同方向的学者和学术灵感相互碰撞,迸发出更多的可能性。 


PaperWeekly 鼓励高校实验室或个人,在我们的平台上分享各类优质内容,可以是最新论文解读,也可以是学术热点剖析科研心得竞赛经验讲解等。我们的目的只有一个,让知识真正流动起来。


📝 稿件基本要求:

• 文章确系个人原创作品,未曾在公开渠道发表,如为其他平台已发表或待发表的文章,请明确标注 

• 稿件建议以 markdown 格式撰写,文中配图以附件形式发送,要求图片清晰,无版权问题

• PaperWeekly 尊重原作者署名权,并将为每篇被采纳的原创首发稿件,提供业内具有竞争力稿酬,具体依据文章阅读量和文章质量阶梯制结算


📬 投稿通道:

• 投稿邮箱:hr@paperweekly.site 

• 来稿请备注即时联系方式(微信),以便我们在稿件选用的第一时间联系作者

• 您也可以直接添加小编微信(pwbot02)快速投稿,备注:姓名-投稿


△长按添加PaperWeekly小编




🔍


现在,在「知乎」也能找到我们了

进入知乎首页搜索「PaperWeekly」

点击「关注」订阅我们的专栏吧



·

登录查看更多
0

相关内容

【ECCV2022】UniNet:具有卷积、Transformer和MLP的统一架构搜索
【CVPR2022】基于知识蒸馏的高效预训练
专知会员服务
31+阅读 · 2022年4月23日
【NeurIPS2021】由相似性迁移的弱样本细粒度分类
专知会员服务
20+阅读 · 2021年11月11日
人大最新《基于Transformer 的视频语言预训练》综述论文
专知会员服务
45+阅读 · 2021年9月27日
专知会员服务
18+阅读 · 2021年9月13日
专知会员服务
20+阅读 · 2021年8月17日
【ICML2020】统一预训练伪掩码语言模型
专知会员服务
25+阅读 · 2020年7月23日
【Amazon】使用预先训练的Transformer模型进行数据增强
专知会员服务
56+阅读 · 2020年3月6日
多项NLP任务新SOTA,Facebook提出预训练模型BART
机器之心
22+阅读 · 2019年11月4日
基于知识蒸馏的BERT模型压缩
大数据文摘
18+阅读 · 2019年10月14日
ACL 2019 | 多语言BERT的语言表征探索
AI科技评论
20+阅读 · 2019年9月6日
中文版-BERT-预训练的深度双向Transformer语言模型-详细介绍
3分钟看懂史上最强NLP模型BERT
新智元
23+阅读 · 2019年2月27日
国家自然科学基金
0+阅读 · 2014年12月31日
国家自然科学基金
0+阅读 · 2013年12月31日
国家自然科学基金
4+阅读 · 2013年12月31日
国家自然科学基金
0+阅读 · 2012年12月31日
国家自然科学基金
0+阅读 · 2012年12月31日
国家自然科学基金
0+阅读 · 2011年12月31日
国家自然科学基金
0+阅读 · 2011年12月31日
国家自然科学基金
0+阅读 · 2011年12月31日
Arxiv
0+阅读 · 2022年9月3日
Arxiv
69+阅读 · 2022年6月13日
Arxiv
19+阅读 · 2021年4月8日
Arxiv
15+阅读 · 2020年2月5日
Arxiv
20+阅读 · 2019年9月7日
VIP会员
相关VIP内容
【ECCV2022】UniNet:具有卷积、Transformer和MLP的统一架构搜索
【CVPR2022】基于知识蒸馏的高效预训练
专知会员服务
31+阅读 · 2022年4月23日
【NeurIPS2021】由相似性迁移的弱样本细粒度分类
专知会员服务
20+阅读 · 2021年11月11日
人大最新《基于Transformer 的视频语言预训练》综述论文
专知会员服务
45+阅读 · 2021年9月27日
专知会员服务
18+阅读 · 2021年9月13日
专知会员服务
20+阅读 · 2021年8月17日
【ICML2020】统一预训练伪掩码语言模型
专知会员服务
25+阅读 · 2020年7月23日
【Amazon】使用预先训练的Transformer模型进行数据增强
专知会员服务
56+阅读 · 2020年3月6日
相关资讯
多项NLP任务新SOTA,Facebook提出预训练模型BART
机器之心
22+阅读 · 2019年11月4日
基于知识蒸馏的BERT模型压缩
大数据文摘
18+阅读 · 2019年10月14日
ACL 2019 | 多语言BERT的语言表征探索
AI科技评论
20+阅读 · 2019年9月6日
中文版-BERT-预训练的深度双向Transformer语言模型-详细介绍
3分钟看懂史上最强NLP模型BERT
新智元
23+阅读 · 2019年2月27日
相关基金
国家自然科学基金
0+阅读 · 2014年12月31日
国家自然科学基金
0+阅读 · 2013年12月31日
国家自然科学基金
4+阅读 · 2013年12月31日
国家自然科学基金
0+阅读 · 2012年12月31日
国家自然科学基金
0+阅读 · 2012年12月31日
国家自然科学基金
0+阅读 · 2011年12月31日
国家自然科学基金
0+阅读 · 2011年12月31日
国家自然科学基金
0+阅读 · 2011年12月31日
相关论文
Arxiv
0+阅读 · 2022年9月3日
Arxiv
69+阅读 · 2022年6月13日
Arxiv
19+阅读 · 2021年4月8日
Arxiv
15+阅读 · 2020年2月5日
Arxiv
20+阅读 · 2019年9月7日
Top
微信扫码咨询专知VIP会员