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

论文标题：Learning to Win Lottery Tickets in BERT Transfer via Task-agnostic Mask Training

论文链接：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）在下游任务（左）和预训练任务（右）的表现

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

基于以上动机，我们提出下游任务无关的掩码训练（Task-Agnostic Mask Training，TAMT）方法。如图 2 所示，TAMT 在预训练任务上优化 BERT 子网络的结构（不改变预训练参数值），从而使子网络在预训练任务上有较好的性能。随后搜索到的子网络将被迁移到多种下游任务进行微调训练。