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

Abstract: Recent studies on the lottery ticket hypothesis (LTH) show that pre-trained language models (PLMs) like BERT contain matching subnetworks that have similar transfer learning performance as the original PLM. These subnetworks are found using magnitude-based pruning. In this paper, we find that the BERT subnetworks have even more potential than these studies have shown. Firstly, we discover that the success of magnitude pruning can be attributed to the preserved pre-training performance, 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 specific downstream tasks. We then fine-tune the subnetworks on the GLUE benchmark and the SQuAD dataset. The results show that, compared with magnitude pruning, mask training can effectively find BERT subnetworks with improved overall performance on downstream tasks. Moreover, our method is also more efficient in searching subnetworks and more advantageous when fine-tuning within a certain range of data scarcity. Our code is available at https://github.com/llyx97/TAMT .

## 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 @xcite 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 corresponding authors. 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) @xcite emerges as an active sub-field of model compression. the lth states that randomly initialized dense networks contain sparse matching subnetworks, i.e., winning tickets (wts), that can be trained in isolation to similar test accuracy as the full model. the original work of lth and subsequent studies have demonstrated that such wts do exist at random initialization or an early point of training @xcite . this implicates that it is possible to reduce training and inference cost via lth. recently, chen et al. (2020) extend the original lth to the pre-training and fine-tuning paradigm, exploring the existence of matching subnetworks in pre-trained bert. such subnetworks are smaller in size, while they can preserve the universal transferability of the full model. encouragingly, @xcite demonstrate that bert indeed contains matching subnetworks that are transferable to multiple downstream tasks without compromising accuracy. these subnetworks are found using iterative magnitude pruning (imp) @xcite on the pre-training task of masked language modeling (mlm), or by directly compressing bert with oneshot magnitude pruning (omp), both of which are agnostic to any specific task. in this paper, we follow @xcite to study the question of lth in bert transfer learning. we find that there is a correlation, to certain extent, between the performance of a bert subnetwork on the pre-training task (right after pruning), and its downstream performance (after fine-tuning). as shown by fig. 1 , the omp subnetworks significantly outperform random subnetworks at 50% sparsity in terms of both mlm loss and downstream score. however, with the increase of model sparsity, the downstream performance and pre-training performance degrade simultaneously. this phenomenon suggests that we might be able to further improve the transferability of bert subnetworks by discovering the structures that better preserve the pre-training performance. to this end, we propose to search transferable bert subnetworks via task-agnostic mask training (tamt), which learns selective binary masks over the model weights on pre-training tasks. in this way, the structure of a subnetwork is directly optimized towards the pre-training objectives, which can preserve the pre-training performance better than heuristically retaining the weights with large magnitudes. the training objective of the masks is a free choice, which can be designed as any loss functions that are agnostic to the downstream tasks. in particular, we investigate the use of mlm loss and a loss based on knowledge distillation (kd) @xcite . to examine the effectiveness of the proposal, we train the masks on the wikitext dataset @xcite for language modeling and then finetune the searched subnetworks on a wide variety of downstream tasks, including the glue benchmark @xcite for natural language understanding (nlu) and the squad dataset @xcite for question answering (qa). the empirical results show that, through mask training, we can indeed find subnetworks with lower pre-training loss and better downstream transferability than omp and imp. compared with imp, which also involves training (the weights) on the pre-training task, mask training requires much fewer training iterations to reach the same performance. moreover, the subnetworks found by mask training are generally more robust when being fine-tuned with reduced data, as long as the training data is not extremely scarce. in summary, our contributions are: ## conclusions in this paper, we address the problem of searching transferable bert subnetworks. we first show that there exist correlations between the pre-training performance and downstream transferablility of a subnetwork. motivated by this, we devise a subnetwork searching method based on task-agnostic mask training (tamt). we empirically show that tamt with mlm loss or kd loss achieve better pre-training and downstream performance than the magnitude pruning, which is recently shown to be successful in finding universal bert subnetworks. tamt is also more efficient in mask searching and produces more robust subnetworks when being fine-tuned within a certain range of data scarsity.