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

Proceedings of the 2022 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pages 5840 - 5857 July 10-15, 2022 ©2022 Association for Computational Linguistics Learning to Win Lottery Tickets in BERT Transfer via Task-agnostic Mask Training Yuanxin Liu1,2, Fandong Meng3, Zheng Lin1,2∗, Peng Fu1∗, Yanan Cao1,2, Weiping Wang1, Jie Zhou3 1Institute of Information Engineering, Chinese Academy of Sciences, Beijing, China 2School of Cyber Security, University of Chinese Academy of Sciences, Beijing, China 3Pattern Recognition Center, WeChat AI, Tencent Inc, China {liuyuanxin,linzheng,fupeng,caoyanan,wangweiping}@iie.ac.cn, {fandongmeng,withtomzhou}@tencent.com Abstract Recent studies on the lottery ticket hypothesis (LTH) show that pre-trained language models (PLMs) like BERT contain matching subnet- works that have similar transfer learning per- formance as the original PLM. These subnet- works are found using magnitude-based prun- ing. In this paper, we find that the BERT sub- networks have even more potential than these studies have shown. Firstly, we discover that the success of magnitude pruning can be at- tributed to the preserved pre-training perfor- mance, which correlates with the downstream transferability. Inspired by this, we propose to directly optimize the subnetwork structure towards the pre-training objectives, which can better preserve the pre-training performance. Specifically, we train binary masks over model weights on the pre-training tasks, with the aim of preserving the universal transferability of the subnetwork, which is agnostic to any spe- cific downstream tasks. We then fine-tune the subnetworks on the GLUE benchmark and the SQuAD dataset. The results show that, com- pared with magnitude pruning, mask training can effectively find BERT subnetworks with improved overall performance on downstream tasks. Moreover, our method is also more ef- ficient in searching subnetworks and more ad- vantageous when fine-tuning within a certain range of data scarcity. Our code is available at https://github.com/llyx97/TAMT. 1 Introduction The NLP community has witnessed a remarkable success of pre-trained language models (PLMs). After being pre-trained on unlabelled corpus in a self-supervised manner, PLMs like BERT (De- vlin et al., 2019) can be fine-tuned as a universal text encoder on a wide range of downstream tasks. Joint work with Pattern Recognition Center, WeChat AI, Tencent Inc, China.∗Zheng Lin and Peng Fu are the corre- sponding authors. 0.5 0.6 0.7 0.8 0.9 Sparsity 40 50 60 70 80 Avg Score Downstream Performance full bert Rand OMP 0.5 0.6 0.7 0.8 0.9 Sparsity 0 2 4 6 8 10 MLM Loss Pre-training Performance 5 10 15 20 25 Score Gap −3 −2 −1 0 1 2 3 MLM Loss Gap Gap between OMP & Rand Figure 1: Average downstream performance (left) and pre-training performance (right) of OMP and random subnetworks of BERTBASE. See Appendix A.1 for the downstream results of each task. However, the growing performance of BERT is driven, to a large extent, by scaling up the model size, which hinders the fine-tuning and deployment of BERT in resource-constrained scenarios. At the same time, the lottery ticket hypothesis (LTH) (Frankle and Carbin, 2019) emerges as an active sub-field of model compression. The LTH states that …

Content truncated to 500 words